# Ultralytics YOLO 🚀, AGPL-3.0 license import torch from ultralytics.data import YOLODataset, converter from ultralytics.data.augment import Compose, Format, v8_transforms from ultralytics.models.yolo.detect import DetectionValidator from ultralytics.utils import colorstr, ops from ultralytics.utils.metrics import DetMetrics, SegmentMetrics, ConfusionMatrix, SegmentationMetric, AverageMeter from ultralytics.utils import LOGGER from ultralytics.utils.plotting import output_to_target, plot_images import os from pathlib import Path import numpy as np import torch.nn.functional as F __all__ = ("MTDETRValidator",) # tuple or list class MTDETRDataset(YOLODataset): """ Real-Time DEtection and TRacking (RT-DETR) dataset class extending the base YOLODataset class. This specialized dataset class is designed for use with the RT-DETR object detection model and is optimized for real-time detection and tracking tasks. """ def __init__(self, *args, data=None, **kwargs): """Initialize the RTDETRDataset class by inheriting from the YOLODataset class.""" super().__init__(*args, data=data, **kwargs) # NOTE: add stretch version load_image for RTDETR mosaic def load_image(self, i, rect_mode=False): """Loads 1 image from dataset index 'i', returns (im, resized hw).""" return super().load_image(i=i, rect_mode=rect_mode) def build_transforms(self, hyp=None): """Temporary, only for evaluation.""" if self.augment: hyp.mosaic = hyp.mosaic if self.augment and not self.rect else 0.0 hyp.mixup = hyp.mixup if self.augment and not self.rect else 0.0 transforms = v8_transforms(self, self.imgsz, hyp, stretch=True) else: # transforms = Compose([LetterBox(new_shape=(self.imgsz, self.imgsz), auto=False, scaleFill=True)]) transforms = Compose([]) transforms.append( Format( bbox_format="xywh", normalize=True, return_mask=self.use_segments, return_keypoint=self.use_keypoints, batch_idx=True, mask_ratio=hyp.mask_ratio, mask_overlap=hyp.overlap_mask, merge_mask=True, imgsz=self.imgsz, seg_nc=len(self.data['type_task']['segmentation']), ) ) return transforms class MTDETRValidator(DetectionValidator): """ RTDETRValidator extends the DetectionValidator class to provide validation capabilities specifically tailored for the RT-DETR (Real-Time DETR) object detection model. The class allows building of an RTDETR-specific dataset for validation, applies Non-maximum suppression for post-processing, and updates evaluation metrics accordingly. Example: ```python from ultralytics.models.rtdetr import RTDETRValidator args = dict(model="rtdetr-l.pt", data="coco8.yaml") validator = RTDETRValidator(args=args) validator() ``` Note: For further details on the attributes and methods, refer to the parent DetectionValidator class. """ def __init__(self, dataloader=None, number_task=None, save_dir=None, pbar=None, args=None, _callbacks=None): """Initialize SegmentationValidator and set task to 'segment', metrics to SegmentMetrics.""" super().__init__(dataloader, save_dir, pbar, args, _callbacks) self.number_task = number_task self.plot_masks = None self.process = None self.args.task = "multi" ### JW init two metics for different task. self.metrics = {'detection': DetMetrics(save_dir=self.save_dir, on_plot=self.on_plot), 'segmentation': SegmentMetrics(save_dir=self.save_dir, on_plot=self.on_plot)} if not isinstance(self.args.mask_threshold, list): self.args.mask_threshold = [self.args.mask_threshold] self.mask_thr = torch.tensor(self.args.mask_threshold).view(1, len(self.args.mask_threshold), 1, 1) ### JW set the mask_threshold def init_metrics(self, model): """Initialize evaluation metrics for YOLO.""" val = self.data.get(self.args.split, "") # validation path self.is_coco = ( isinstance(val, str) and "coco" in val and (val.endswith(f"{os.sep}val2017.txt") or val.endswith(f"{os.sep}test-dev2017.txt")) ) # is COCO self.is_lvis = isinstance(val, str) and "lvis" in val and not self.is_coco # is LVIS self.class_map = converter.coco80_to_coco91_class() if self.is_coco else list(range(len(model.names))) self.args.save_json |= (self.is_coco or self.is_lvis) and not self.training # run on final val if training COCO self.names = model.names self.nc = self.number_task self.filtered_names = {task: {i: self.names[i] for i in indices if i in self.names} for task, indices in self.data['type_task'].items()} self.metrics['detection'].names = self.filtered_names['detection'] self.metrics['detection'].plot = self.args.plots self.metrics['segmentation'].names = self.filtered_names['segmentation'] self.metrics['segmentation'].plot = self.args.plots ### JW init the segmentation task self.seg_metrics = {self.dataloader.dataset.data['names'][i]: SegmentationMetric(2) for i in self.dataloader.dataset.data['type_task']['segmentation']} self.seg_result = {self.dataloader.dataset.data['names'][i]: {'pixacc': AverageMeter(), 'subacc': AverageMeter(), 'IoU': AverageMeter(), 'mIoU': AverageMeter()} for i in self.dataloader.dataset.data['type_task']['segmentation']} self.plot_masks = {self.dataloader.dataset.data['names'][i]: [] for i in self.dataloader.dataset.data['type_task']['segmentation']} self.confusion_matrix_detection = ConfusionMatrix(nc=self.nc['detection'], conf=self.args.conf) self.confusion_matrix_segmentation = ConfusionMatrix(nc=self.nc['segmentation'], conf=self.args.conf) self.seen = 0 self.jdict = [] self.stats = dict(tp=[], conf=[], pred_cls=[], target_cls=[], target_img=[]) def preprocess(self, batch): """Preprocesses batch by converting masks to float and sending to device.""" batch = super().preprocess(batch) # batch["masks"] = batch["masks"].to(self.device).float() return batch def build_dataset(self, img_path, mode="val", batch=None): """ Build an RTDETR Dataset. Args: img_path (str): Path to the folder containing images. mode (str): `train` mode or `val` mode, users are able to customize different augmentations for each mode. batch (int, optional): Size of batches, this is for `rect`. Defaults to None. """ return MTDETRDataset( img_path=img_path, imgsz=self.args.imgsz, batch_size=batch, augment=False, # no augmentation hyp=self.args, rect=False, # no rect cache=self.args.cache or None, prefix=colorstr(f"{mode}: "), data=self.data, ) def postprocess(self, preds): """Apply Non-maximum suppression to prediction outputs.""" if not isinstance(preds, (list, tuple)): # list for PyTorch inference but list[0] Tensor for export inference preds = [preds, None] bs, _, nd = preds[0].shape bboxes, scores = preds[0].split((4, nd - 4), dim=-1) bboxes *= self.args.imgsz outputs = [torch.zeros((0, 6), device=bboxes.device)] * bs for i, bbox in enumerate(bboxes): # (300, 4) bbox = ops.xywh2xyxy(bbox) score, cls = scores[i].max(-1) # (300, ) # Do not need threshold for evaluation as only got 300 boxes here # idx = score > self.args.conf pred = torch.cat([bbox, score[..., None], cls[..., None]], dim=-1) # filter # Sort by confidence to correctly get internal metrics pred = pred[score.argsort(descending=True)] outputs[i] = pred # [idx] mask_thr_tensor = self.mask_thr.to(device=preds[1][5].device) mask = (torch.sigmoid(preds[1][5]) > mask_thr_tensor).float() return outputs, mask def _prepare_batch(self, si, batch, detection_indices): """Prepares a batch for training or inference by applying transformations.""" idx_detection = detection_indices[batch["batch_idx"][detection_indices] == si] cls_detection = batch["cls"][idx_detection].squeeze(-1) bbox_detection = batch["bboxes"][idx_detection] ori_shape = batch["ori_shape"][si] imgsz = batch["img"].shape[2:] ratio_pad = batch["ratio_pad"][si] if len(cls_detection): bbox_detection = ops.xywh2xyxy(bbox_detection) # target boxes bbox_detection[..., [0, 2]] *= ori_shape[1] # native-space pred bbox_detection[..., [1, 3]] *= ori_shape[0] # native-space pred return {"cls": cls_detection, "bbox": bbox_detection, "ori_shape": ori_shape, "imgsz": imgsz, "ratio_pad": ratio_pad} def _prepare_pred(self, pred, pbatch): """Prepares and returns a batch with transformed bounding boxes and class labels.""" predn = pred.clone() predn[..., [0, 2]] *= pbatch["ori_shape"][1] / self.args.imgsz # native-space pred predn[..., [1, 3]] *= pbatch["ori_shape"][0] / self.args.imgsz # native-space pred return predn.float() def update_metrics(self, preds_list, batch): """Metrics.""" preds = preds_list[0] merge_mask = preds_list[1] type_task = batch['type_task'][0] detection_classes = torch.tensor(type_task['detection']) detection_indices = torch.where(torch.isin(batch["cls"], detection_classes.to(device=batch["cls"].device)))[0] for si, pred in enumerate(preds): self.seen += 1 npr = len(pred) stat = dict( conf=torch.zeros(0, device=self.device), pred_cls=torch.zeros(0, device=self.device), tp=torch.zeros(npr, self.niou, dtype=torch.bool, device=self.device), ) pbatch_detection= self._prepare_batch(si, batch, detection_indices) cls, bbox = pbatch_detection.pop("cls"), pbatch_detection.pop("bbox") nl = len(cls) stat["target_cls"] = cls stat["target_img"] = cls.unique() if npr == 0: if nl: for k in self.stats.keys(): self.stats[k].append(stat[k]) if self.args.plots: self.confusion_matrix.process_batch(detections=None, gt_bboxes=bbox, gt_cls=cls) continue # Predictions if self.args.single_cls: pred[:, 5] = 0 predn = self._prepare_pred(pred, pbatch_detection) stat["conf"] = predn[:, 4] stat["pred_cls"] = predn[:, 5] ### JW segmentataion evaluate _, nc, _, _ = merge_mask.shape for seg_nc in range(nc): task_name = self.data['names'][self.data['type_task']['segmentation'][seg_nc]] pred_mask = merge_mask[si][seg_nc].squeeze() gt_mask = batch['merge_mask'][si][seg_nc].to(device=pred_mask.device) self.seg_metrics[task_name].reset() if gt_mask.shape != pred_mask.shape: pred_mask = pred_mask.unsqueeze(0).unsqueeze(0) # (1, 1, 64, 64) gt_mask = gt_mask.unsqueeze(0).unsqueeze(0) # (1, 1, 320, 320) pred_mask = F.interpolate( pred_mask, size=gt_mask.shape[-2:], mode='bilinear', ) pred_mask = pred_mask.squeeze(0).squeeze(0) gt_mask = gt_mask.squeeze(0).squeeze(0) self.seg_metrics[task_name].addBatch(pred_mask.cpu(), gt_mask.cpu()) self.seg_result[task_name]['pixacc'].update(self.seg_metrics[task_name].pixelAccuracy()) self.seg_result[task_name]['subacc'].update(self.seg_metrics[task_name].lineAccuracy()) self.seg_result[task_name]['IoU'].update(self.seg_metrics[task_name].IntersectionOverUnion()) self.seg_result[task_name]['mIoU'].update(self.seg_metrics[task_name].meanIntersectionOverUnion()) if self.args.plots: self.plot_masks[task_name].append(pred_mask.cpu()) ### JW TODO Need to adapt to the segmentation task plot # Evaluate if nl: stat["tp"] = self._process_batch(predn, bbox, cls) if self.args.plots: self.confusion_matrix_detection.process_batch(predn, bbox, cls) for k in self.stats.keys(): self.stats[k].append(stat[k]) # Save if self.args.save_json: self.pred_to_json(predn, batch["im_file"][si]) if self.args.save_txt: self.save_one_txt( predn, self.args.save_conf, pbatch_detection["ori_shape"], self.save_dir / "labels" / f'{Path(batch["im_file"][si]).stem}.txt', ) def get_stats(self): """Returns metrics statistics and results dictionary.""" stats = {k: torch.cat(v, 0).cpu().numpy() for k, v in self.stats.items()} # to numpy self.nt_per_class = np.bincount(stats["target_cls"].astype(int), minlength=self.nc['detection']) self.nt_per_image = np.bincount(stats["target_img"].astype(int), minlength=self.nc['detection']) stats.pop("target_img", None) if len(stats) and stats["tp"].any(): self.metrics['detection'].process(**stats) return self.metrics['detection'].results_dict def finalize_metrics(self, *args, **kwargs): """Set final values for metrics speed and confusion matrix.""" self.metrics['detection'].confusion_matrix = self.confusion_matrix def print_results(self): """Prints training/validation set metrics per class.""" pf = "%22s" + "%11i" * 2 + "%11.3g" * len(self.metrics['detection'].keys) # print format LOGGER.info(pf % ("all", self.seen, self.nt_per_class.sum(), *self.metrics['detection'].mean_results())) if self.nt_per_class.sum() == 0: LOGGER.warning(f"WARNING ⚠️ no labels found in {self.args.task} set, can not compute metrics without labels") # Print results per class if self.args.verbose and not self.training and self.nc['detection'] > 1 and len(self.stats): for i, c in enumerate(self.metrics.ap_class_index): LOGGER.info( pf % (self.names[c], self.nt_per_image[c], self.nt_per_class[c], *self.metrics['detection'].class_result(i)) ) pf = '%22s' + ('%11s' + '%11.3g') * 4 seg_index = self.data['type_task']['segmentation'] for i in seg_index: key_values = [(key, value.avg) for key, value in self.seg_result[self.data['names'][i]].items()] LOGGER.info(pf % (self.data['names'][i], *sum(key_values, ()))) if self.args.plots: for normalize in True, False: self.confusion_matrix_detection.plot( save_dir=self.save_dir, names=self.names.values(), normalize=normalize, on_plot=self.on_plot ) def plot_val_samples(self, batch, ni): """Plot validation image samples.""" plot_images( batch["img"], batch["batch_idx"], batch["cls"].squeeze(-1), batch["bboxes"], paths=batch["im_file"], fname=self.save_dir / f"val_batch{ni}_labels.jpg", names=self.names, on_plot=self.on_plot, ) def plot_predictions(self, batch, preds, ni): """Plots predicted bounding boxes on input images and saves the result.""" plot_images( batch["img"], *output_to_target(preds, max_det=self.args.max_det), paths=batch["im_file"], fname=self.save_dir / f"val_batch{ni}_pred.jpg", names=self.names, on_plot=self.on_plot, ) # pred