# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license from __future__ import annotations from pathlib import Path from typing import Any import numpy as np import torch import torch.nn.functional as F from ultralytics.models.yolo.detect import DetectionValidator from ultralytics.utils import LOGGER, ops from ultralytics.utils.checks import check_requirements from ultralytics.utils.metrics import SegmentMetrics, mask_iou class SegmentationValidator(DetectionValidator): """A class extending the DetectionValidator class for validation based on a segmentation model. This validator handles the evaluation of segmentation models, processing both bounding box and mask predictions to compute metrics such as mAP for both detection and segmentation tasks. Attributes: plot_masks (list): List to store masks for plotting. process (callable): Function to process masks based on save_json and save_txt flags. args (SimpleNamespace): Arguments for the validator. metrics (SegmentMetrics): Metrics calculator for segmentation tasks. stats (dict): Dictionary to store statistics during validation. Examples: >>> from ultralytics.models.yolo.segment import SegmentationValidator >>> args = dict(model="yolo26n-seg.pt", data="coco8-seg.yaml") >>> validator = SegmentationValidator(args=args) >>> validator() """ def __init__(self, dataloader=None, save_dir=None, args=None, _callbacks: dict | None = None) -> None: """Initialize SegmentationValidator and set task to 'segment', metrics to SegmentMetrics. Args: dataloader (torch.utils.data.DataLoader, optional): DataLoader to use for validation. save_dir (Path, optional): Directory to save results. args (dict, optional): Arguments for the validator. _callbacks (dict, optional): Dictionary of callback functions. """ super().__init__(dataloader, save_dir, args, _callbacks) self.process = None self.args.task = "segment" self.metrics = SegmentMetrics() def preprocess(self, batch: dict[str, Any]) -> dict[str, Any]: """Preprocess batch of images for YOLO segmentation validation. Args: batch (dict[str, Any]): Batch containing images and annotations. Returns: (dict[str, Any]): Preprocessed batch. """ batch = super().preprocess(batch) batch["masks"] = batch["masks"].float() return batch def init_metrics(self, model: torch.nn.Module) -> None: """Initialize metrics and select mask processing function based on save_json flag. Args: model (torch.nn.Module): Model to validate. """ super().init_metrics(model) if self.args.save_json: check_requirements("faster-coco-eval>=1.6.7") # More accurate vs faster self.process = ops.process_mask_native if self.args.save_json or self.args.save_txt else ops.process_mask def get_desc(self) -> str: """Return a formatted description of evaluation metrics.""" return ("%22s" + "%11s" * 10) % ( "Class", "Images", "Instances", "Box(P", "R", "mAP50", "mAP50-95)", "Mask(P", "R", "mAP50", "mAP50-95)", ) def postprocess(self, preds: list[torch.Tensor]) -> list[dict[str, torch.Tensor]]: """Post-process YOLO predictions and return output detections with proto. Args: preds (list[torch.Tensor]): Raw predictions from the model. Returns: (list[dict[str, torch.Tensor]]): Processed detection predictions with masks. """ proto = preds[0][1] if isinstance(preds[0], tuple) else preds[1] preds = super().postprocess(preds[0]) imgsz = [4 * x for x in proto.shape[2:]] # get image size from proto for i, pred in enumerate(preds): coefficient = pred.pop("extra") pred["masks"] = ( self.process(proto[i], coefficient, pred["bboxes"], shape=imgsz) if coefficient.shape[0] else torch.zeros( (0, *(imgsz if self.process is ops.process_mask_native else proto.shape[2:])), dtype=torch.uint8, device=pred["bboxes"].device, ) ) return preds def _prepare_batch(self, si: int, batch: dict[str, Any]) -> dict[str, Any]: """Prepare a batch for validation by processing images and targets. Args: si (int): Sample index within the batch. batch (dict[str, Any]): Batch data containing images and annotations. Returns: (dict[str, Any]): Prepared batch with processed annotations. """ prepared_batch = super()._prepare_batch(si, batch) nl = prepared_batch["cls"].shape[0] if self.args.overlap_mask: masks = batch["masks"][si] index = torch.arange(1, nl + 1, device=masks.device).view(nl, 1, 1) masks = (masks == index).float() else: masks = batch["masks"][batch["batch_idx"] == si] if nl: mask_size = [s if self.process is ops.process_mask_native else s // 4 for s in prepared_batch["imgsz"]] if masks.shape[1:] != mask_size: masks = F.interpolate(masks[None], mask_size, mode="bilinear", align_corners=False)[0] masks = masks.gt_(0.5) prepared_batch["masks"] = masks return prepared_batch def _process_batch(self, preds: dict[str, torch.Tensor], batch: dict[str, Any]) -> dict[str, np.ndarray]: """Compute correct prediction matrix for a batch based on bounding boxes and optional masks. Args: preds (dict[str, torch.Tensor]): Dictionary containing predictions with keys like 'cls' and 'masks'. batch (dict[str, Any]): Dictionary containing batch data with keys like 'cls' and 'masks'. Returns: (dict[str, np.ndarray]): A dictionary containing correct prediction matrices including 'tp_m' for mask IoU. Examples: >>> preds = {"cls": torch.tensor([1, 0]), "masks": torch.rand(2, 640, 640), "bboxes": torch.rand(2, 4)} >>> batch = {"cls": torch.tensor([1, 0]), "masks": torch.rand(2, 640, 640), "bboxes": torch.rand(2, 4)} >>> correct_preds = validator._process_batch(preds, batch) Notes: - This method computes IoU between predicted and ground truth masks. - Overlapping masks are handled based on the overlap_mask argument setting. """ tp = super()._process_batch(preds, batch) gt_cls = batch["cls"] if gt_cls.shape[0] == 0 or preds["cls"].shape[0] == 0: tp_m = np.zeros((preds["cls"].shape[0], self.niou), dtype=bool) else: iou = mask_iou(batch["masks"].flatten(1), preds["masks"].flatten(1).float()) # float, uint8 tp_m = self.match_predictions(preds["cls"], gt_cls, iou).cpu().numpy() tp.update({"tp_m": tp_m}) # update tp with mask IoU return tp def plot_predictions(self, batch: dict[str, Any], preds: list[dict[str, torch.Tensor]], ni: int) -> None: """Plot batch predictions with masks and bounding boxes. Args: batch (dict[str, Any]): Batch containing images and annotations. preds (list[dict[str, torch.Tensor]]): List of predictions from the model. ni (int): Batch index. """ for p in preds: masks = p["masks"] if masks.shape[0] > self.args.max_det: LOGGER.warning(f"Limiting validation plots to 'max_det={self.args.max_det}' items.") p["masks"] = torch.as_tensor(masks[: self.args.max_det], dtype=torch.uint8).cpu() super().plot_predictions(batch, preds, ni, max_det=self.args.max_det) # plot bboxes def save_one_txt(self, predn: dict[str, torch.Tensor], save_conf: bool, shape: tuple[int, int], file: Path) -> None: """Save YOLO detections to a txt file in normalized coordinates in a specific format. Args: predn (dict[str, torch.Tensor]): Prediction dictionary containing 'bboxes', 'conf', 'cls', and 'masks' keys. save_conf (bool): Whether to save confidence scores. shape (tuple[int, int]): Shape of the original image. file (Path): File path to save the detections. """ from ultralytics.engine.results import Results Results( np.zeros((shape[0], shape[1]), dtype=np.uint8), path=None, names=self.names, boxes=torch.cat([predn["bboxes"], predn["conf"].unsqueeze(-1), predn["cls"].unsqueeze(-1)], dim=1), masks=torch.as_tensor(predn["masks"], dtype=torch.uint8), ).save_txt(file, save_conf=save_conf) def pred_to_json(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> None: """Save one JSON result for COCO evaluation. Args: predn (dict[str, torch.Tensor]): Predictions containing bboxes, masks, confidence scores, and classes. pbatch (dict[str, Any]): Batch dictionary containing 'imgsz', 'ori_shape', 'ratio_pad', and 'im_file'. """ def to_string(counts: list[int]) -> str: """Converts the RLE object into a compact string representation. Each count is delta-encoded and variable-length encoded as a string. Args: counts (list[int]): List of RLE counts. """ result = [] for i in range(len(counts)): x = int(counts[i]) # Apply delta encoding for all counts after the second entry if i > 2: x -= int(counts[i - 2]) # Variable-length encode the value while True: c = x & 0x1F # Take 5 bits x >>= 5 # If the sign bit (0x10) is set, continue if x != -1; # otherwise, continue if x != 0 more = (x != -1) if (c & 0x10) else (x != 0) if more: c |= 0x20 # Set continuation bit c += 48 # Shift to ASCII result.append(chr(c)) if not more: break return "".join(result) def multi_encode(pixels: torch.Tensor) -> list[int]: """Convert multiple binary masks using Run-Length Encoding (RLE). Args: pixels (torch.Tensor): A 2D tensor where each row represents a flattened binary mask with shape [N, H*W]. Returns: (list[list[int]]): A list of RLE counts for each mask. """ transitions = pixels[:, 1:] != pixels[:, :-1] row_idx, col_idx = torch.where(transitions) col_idx = col_idx + 1 # Compute run lengths counts = [] for i in range(pixels.shape[0]): positions = col_idx[row_idx == i] if len(positions): count = torch.diff(positions).tolist() count.insert(0, positions[0].item()) count.append(len(pixels[i]) - positions[-1].item()) else: count = [len(pixels[i])] # Ensure starting with background (0) count if pixels[i][0].item() == 1: count = [0, *count] counts.append(count) return counts pred_masks = predn["masks"].transpose(2, 1).contiguous().view(len(predn["masks"]), -1) # N, H*W h, w = predn["masks"].shape[1:3] counts = multi_encode(pred_masks) rles = [] for c in counts: rles.append({"size": [h, w], "counts": to_string(c)}) super().pred_to_json(predn, pbatch) for i, r in enumerate(rles): self.jdict[-len(rles) + i]["segmentation"] = r # segmentation def scale_preds(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> dict[str, torch.Tensor]: """Scales predictions to the original image size.""" return { **super().scale_preds(predn, pbatch), "masks": ops.scale_masks(predn["masks"][None], pbatch["ori_shape"], ratio_pad=pbatch["ratio_pad"])[ 0 ].byte(), } def eval_json(self, stats: dict[str, Any]) -> dict[str, Any]: """Return COCO-style instance segmentation evaluation metrics.""" pred_json = self.save_dir / "predictions.json" # predictions anno_json = ( self.data["path"] / "annotations" / ("instances_val2017.json" if self.is_coco else f"lvis_v1_{self.args.split}.json") ) # annotations return super().coco_evaluate(stats, pred_json, anno_json, ["bbox", "segm"], suffix=["Box", "Mask"])