测试 yolov8 分割模型 边缘检测

发现 cfg/default.yaml 参数 mask_ratio 等于4 直接训练如下边缘分割标签,推理时mask 稀疏，训练时分数偏低,mask_ratio 改为1训练时打印的mask 的 P指标一直为0,将imgsz=原图size 训练分数也不高

标注用的是labelme多边形

阅读源码发现可能是因为mask缩放导致

且出现上边缘mask被box过度剪裁的情况

修改了源码中的两处,还是保持mask_ratio等于4,重新训练,推理如下,虽然mask粗糙但几乎不产生断裂

修改如下:

1.ultralytics/data/utils.py

从

def polygon2mask(imgsz, polygons, color=1, downsample_ratio=1):
    mask = np.zeros(imgsz, dtype=np.uint8)
    polygons = np.asarray(polygons, dtype=np.int32)
    polygons = polygons.reshape((polygons.shape[0], -1, 2))
    cv2.fillPoly(mask, polygons, color=color)
    nh, nw = (imgsz[0] // downsample_ratio, imgsz[1] // downsample_ratio)
    # Note: fillPoly first then resize is trying to keep the same loss calculation method when mask-ratio=1
    return cv2.resize(mask, (nw, nh))

到

def polygon2mask(imgsz, polygons, color=1, downsample_ratio=1):
    mask = np.zeros((imgsz[0]// downsample_ratio,imgsz[1]// downsample_ratio), dtype=np.uint8)
    polygons=[[j*0.25 for j in i] for i in polygons]
    polygons = np.asarray(polygons, dtype=np.int32)
    polygons = polygons.reshape((polygons.shape[0], -1, 2))
    cv2.fillPoly(mask, polygons, color=color)
    return mask

2.ultralytics/models/yolo/segment/predict.py

类SegmentationPredictor postprocess方法 外扩mask 1个像素

def postprocess(self, preds, img, orig_imgs):
    """Applies non-max suppression and processes detections for each image in an input batch."""
    p = ops.non_max_suppression(
        preds[0],
        self.args.conf,
        self.args.iou,
        agnostic=self.args.agnostic_nms,
        max_det=self.args.max_det,
        nc=len(self.model.names),
        classes=self.args.classes,
    )

    if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
        orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)

    results = []
    proto = preds[1][-1] if len(preds[1]) == 3 else preds[1]  # second output is len 3 if pt, but only 1 if exported
    for i, pred in enumerate(p):
        orig_img = orig_imgs[i]
        img_path = self.batch[0][i]
        if not len(pred):  # save empty boxes
            masks = None
        elif self.args.retina_masks:
            pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
            masks = ops.process_mask_native(proto[i], pred[:, 6:], pred[:, :4], orig_img.shape[:2])  # HWC
        else:
            #外扩盒子
            c, mh, mw = proto[i].shape  # CHW
            ih, iw = img.shape[2:]
            pred[:, :4][:, 0] -= iw / mw*1
            pred[:, :4][:, 1] -= ih / mh*1
            pred[:, :4][:, 2] += iw / mw*1
            pred[:, :4][:, 3] += ih / mh*1

            masks = ops.process_mask(proto[i], pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True)  # HWC
            pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
        results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks))
    return results


更正第2点:
 取消 SegmentationPredictor postprocess方法 外扩mask 1个像素 的更改,我的理解是这里值影响推理,所以将训练和验证时的标签的盒子外扩才能使模型更不容易出现裁切边缘mask的情况,尤其是水平和铅直的边缘

复制 ultralytics/data/augment.py 中的 Format类 为

FormatExtendBoxBorder

修改FormatExtendBoxBorder的 

def __call__(self, labels):

方法，增加外扩逻辑

        ...

#-------------------#
instances.denormalize(w, h)
ori_format = instances._bboxes.format
instances.convert_bbox(format="xyxy")
mask_ratio_h=h//list(labels["masks"].shape)[1]
mask_ratio_w = w // list(labels["masks"].shape)[2]
instances._bboxes.add(offset=(-2*mask_ratio_w, -2*mask_ratio_h, 2*mask_ratio_w, 2*mask_ratio_h))
if ori_format != "xyxy":
    instances.convert_bbox(format=ori_format)

instances.clip(w, h)
#-------------------#
if self.normalize:
    instances.normalize(w, h)

将 ultralytics/data/dataset.py

YOLODataset的

build_transforms方法中

'''transforms.append(
    Format(
        bbox_format="xywh",
        normalize=True,
        return_mask=self.use_segments,
        return_keypoint=self.use_keypoints,
        return_obb=self.use_obb,
        batch_idx=True,
        mask_ratio=hyp.mask_ratio,
        mask_overlap=hyp.overlap_mask,
    )
)'''
#改为
transforms.append(
    FormatExtendBoxBorder(
        bbox_format="xywh",
        normalize=True,
        return_mask=self.use_segments,
        return_keypoint=self.use_keypoints,
        return_obb=self.use_obb,
        batch_idx=True,
        mask_ratio=hyp.mask_ratio,
        mask_overlap=hyp.overlap_mask,
    )

#*由于边缘检测对mask要求越精细越好,要在图像宽高在1024以上预测mask

尝试在yolov8m-seg.pt上imgsz=640的情况下做微调,可以得到不错的分割结果，但是因为缩放的原因，预测的mask的线宽达到十几像素

尝试在yolov8m-seg.pt上imgsz=1440的情况下做微调,出现明显漏检，以及预测框偏小的情况,从0开始训练似乎效果也不理想

后来在网友的建议下，换成p6的分割模型，从0开始训练，还是imgsz=1440，实际是1472, 结果相当不错，mask线宽可能只有5,6个像素