67 lines
2.2 KiB
Python
67 lines
2.2 KiB
Python
from ultralytics import YOLO
|
|
import os
|
|
from tkinter import Tk, filedialog
|
|
import cv2
|
|
import numpy as np
|
|
|
|
# 載入模型
|
|
model = YOLO('E:/code/image_segmentation/runs/segment/train14/weights/best.pt')
|
|
|
|
# 使用 tkinter 開啟檔案選擇對話框
|
|
def select_files():
|
|
root = Tk()
|
|
root.withdraw() # 隱藏主視窗
|
|
file_paths = filedialog.askopenfilenames(
|
|
title="Select Images",
|
|
filetypes=[("Image files", "*.bmp;*.png;*.jpg;*.jpeg")]
|
|
)
|
|
return file_paths
|
|
|
|
# 選擇影像檔案
|
|
image_paths = select_files()
|
|
if not image_paths:
|
|
print("No files selected.")
|
|
else:
|
|
# 執行推論
|
|
for image_path in image_paths:
|
|
# 推論
|
|
results = model(
|
|
source=image_path,
|
|
device='0', # 使用 GPU 若發生錯誤改成CPU
|
|
max_det=8,
|
|
iou=0.1,
|
|
conf=0.7
|
|
)
|
|
|
|
# 取得圖片名稱和目錄
|
|
img_name = os.path.basename(image_path)
|
|
output_dir = os.path.join(os.path.dirname(image_path), "segmented_results")
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# 處理每張圖片的結果
|
|
for result in results:
|
|
original_image = result.orig_img
|
|
masks = result.masks.data.cpu().numpy() if result.masks else None
|
|
|
|
if masks is not None:
|
|
# 初始化遮罩影像
|
|
mask_image = np.zeros_like(original_image)
|
|
|
|
for mask in masks:
|
|
mask = mask.astype(np.uint8) * 255
|
|
|
|
# 調整遮罩大小與原始影像一致
|
|
mask_resized = cv2.resize(mask, (original_image.shape[1], original_image.shape[0]))
|
|
|
|
# 將單通道遮罩轉換為三通道
|
|
mask_colored = cv2.merge([mask_resized, mask_resized, mask_resized])
|
|
|
|
# 疊加遮罩到影像
|
|
mask_image = cv2.addWeighted(mask_image, 1, mask_colored, 1, 0)
|
|
|
|
# 保存結果
|
|
output_path = os.path.join(output_dir, f"segmented_{img_name}")
|
|
cv2.imwrite(output_path, mask_image)
|
|
print(f"Saved segmented result to: {output_path}")
|
|
print("Inference completed!")
|