utils/segment/general.py - master - githubmirror/yolov5

utils/segment/general.py
113 строк5 KB
Glenn Jocher
Re-export 17 duplicated utilities from ultralytics (#13835)
29 июл 2026, 03:45
29 июл 2026, 03:452e3ec55
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
"""Segmentation mask utils."""

import cv2
import numpy as np
import torch
import torch.nn.functional as F


def crop_mask(masks, boxes):
    """Crop predicted masks by zeroing out everything not in the predicted bbox.

    Args:
        masks (torch.Tensor): Masks with shape (n, h, w).
        boxes (torch.Tensor): Box coordinates with shape (n, 4) in xyxy pixel format (mask resolution).
    """
    _n, h, w = masks.shape
    x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1)  # x1 shape(1,1,n)
    r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :]  # rows shape(1,w,1)
    c = torch.arange(h, device=masks.device, dtype=x1.dtype)[None, :, None]  # cols shape(h,1,1)

    return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))


def process_mask(protos, masks_in, bboxes, shape, upsample=False):
    """Crop mask prototypes with the predicted boxes, then optionally upsample to the input image size.

    Args:
        protos (torch.Tensor): Mask prototypes with shape (mask_dim, mask_h, mask_w).
        masks_in (torch.Tensor): Mask coefficients with shape (n, mask_dim), n is number of masks after NMS.
        bboxes (torch.Tensor): Box coordinates with shape (n, 4).
        shape (tuple): Input image size as (h, w).
        upsample (bool): Whether to upsample the masks to the input image size.

    Returns:
        (torch.Tensor): Binary masks with shape (n, mask_h, mask_w), upsampled to (n, h, w) when upsample=True.
    """
    c, mh, mw = protos.shape  # CHW
    ih, iw = shape
    masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)  # CHW

    downsampled_bboxes = bboxes.clone()
    downsampled_bboxes[:, 0] *= mw / iw
    downsampled_bboxes[:, 2] *= mw / iw
    downsampled_bboxes[:, 3] *= mh / ih
    downsampled_bboxes[:, 1] *= mh / ih

    masks = crop_mask(masks, downsampled_bboxes)  # CHW
    if upsample:
        masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0]  # CHW
    return masks.gt_(0.5)


def process_mask_native(protos, masks_in, bboxes, shape):
    """Crop mask prototypes to the input image size (native), upsampling and then cropping by the predicted boxes.

    Args:
        protos (torch.Tensor): Mask prototypes with shape (mask_dim, mask_h, mask_w).
        masks_in (torch.Tensor): Mask coefficients with shape (n, mask_dim), n is number of masks after NMS.
        bboxes (torch.Tensor): Box coordinates with shape (n, 4).
        shape (tuple): Input image size as (h, w).

    Returns:
        (torch.Tensor): Binary masks with shape (n, h, w).
    """
    c, mh, mw = protos.shape  # CHW
    masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
    gain = min(mh / shape[0], mw / shape[1])  # gain  = old / new
    pad = (mw - shape[1] * gain) / 2, (mh - shape[0] * gain) / 2  # wh padding
    top, left = int(pad[1]), int(pad[0])  # y, x
    bottom, right = int(mh - pad[1]), int(mw - pad[0])
    masks = masks[:, top:bottom, left:right]

    masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0]  # CHW
    masks = crop_mask(masks, bboxes)  # CHW
    return masks.gt_(0.5)


def scale_image(im1_shape, masks, im0_shape, ratio_pad=None):
    """Rescale masks from the model input shape (im1_shape) to the original image shape (im0_shape).

    Args:
        im1_shape (tuple): Model input shape as (h, w).
        masks (np.ndarray): Masks with shape (h, w, num).
        im0_shape (tuple): Original image shape as (h, w, 3).
        ratio_pad (tuple, optional): Ratio and padding for scaling. If None, calculated from the shapes.

    Returns:
        (np.ndarray): Rescaled masks resized to im0_shape.
    """
    # Rescale coordinates (xyxy) from im1_shape to im0_shape
    if ratio_pad is None:  # calculate from im0_shape
        gain = min(im1_shape[0] / im0_shape[0], im1_shape[1] / im0_shape[1])  # gain  = old / new
        pad = (im1_shape[1] - im0_shape[1] * gain) / 2, (im1_shape[0] - im0_shape[0] * gain) / 2  # wh padding
    else:
        pad = ratio_pad[1]
    top, left = int(pad[1]), int(pad[0])  # y, x
    bottom, right = int(im1_shape[0] - pad[1]), int(im1_shape[1] - pad[0])

    if len(masks.shape) < 2:
        raise ValueError(f'"len of masks shape" should be 2 or 3, but got {len(masks.shape)}')
    masks = masks[top:bottom, left:right]
    if masks.ndim == 3 and masks.shape[2] > 128:  # OpenCV 5 lowered CV_CN_MAX from 512 to 128
        masks = [
            cv2.resize(masks[:, :, i : i + 128], (im0_shape[1], im0_shape[0])) for i in range(0, masks.shape[2], 128)
        ]
        masks = np.concatenate([x if x.ndim == 3 else x[:, :, None] for x in masks], axis=2)
    else:
        masks = cv2.resize(masks, (im0_shape[1], im0_shape[0]))

    if len(masks.shape) == 2:
        masks = masks[:, :, None]
    return masks