加载自定义数据集训练的YOLOv9实例分割模型遇AttributeError求助

错误原因 #

• 警告已明确提示：YOLO SegmentationModel is not yet AutoShape compatible，说明该分割模型不支持自动处理图片路径字符串，必须手动将图片预处理为模型可接受的张量输入。
• 报错'str' object has no attribute 'shape'是因为直接将图片路径字符串传入模型的forward方法，而该方法需要的是带shape属性的张量/数组类型输入。

修正后的推理代码 #

from pathlib import Path
import numpy as np
import torch
import cv2

# 加载模型
model = torch.hub.load('.', 'custom', path='yolov9-inst/runs/train-seg/gelan-c-seg15/weights/best.pt', source='local')
# 设置模型为评估模式
model.eval()

# 加载并预处理图片
img_path = 'WALL-INSTANCEE-2/test/images/5a243513a69b150001f56c31_emptyroom6_jpeg_jpg.rf.7aa8f6a9aefbb1c76adc60a7b392dcd6.jpg'
img = cv2.imread(img_path)
# 转换颜色空间（YOLO默认使用RGB，cv2读取的是BGR）
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# 调整维度：(H, W, C) -> (C, H, W)，添加batch维度转为(1, C, H, W)
img_tensor = torch.from_numpy(img_rgb).permute(2, 0, 1).unsqueeze(0).float()
# 归一化到0-1范围
img_tensor /= 255.0
# 如果有GPU，将张量转到cuda
if torch.cuda.is_available():
    img_tensor = img_tensor.cuda()

# 推理，关闭梯度计算节省资源
with torch.no_grad():
    results = model(img_tensor)

# 提取分割掩码
# 模型输出的masks形状一般为(batch, num_masks, H, W)
masks = results[0].masks.data.cpu().numpy()

# 处理每个分割对象的区域
for idx, mask in enumerate(masks[0]):  # 取batch中的第一张图
    # 将掩码二值化（阈值可根据实际情况调整）
    binary_mask = (mask > 0.5).astype(np.uint8)
    # 提取原图中的对应分割区域
    segmented_region = cv2.bitwise_and(img, img, mask=binary_mask)
    # 保存分割区域（可根据需求替换为其他处理逻辑）
    cv2.imwrite(f'segmented_region_{idx}.png', segmented_region)

关键说明 #

• 必须手动完成颜色空间转换、维度调整、归一化三步预处理，确保输入符合模型要求。
• 使用torch.no_grad()关闭梯度计算，避免不必要的内存占用。
• 分割掩码需二值化处理，0.5为通用阈值，可根据模型输出的置信度调整。

内容的提问来源于stack exchange，提问作者Plutus