223 lines
7.6 KiB
Python
223 lines
7.6 KiB
Python
import os
|
|
import cv2
|
|
import torch
|
|
import argparse
|
|
import numpy as np
|
|
from ultralytics import YOLO
|
|
from pathlib import Path
|
|
|
|
# ====================== 配置参数 ======================
|
|
MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/seg/exp7/weights/best.pt"
|
|
SOURCE_IMG_DIR = "/home/hx/yolo/output_masks"
|
|
OUTPUT_DIR = "/home/hx/yolo/output_masks2"
|
|
CONF_THRESHOLD = 0.25
|
|
IOU_THRESHOLD = 0.45
|
|
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
|
|
SAVE_TXT = True
|
|
SAVE_MASKS = True
|
|
VIEW_IMG = False
|
|
LINE_WIDTH = 2
|
|
|
|
def plot_result_with_opacity(result, line_width=2, mask_opacity=0.5):
|
|
"""
|
|
手动绘制 YOLO 分割结果，支持掩码透明度叠加，并修复掩码尺寸不匹配问题
|
|
"""
|
|
img = result.orig_img.copy() # HWC, BGR
|
|
h, w = img.shape[:2]
|
|
|
|
# 获取原始图像尺寸
|
|
orig_shape = img.shape[:2] # (height, width)
|
|
|
|
if result.masks is not None and len(result.boxes) > 0:
|
|
# 将掩码从 GPU 移到 CPU 并转为 numpy
|
|
masks = result.masks.data.cpu().numpy() # (N, H_mask, W_mask)
|
|
|
|
# resize 掩码到原始图像尺寸
|
|
resized_masks = []
|
|
for mask in masks:
|
|
mask_resized = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
|
|
mask_resized = (mask_resized > 0.5).astype(np.uint8) # 二值化
|
|
resized_masks.append(mask_resized)
|
|
resized_masks = np.array(resized_masks)
|
|
|
|
# 随机颜色 (BGR)
|
|
num_masks = len(result.boxes)
|
|
colors = np.random.randint(0, 255, size=(num_masks, 3), dtype=np.uint8)
|
|
|
|
# 创建叠加层
|
|
overlay = img.copy()
|
|
for i in range(num_masks):
|
|
color = colors[i].tolist()
|
|
mask_resized = resized_masks[i]
|
|
overlay[mask_resized == 1] = color
|
|
|
|
# 透明叠加
|
|
cv2.addWeighted(overlay, mask_opacity, img, 1 - mask_opacity, 0, img)
|
|
|
|
# 绘制边界框和标签（保持不变）
|
|
if result.boxes is not None:
|
|
boxes = result.boxes.xyxy.cpu().numpy()
|
|
classes = result.boxes.cls.cpu().numpy().astype(int)
|
|
confidences = result.boxes.conf.cpu().numpy()
|
|
|
|
colors = np.random.randint(0, 255, size=(len(classes), 3), dtype=np.uint8)
|
|
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
font_scale = 0.6
|
|
thickness = 1
|
|
|
|
for i in range(len(boxes)):
|
|
box = boxes[i].astype(int)
|
|
cls_id = classes[i]
|
|
conf = confidences[i]
|
|
color = colors[i].tolist()
|
|
|
|
# 绘制矩形框
|
|
cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), color, line_width)
|
|
|
|
# 标签文本
|
|
label = f"{cls_id} {conf:.2f}"
|
|
|
|
# 获取文本大小
|
|
(text_w, text_h), baseline = cv2.getTextSize(label, font, font_scale, thickness)
|
|
text_h += baseline
|
|
|
|
# 绘制标签背景
|
|
cv2.rectangle(img, (box[0], box[1] - text_h - 6), (box[0] + text_w, box[1]), color, -1)
|
|
# 绘制文本
|
|
cv2.putText(img, label, (box[0], box[1] - 4), font, font_scale,
|
|
(255, 255, 255), thickness, cv2.LINE_AA)
|
|
|
|
return img
|
|
|
|
|
|
def run_segmentation_inference(
|
|
model_path,
|
|
source,
|
|
output_dir,
|
|
conf_threshold=0.25,
|
|
iou_threshold=0.45,
|
|
device="cuda:0",
|
|
save_txt=True,
|
|
save_masks=True,
|
|
view_img=False,
|
|
line_width=2,
|
|
):
|
|
print(f"🚀 加载模型: {model_path}")
|
|
print(f"💻 使用设备: {device}")
|
|
|
|
# 加载模型
|
|
model = YOLO(model_path)
|
|
|
|
# 创建输出目录
|
|
output_dir = Path(output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
txt_dir = output_dir / "labels"
|
|
mask_dir = output_dir / "masks"
|
|
if save_txt:
|
|
txt_dir.mkdir(exist_ok=True)
|
|
if save_masks:
|
|
mask_dir.mkdir(exist_ok=True)
|
|
|
|
# 获取图像文件列表
|
|
source = Path(source)
|
|
if source.is_file():
|
|
img_files = [source]
|
|
else:
|
|
img_files = list(source.glob("*.jpg")) + \
|
|
list(source.glob("*.jpeg")) + \
|
|
list(source.glob("*.png")) + \
|
|
list(source.glob("*.bmp"))
|
|
|
|
if not img_files:
|
|
print(f"❌ 在 {source} 中未找到图像文件")
|
|
return
|
|
|
|
print(f"🖼️ 共 {len(img_files)} 张图片待推理...")
|
|
|
|
# 推理循环
|
|
for img_path in img_files:
|
|
print(f"🔍 推理: {img_path.name}")
|
|
|
|
# 执行推理
|
|
results = model(
|
|
source=str(img_path),
|
|
conf=conf_threshold,
|
|
iou=iou_threshold,
|
|
imgsz=640,
|
|
device=device,
|
|
verbose=True
|
|
)
|
|
|
|
result = results[0]
|
|
orig_img = result.orig_img # 原始图像
|
|
|
|
# ✅ 使用自定义绘制函数（支持透明度）
|
|
plotted = plot_result_with_opacity(result, line_width=line_width, mask_opacity=0.5)
|
|
|
|
# 保存可视化图像
|
|
save_path = output_dir / f"seg_{img_path.name}"
|
|
cv2.imwrite(str(save_path), plotted)
|
|
print(f"✅ 保存结果: {save_path}")
|
|
|
|
# 保存 YOLO 格式标签（多边形）
|
|
if save_txt and result.masks is not None:
|
|
txt_path = txt_dir / (img_path.stem + ".txt")
|
|
with open(txt_path, 'w') as f:
|
|
for i in range(len(result.boxes)):
|
|
cls_id = int(result.boxes.cls[i])
|
|
seg = result.masks.xy[i] # 多边形点 (N, 2)
|
|
seg = seg.flatten()
|
|
seg = seg / [orig_img.shape[1], orig_img.shape[0]] # 归一化
|
|
seg = seg.tolist()
|
|
line = f"{cls_id} {' '.join(f'{x:.6f}' for x in seg)}\n"
|
|
f.write(line)
|
|
print(f"📝 保存标签: {txt_path}")
|
|
|
|
# 保存合并的掩码图
|
|
if save_masks and result.masks is not None:
|
|
mask = result.masks.data.cpu().numpy()
|
|
combined_mask = (mask.sum(axis=0) > 0).astype(np.uint8) * 255 # 合并所有掩码
|
|
mask_save_path = mask_dir / f"mask_{img_path.stem}.png"
|
|
cv2.imwrite(str(mask_save_path), combined_mask)
|
|
print(f"🎨 保存掩码: {mask_save_path}")
|
|
|
|
# 实时显示（可选）
|
|
if view_img:
|
|
cv2.imshow("Segmentation Result", plotted)
|
|
if cv2.waitKey(0) == 27: # ESC 退出
|
|
cv2.destroyAllWindows()
|
|
break
|
|
|
|
if view_img:
|
|
cv2.destroyAllWindows()
|
|
print(f"\n🎉 推理完成！结果保存在: {output_dir}")
|
|
|
|
|
|
# ====================== 主程序 ======================
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model", default=MODEL_PATH, help="模型权重路径")
|
|
parser.add_argument("--source", default=SOURCE_IMG_DIR, help="图片路径或文件夹")
|
|
parser.add_argument("--output", default=OUTPUT_DIR, help="输出目录")
|
|
parser.add_argument("--conf", type=float, default=CONF_THRESHOLD, help="置信度阈值")
|
|
parser.add_argument("--iou", type=float, default=IOU_THRESHOLD, help="IoU 阈值")
|
|
parser.add_argument("--device", default=DEVICE, help="设备: cuda:0, cpu")
|
|
parser.add_argument("--view-img", action="store_true", help="显示图像")
|
|
parser.add_argument("--save-txt", action="store_true", help="保存标签")
|
|
parser.add_argument("--save-masks", action="store_true", help="保存掩码")
|
|
|
|
opt = parser.parse_args()
|
|
|
|
run_segmentation_inference(
|
|
model_path=opt.model,
|
|
source=opt.source,
|
|
output_dir=opt.output,
|
|
conf_threshold=opt.conf,
|
|
iou_threshold=opt.iou,
|
|
device=opt.device,
|
|
save_txt=opt.save_txt,
|
|
save_masks=opt.save_masks,
|
|
view_img=opt.view_img,
|
|
line_width=LINE_WIDTH,
|
|
) |