在YOLOv8-seg的实例分割任务中，mask_coeff @ proto（掩码系数与原型掩码的线性组合） 是一种高效的分割掩码生成方法，其核心思想是通过共享原型掩码（Proto） 和实例特定的掩码系数 来生成最终的分割结果。这种方法与直接对每个像素进行分类相比，具有以下显著优势：
核心机制： 原型掩码（Proto）与掩码系数（Mask Coefficients）
- (1) 原型掩码（Proto） 定义： 原型掩码是一个共享的基函数矩阵（形状如 [B, C, H, W]），其中 C 是基的数量（如32），H×W 是下采样后的特征图尺寸。 作用：类似于一组通用的形状模板，可以线性组合生成不同实例的掩码。
- (2)掩码系数（Mask Coefficients） 定义： 每个实例的掩码系数是一个向量（形状如 [num_instances, C]），表示该实例在原型掩码上的权重。
作用：通过线性组合（mask_coeff @ proto）生成该实例的掩码。
代码部分讲解
分割算法label获取
class Format:
"""
A class for formatting image annotations for object detection, instance segmentation, and pose estimation tasks.
This class standardizes image and instance annotations to be used by the `collate_fn` in PyTorch DataLoader.
Attributes:
bbox_format (str): Format for bounding boxes. Options are 'xywh' or 'xyxy'.
normalize (bool): Whether to normalize bounding boxes.
return_mask (bool): Whether to return instance masks for segmentation.
return_keypoint (bool): Whether to return keypoints for pose estimation.
return_obb (bool): Whether to return oriented bounding boxes.
mask_ratio (int): Downsample ratio for masks.
mask_overlap (bool): Whether to overlap masks.
batch_idx (bool): Whether to keep batch indexes.
bgr (float): The probability to return BGR images.
Methods:
__call__: Formats labels dictionary with image, classes, bounding boxes, and optionally masks and keypoints.
_format_img: Converts image from Numpy array to PyTorch tensor.
_format_segments: Converts polygon points to bitmap masks.
Examples:
>>> formatter = Format(bbox_format="xywh", normalize=True, return_mask=True)
>>> formatted_labels = formatter(labels)
>>> img = formatted_labels["img"]
>>> bboxes = formatted_labels["bboxes"]
>>> masks = formatted_labels["masks"]
"""
def __init__(
self,
bbox_format="xywh",
normalize=True,
return_mask=False,
return_keypoint=False,
return_obb=False,
mask_ratio=4,
mask_overlap=True,
batch_idx=True,
bgr=0.0,
):
"""
Initializes the Format class with given parameters for image and instance annotation formatting.
This class standardizes image and instance annotations for object detection, instance segmentation, and pose
estimation tasks, preparing them for use in PyTorch DataLoader's `collate_fn`.
Args:
bbox_format (str): Format for bounding boxes. Options are 'xywh', 'xyxy', etc.
normalize (bool): Whether to normalize bounding boxes to [0,1].
return_mask (bool): If True, returns instance masks for segmentation tasks.
return_keypoint (bool): If True, returns keypoints for pose estimation tasks.
return_obb (bool): If True, returns oriented bounding boxes.
mask_ratio (int): Downsample ratio for masks.
mask_overlap (bool): If True, allows mask overlap.
batch_idx (bool): If True, keeps batch indexes.
bgr (float): Probability of returning BGR images instead of RGB.
Attributes:
bbox_format (str): Format for bounding boxes.
normalize (bool): Whether bounding boxes are normalized.
return_mask (bool): Whether to return instance masks.
return_keypoint (bool): Whether to return keypoints.
return_obb (bool): Whether to return oriented bounding boxes.
mask_ratio (int): Downsample ratio for masks.
mask_overlap (bool): Whether masks can overlap.
batch_idx (bool): Whether to keep batch indexes.
bgr (float): The probability to return BGR images.
Examples:
>>> format = Format(bbox_format="xyxy", return_mask=True, return_keypoint=False)
>>> print(format.bbox_format)
xyxy
"""
self.bbox_format = bbox_format
self.normalize = normalize
self.return_mask = return_mask # set False when training detection only
self.return_keypoint = return_keypoint
self.return_obb = return_obb
self.mask_ratio = mask_ratio
self.mask_overlap = mask_overlap
self.batch_idx = batch_idx # keep the batch indexes
self.bgr = bgr
def __call__(self, labels):
"""
Formats image annotations for object detection, instance segmentation, and pose estimation tasks.
This method standardizes the image and instance annotations to be used by the `collate_fn` in PyTorch
DataLoader. It processes the input labels dictionary, converting annotations to the specified format and
applying normalization if required.
Args:
labels (dict): A dictionary containing image and annotation data with the following keys:
- 'img': The input image as a numpy array.
- 'cls': Class labels for instances.
- 'instances': An Instances object containing bounding boxes, segments, and keypoints.
Returns:
(dict): A dictionary with formatted data, including:
- 'img': Formatted image tensor.
- 'cls': Class label's tensor.
- 'bboxes': Bounding boxes tensor in the specified format.
- 'masks': Instance masks tensor (if return_mask is True).
- 'keypoints': Keypoints tensor (if return_keypoint is True).
- 'batch_idx': Batch index tensor (if batch_idx is True).
Examples:
>>> formatter = Format(bbox_format="xywh", normalize=True, return_mask=True)
>>> labels = {"img": np.random.rand(640, 640, 3), "cls": np.array([0, 1]), "instances": Instances(...)}
>>> formatted_labels = formatter(labels)
>>> print(formatted_labels.keys())
"""
img = labels.pop("img")
h, w = img.shape[:2]
cls = labels.pop("cls")
instances = labels.pop("instances")
instances.convert_bbox(format=self.bbox_format)
instances.denormalize(w, h)
nl = len(instances)
if self.return_mask:
if nl:
masks, instances, cls = self._format_segments(instances, cls, w, h)
masks = torch.from_numpy(masks)
else:
masks = torch.zeros(
1 if self.mask_overlap else nl, img.shape[0] // self.mask_ratio, img.shape[1] // self.mask_ratio
)
labels["masks"] = masks
labels["img"] = self._format_img(img)
labels["cls"] = torch.from_numpy(cls) if nl else torch.zeros(nl)
labels["bboxes"] = torch.from_numpy(instances.bboxes) if nl else torch.zeros((nl, 4))
if self.return_keypoint:
labels["keypoints"] = torch.from_numpy(instances.keypoints)
if self.normalize:
labels["keypoints"][..., 0] /= w
labels["keypoints"][..., 1] /= h
if self.return_obb:
labels["bboxes"] = (
xyxyxyxy2xywhr(torch.from_numpy(instances.segments)) if len(instances.segments) else torch.zeros((0, 5))
)
# NOTE: need to normalize obb in xywhr format for width-height consistency
if self.normalize:
labels["bboxes"][:, [0, 2]] /= w
labels["bboxes"][:, [1, 3]] /= h
# Then we can use collate_fn
if self.batch_idx:
labels["batch_idx"] = torch.zeros(nl)
return labels
1.1 polygons的处理
该函数将多边形坐标列表转换为二值掩膜（Binary Mask），并支持下采样。其核心步骤如下：
- 初始化掩膜：创建全零的掩膜数组（形状为 imgsz）。
- 处理多边形坐标：
(1) 将输入的多边形列表转换为 np.int32 类型的 numpy 数组。
(2) 将多边形坐标 reshape 为 (N_polygons, N_points, 2) 格式。 - 填充多边形：使用 cv2.fillPoly 将多边形区域填充为指定 color。
- 下采样：通过 cv2.resize 将掩膜调整为 imgsz // downsample_ratio 的尺寸。
def polygon2mask(imgsz, polygons, color=1, downsample_ratio=1):
"""
Convert a list of polygons to a binary mask of the specified image size.
Args:
imgsz (tuple): The size of the image as (height, width).
polygons (list[np.ndarray]): A list of polygons. Each polygon is an array with shape [N, M], where
N is the number of polygons, and M is the number of points such that M % 2 = 0.
color (int, optional): The color value to fill in the polygons on the mask.
downsample_ratio (int, optional): Factor by which to downsample the mask.
Returns:
(np.ndarray): A binary mask of the specified image size with the polygons filled in.
"""
# 初始化掩膜
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))
1.2 重叠掩膜
该函数将多边形列表转换为一个 重叠掩膜，每个像素的值表示覆盖它的实例的排序后的索引（按面积降序）。返回的 masks 数组中：
- 值 0 表示无实例覆盖。
- 值 i+1 表示该像素属于按面积降序排列的第 i 个实例。
def polygons2masks_overlap(imgsz, segments, downsample_ratio=1):
"""Return a (640, 640) overlap mask."""
masks = np.zeros(
(imgsz[0] // downsample_ratio, imgsz[1] // downsample_ratio),
dtype=np.int32 if len(segments) > 255 else np.uint8,
)
areas = []
ms = []
for si in range(len(segments)):
mask = polygon2mask(imgsz, [segments[si].reshape(-1)], downsample_ratio=downsample_ratio, color=1)
ms.append(mask.astype(masks.dtype))
areas.append(mask.sum())
# 按面积降序排序
areas = np.asarray(areas)
index = np.argsort(-areas)
ms = np.array(ms)[index]
# 修正叠加逻辑：大实例优先，小实例仅覆盖未被覆盖区域
for i in range(len(segments)):
mask = ms[i] * (i + 1)
masks = masks + mask
masks = np.clip(masks, a_min=0, a_max=i + 1)
return masks, index
seg_predit 代码部分
class Segment(Detect):
"""YOLO Segment head for segmentation models."""
def __init__(self, nc=80, nm=32, npr=256, ch=()):
"""Initialize the YOLO model attributes such as the number of masks, prototypes, and the convolution layers."""
super().__init__(nc, ch)
self.nm = nm # number of masks
self.npr = npr # number of protos
self.proto = Proto(ch[0], self.npr, self.nm) # protos
c4 = max(ch[0] // 4, self.nm)
self.cv4 = nn.ModuleList(nn.Sequential(Conv(x, c4, 3), Conv(c4, c4, 3), nn.Conv2d(c4, self.nm, 1)) for x in ch)
def forward(self, x):
"""Return model outputs and mask coefficients if training, otherwise return outputs and mask coefficients."""
p = self.proto(x[0]) # mask protos
bs = p.shape[0] # batch size
# 计算掩码系数 mc
# self.cv4[i](x[i])：对每个特征图 x[i]（如P3、P4、P5）应用 cv4 模块，输出形状为 [bs, nm, H_i, W_i]。
# view(bs, self.nm, -1)：将空间维度 H_i*W_i 展平，得到 [bs, nm, H_i*W_i]。
# total_positions]，其中 total_positions 是所有特征图的像素总和。
mc = torch.cat([self.cv4[i](x[i]).view(bs, self.nm, -1) for i in range(self.nl)], 2) # mask coefficients
x = Detect.forward(self, x)
if self.training:
return x, mc, p
return (torch.cat([x, mc], 1), p) if self.export else (torch.cat([x[0], mc], 1), (x[1], mc, p))
2.1 seg_loss的实现细节
def calculate_segmentation_loss(
self,
fg_mask: torch.Tensor, # (BS, N_anchors)
masks: torch.Tensor, # overlap=False：形状 (BS, H, W) 若overlap=True: 形状 (BS, num_instances, H, W)
target_gt_idx: torch.Tensor, # 每个锚框对应的真实目标索引 (BS, N_anchors)
target_bboxes: torch.Tensor, # 每个锚框对应的真实边界框坐标 (BS, N_anchors, 4)
batch_idx: torch.Tensor, # 标记每个真实目标属于哪个批次的索引 (N_labels_in_batch, 1)
proto: torch.Tensor, # 原型掩码特征图，用于生成所有实例的掩码 (BS, 32, H_p, W_p)
pred_masks: torch.Tensor, # 预测的掩码系数，用于线性组合 proto 生成最终掩码 (BS, N_anchors, 32)
imgsz: torch.Tensor,
overlap: bool, # 是否处理重叠掩码（True 表示多实例掩码，False 表示单通道掩码）
) -> torch.Tensor:
"""
Calculate the loss for instance segmentation.
Args:
fg_mask (torch.Tensor): A binary tensor of shape (BS, N_anchors) indicating which anchors are positive.
masks (torch.Tensor): Ground truth masks of shape (BS, H, W) if `overlap` is False, otherwise (BS, ?, H, W).
target_gt_idx (torch.Tensor): Indexes of ground truth objects for each anchor of shape (BS, N_anchors).
target_bboxes (torch.Tensor): Ground truth bounding boxes for each anchor of shape (BS, N_anchors, 4).
batch_idx (torch.Tensor): Batch indices of shape (N_labels_in_batch, 1).
proto (torch.Tensor): Prototype masks of shape (BS, 32, H, W).
pred_masks (torch.Tensor): Predicted masks for each anchor of shape (BS, N_anchors, 32).
imgsz (torch.Tensor): Size of the input image as a tensor of shape (2), i.e., (H, W).
overlap (bool): Whether the masks in `masks` tensor overlap.
Returns:
(torch.Tensor): The calculated loss for instance segmentation.
Notes:
The batch loss can be computed for improved speed at higher memory usage.
For example, pred_mask can be computed as follows:
pred_mask = torch.einsum('in,nhw->ihw', pred, proto) # (i, 32) @ (32, 160, 160) -> (i, 160, 160)
"""
_, _, mask_h, mask_w = proto.shape
loss = 0
# Normalize to 0-1
target_bboxes_normalized = target_bboxes / imgsz[[1, 0, 1, 0]] # 将坐标归一化到 [0, 1]
# Areas of target bboxes
marea = xyxy2xywh(target_bboxes_normalized)[..., 2:].prod(2) # 计算目标框的面积（宽 × 高）
# Normalize to mask size
# mxyxy 将目标框坐标转换为原型掩码的尺寸（mask_h 和 mask_w 是 proto 的特征图尺寸）
mxyxy = target_bboxes_normalized * torch.tensor([mask_w, mask_h, mask_w, mask_h], device=proto.device)
# fg_mask_i 是当前批次的正样本掩码
# masks_i 的形状为 (num_instances, H, W)，每个实例对应一个掩码
# gt_mask 二值掩码 目标存在的区域（即真实目标的像素坐标（x,y））对应的值确实是 1，而背景区域的值为 0
for i, single_i in enumerate(zip(fg_mask, target_gt_idx, pred_masks, proto, mxyxy, marea, masks)):
fg_mask_i, target_gt_idx_i, pred_masks_i, proto_i, mxyxy_i, marea_i, masks_i = single_i
if fg_mask_i.any():
mask_idx = target_gt_idx_i[fg_mask_i]
if overlap:
gt_mask = masks_i == (mask_idx + 1).view(-1, 1, 1)
gt_mask = gt_mask.float()
else:
gt_mask = masks[batch_idx.view(-1) == i][mask_idx]
loss += self.single_mask_loss(
gt_mask, pred_masks_i[fg_mask_i], proto_i, mxyxy_i[fg_mask_i], marea_i[fg_mask_i]
)
# WARNING: lines below prevents Multi-GPU DDP 'unused gradient' PyTorch errors, do not remove
else:
loss += (proto * 0).sum() + (pred_masks * 0).sum() # inf sums may lead to nan loss
return loss / fg_mask.sum()
2.2 single_mask_loss 的实现细节
def single_mask_loss(
gt_mask: torch.Tensor, pred: torch.Tensor, proto: torch.Tensor, xyxy: torch.Tensor, area: torch.Tensor
) -> torch.Tensor:
"""
Compute the instance segmentation loss for a single image.
Args:
gt_mask (torch.Tensor): Ground truth mask of shape (n, H, W), where n is the number of objects.
pred (torch.Tensor): Predicted mask coefficients of shape (n, 32).
proto (torch.Tensor): Prototype masks of shape (32, H, W).
xyxy (torch.Tensor): Ground truth bounding boxes in xyxy format, normalized to [0, 1], of shape (n, 4).
area (torch.Tensor): Area of each ground truth bounding box of shape (n,).
Returns:
(torch.Tensor): The calculated mask loss for a single image.
Notes:
The function uses the equation pred_mask = torch.einsum('in,nhw->ihw', pred, proto) to produce the
predicted masks from the prototype masks and predicted mask coefficients.
"""
pred_mask = torch.einsum("in,nhw->ihw", pred, proto) # (n, 32) @ (32, 80, 80) -> (n, 80, 80)
loss = F.binary_cross_entropy_with_logits(pred_mask, gt_mask, reduction="none")
# 裁剪目标框区域，仅计算目标内的损失， 然后按面积加权
return (crop_mask(loss, xyxy).mean(dim=(1, 2)) / area).sum()
2.3 crop_mask 实现细节
目的： 该函数将输入的 掩码（masks） 裁剪到对应的 边界框（bounding boxes） 内，仅保留边界框内的区域，其余区域置零。
def crop_mask(masks, boxes):
"""
Crop masks to bounding boxes.
Args:
masks (torch.Tensor): [n, h, w] tensor of masks.
boxes (torch.Tensor): [n, 4] tensor of bbox coordinates in relative point form.
Returns:
(torch.Tensor): Cropped masks.
"""
_, h, w = masks.shape
x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1) # x1 shape(n,1,1)
r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :] # rows shape(1,1,w)
c = torch.arange(h, device=masks.device, dtype=x1.dtype)[None, :, None] # cols shape(1,h,1)
# mask 的形状为 [n,h,w]，每个元素为 True（1）或 False（0），表示该像素是否在边界框内
return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
扫一扫
3026
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
