Skip to content

Register now for

# Reference for `ultralytics/models/yolo/obb/val.py`

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/yolo/obb/val.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!

## `` ultralytics.models.yolo.obb.val.OBBValidator

    OBBValidator(dataloader=None, save_dir=None, args=None, _callbacks=None)

Bases: `DetectionValidator`

A class extending the DetectionValidator class for validation based on an Oriented Bounding Box (OBB) model.

This validator specializes in evaluating models that predict rotated bounding boxes, commonly used for aerial and satellite imagery where objects can appear at various orientations.

Attributes:

Name | Type | Description
--- | --- | ---
`args` | `dict` | Configuration arguments for the validator.
`metrics` | `OBBMetrics` | Metrics object for evaluating OBB model performance.
`is_dota` | `bool` | Flag indicating whether the validation dataset is in DOTA format.

Methods:

Name | Description
--- | ---
`init_metrics` | Initialize evaluation metrics for YOLO.
`_process_batch` | Process batch of detections and ground truth boxes to compute IoU matrix.
`_prepare_batch` | Prepare batch data for OBB validation.
`_prepare_pred` | Prepare predictions with scaled and padded bounding boxes.
`plot_predictions` | Plot predicted bounding boxes on input images.
`pred_to_json` | Serialize YOLO predictions to COCO json format.
`save_one_txt` | Save YOLO detections to a txt file in normalized coordinates.
`eval_json` | Evaluate YOLO output in JSON format and return performance statistics.

Examples:

    >>> from ultralytics.models.yolo.obb import OBBValidator
    >>> args = dict(model="yolo11n-obb.pt", data="dota8.yaml")
    >>> validator = OBBValidator(args=args)
    >>> validator(model=args["model"])

This constructor initializes an OBBValidator instance for validating Oriented Bounding Box (OBB) models. It extends the DetectionValidator class and configures it specifically for the OBB task.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`dataloader` | `DataLoader` | Dataloader to be used for validation. | `None`
`save_dir` | `str | Path` | Directory to save results. | `None`
`args` | `dict | SimpleNamespace` | Arguments containing validation parameters. | `None`
`_callbacks` | `list` | List of callback functions to be called during validation. | `None`

Source code in `ultralytics/models/yolo/obb/val.py`

    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    def __init__(self, dataloader=None, save_dir=None, args=None, _callbacks=None) -> None:
        """
        Initialize OBBValidator and set task to 'obb', metrics to OBBMetrics.

        This constructor initializes an OBBValidator instance for validating Oriented Bounding Box (OBB) models.
        It extends the DetectionValidator class and configures it specifically for the OBB task.

        Args:
            dataloader (torch.utils.data.DataLoader, optional): Dataloader to be used for validation.
            save_dir (str | Path, optional): Directory to save results.
            args (dict | SimpleNamespace, optional): Arguments containing validation parameters.
            _callbacks (list, optional): List of callback functions to be called during validation.
        """
        super().__init__(dataloader, save_dir, args, _callbacks)
        self.args.task = "obb"
        self.metrics = OBBMetrics()

### `` eval_json

    eval_json(stats: dict[str, Any]) -> dict[str, Any]

Evaluate YOLO output in JSON format and save predictions in DOTA format.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`stats` | `dict[str, Any]` | Performance statistics dictionary. | required

Returns:

Type | Description
--- | ---
`dict[str, Any]` | Updated performance statistics.

Source code in `ultralytics/models/yolo/obb/val.py`

    235
    236
    237
    238
    239
    240
    241
    242
    243
    244
    245
    246
    247
    248
    249
    250
    251
    252
    253
    254
    255
    256
    257
    258
    259
    260
    261
    262
    263
    264
    265
    266
    267
    268
    269
    270
    271
    272
    273
    274
    275
    276
    277
    278
    279
    280
    281
    282
    283
    284
    285
    286
    287
    288
    289
    290
    291
    292
    293
    294
    295
    296
    297
    298
    299
    def eval_json(self, stats: dict[str, Any]) -> dict[str, Any]:
        """
        Evaluate YOLO output in JSON format and save predictions in DOTA format.

        Args:
            stats (dict[str, Any]): Performance statistics dictionary.

        Returns:
            (dict[str, Any]): Updated performance statistics.
        """
        if self.args.save_json and self.is_dota and len(self.jdict):
            import json
            import re
            from collections import defaultdict

            pred_json = self.save_dir / "predictions.json"  # predictions
            pred_txt = self.save_dir / "predictions_txt"  # predictions
            pred_txt.mkdir(parents=True, exist_ok=True)
            data = json.load(open(pred_json))
            # Save split results
            LOGGER.info(f"Saving predictions with DOTA format to {pred_txt}...")
            for d in data:
                image_id = d["image_id"]
                score = d["score"]
                classname = self.names[d["category_id"] - 1].replace(" ", "-")
                p = d["poly"]

                with open(f"{pred_txt / f'Task1_{classname}'}.txt", "a", encoding="utf-8") as f:
                    f.writelines(f"{image_id} {score} {p[0]} {p[1]} {p[2]} {p[3]} {p[4]} {p[5]} {p[6]} {p[7]}\n")
            # Save merged results, this could result slightly lower map than using official merging script,
            # because of the probiou calculation.
            pred_merged_txt = self.save_dir / "predictions_merged_txt"  # predictions
            pred_merged_txt.mkdir(parents=True, exist_ok=True)
            merged_results = defaultdict(list)
            LOGGER.info(f"Saving merged predictions with DOTA format to {pred_merged_txt}...")
            for d in data:
                image_id = d["image_id"].split("__", 1)[0]
                pattern = re.compile(r"\d+___\d+")
                x, y = (int(c) for c in re.findall(pattern, d["image_id"])[0].split("___"))
                bbox, score, cls = d["rbox"], d["score"], d["category_id"] - 1
                bbox[0] += x
                bbox[1] += y
                bbox.extend([score, cls])
                merged_results[image_id].append(bbox)
            for image_id, bbox in merged_results.items():
                bbox = torch.tensor(bbox)
                max_wh = torch.max(bbox[:, :2]).item() * 2
                c = bbox[:, 6:7] * max_wh  # classes
                scores = bbox[:, 5]  # scores
                b = bbox[:, :5].clone()
                b[:, :2] += c
                # 0.3 could get results close to the ones from official merging script, even slightly better.
                i = TorchNMS.fast_nms(b, scores, 0.3, iou_func=batch_probiou)
                bbox = bbox[i]

                b = ops.xywhr2xyxyxyxy(bbox[:, :5]).view(-1, 8)
                for x in torch.cat([b, bbox[:, 5:7]], dim=-1).tolist():
                    classname = self.names[int(x[-1])].replace(" ", "-")
                    p = [round(i, 3) for i in x[:-2]]  # poly
                    score = round(x[-2], 3)

                    with open(f"{pred_merged_txt / f'Task1_{classname}'}.txt", "a", encoding="utf-8") as f:
                        f.writelines(f"{image_id} {score} {p[0]} {p[1]} {p[2]} {p[3]} {p[4]} {p[5]} {p[6]} {p[7]}\n")

        return stats

### `` init_metrics

    init_metrics(model: Module) -> None

Initialize evaluation metrics for YOLO obb validation.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`model` | `Module` | Model to validate. | required

Source code in `ultralytics/models/yolo/obb/val.py`

    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    def init_metrics(self, model: torch.nn.Module) -> None:
        """
        Initialize evaluation metrics for YOLO obb validation.

        Args:
            model (torch.nn.Module): Model to validate.
        """
        super().init_metrics(model)
        val = self.data.get(self.args.split, "")  # validation path
        self.is_dota = isinstance(val, str) and "DOTA" in val  # check if dataset is DOTA format
        self.confusion_matrix.task = "obb"  # set confusion matrix task to 'obb'

### `` plot_predictions

    plot_predictions(batch: dict[str, Any], preds: list[Tensor], ni: int) -> None

Plot predicted bounding boxes on input images and save the result.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`batch` | `dict[str, Any]` | Batch data containing images, file paths, and other metadata. | required
`preds` | `list[Tensor]` | List of prediction tensors for each image in the batch. | required
`ni` | `int` | Batch index used for naming the output file. | required

Examples:

    >>> validator = OBBValidator()
    >>> batch = {"img": images, "im_file": paths}
    >>> preds = [torch.rand(10, 7)]  # Example predictions for one image
    >>> validator.plot_predictions(batch, preds, 0)

Source code in `ultralytics/models/yolo/obb/val.py`

    148
    149
    150
    151
    152
    153
    154
    155
    156
    157
    158
    159
    160
    161
    162
    163
    164
    165
    166
    def plot_predictions(self, batch: dict[str, Any], preds: list[torch.Tensor], ni: int) -> None:
        """
        Plot predicted bounding boxes on input images and save the result.

        Args:
            batch (dict[str, Any]): Batch data containing images, file paths, and other metadata.
            preds (list[torch.Tensor]): List of prediction tensors for each image in the batch.
            ni (int): Batch index used for naming the output file.

        Examples:
            >>> validator = OBBValidator()
            >>> batch = {"img": images, "im_file": paths}
            >>> preds = [torch.rand(10, 7)]  # Example predictions for one image
            >>> validator.plot_predictions(batch, preds, 0)
        """
        for p in preds:
            # TODO: fix this duplicated `xywh2xyxy`
            p["bboxes"][:, :4] = ops.xywh2xyxy(p["bboxes"][:, :4])  # convert to xyxy format for plotting
        super().plot_predictions(batch, preds, ni)  # plot bboxes

### `` postprocess

    postprocess(preds: Tensor) -> list[dict[str, torch.Tensor]]

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`preds` | `Tensor` | Raw predictions from the model. | required

Returns:

Type | Description
--- | ---
`list[dict[str, Tensor]]` | Processed predictions with angle information concatenated to bboxes.

Source code in `ultralytics/models/yolo/obb/val.py`

    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    def postprocess(self, preds: torch.Tensor) -> list[dict[str, torch.Tensor]]:
        """
        Args:
            preds (torch.Tensor): Raw predictions from the model.

        Returns:
            (list[dict[str, torch.Tensor]]): Processed predictions with angle information concatenated to bboxes.
        """
        preds = super().postprocess(preds)
        for pred in preds:
            pred["bboxes"] = torch.cat([pred["bboxes"], pred.pop("extra")], dim=-1)  # concatenate angle
        return preds

### `` pred_to_json

    pred_to_json(predn: dict[str, Tensor], pbatch: dict[str, Any]) -> None

Convert YOLO predictions to COCO JSON format with rotated bounding box information.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`predn` | `dict[str, Tensor]` | Prediction dictionary containing 'bboxes', 'conf', and 'cls' keys with bounding box coordinates, confidence scores, and class predictions. | required
`pbatch` | `dict[str, Any]` | Batch dictionary containing 'imgsz', 'ori_shape', 'ratio_pad', and 'im_file'. | required

Notes

This method processes rotated bounding box predictions and converts them to both rbox format (x, y, w, h, angle) and polygon format (x1, y1, x2, y2, x3, y3, x4, y4) before adding them to the JSON dictionary.

Source code in `ultralytics/models/yolo/obb/val.py`

    168
    169
    170
    171
    172
    173
    174
    175
    176
    177
    178
    179
    180
    181
    182
    183
    184
    185
    186
    187
    188
    189
    190
    191
    192
    193
    194
    195
    196
    197
    def pred_to_json(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> None:
        """
        Convert YOLO predictions to COCO JSON format with rotated bounding box information.

        Args:
            predn (dict[str, torch.Tensor]): Prediction dictionary containing 'bboxes', 'conf', and 'cls' keys
                with bounding box coordinates, confidence scores, and class predictions.
            pbatch (dict[str, Any]): Batch dictionary containing 'imgsz', 'ori_shape', 'ratio_pad', and 'im_file'.

        Notes:
            This method processes rotated bounding box predictions and converts them to both rbox format
            (x, y, w, h, angle) and polygon format (x1, y1, x2, y2, x3, y3, x4, y4) before adding them
            to the JSON dictionary.
        """
        path = Path(pbatch["im_file"])
        stem = path.stem
        image_id = int(stem) if stem.isnumeric() else stem
        rbox = predn["bboxes"]
        poly = ops.xywhr2xyxyxyxy(rbox).view(-1, 8)
        for r, b, s, c in zip(rbox.tolist(), poly.tolist(), predn["conf"].tolist(), predn["cls"].tolist()):
            self.jdict.append(
                {
                    "image_id": image_id,
                    "file_name": path.name,
                    "category_id": self.class_map[int(c)],
                    "score": round(s, 5),
                    "rbox": [round(x, 3) for x in r],
                    "poly": [round(x, 3) for x in b],
                }
            )

### `` save_one_txt

    save_one_txt(
        predn: dict[str, Tensor],
        save_conf: bool,
        shape: tuple[int, int],
        file: Path,
    ) -> None

Save YOLO OBB detections to a text file in normalized coordinates.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`predn` | `Tensor` | Predicted detections with shape (N, 7) containing bounding boxes, confidence scores, class predictions, and angles in format (x, y, w, h, conf, cls, angle). | required
`save_conf` | `bool` | Whether to save confidence scores in the text file. | required
`shape` | `tuple[int, int]` | Original image shape in format (height, width). | required
`file` | `Path` | Output file path to save detections. | required

Examples:

    >>> validator = OBBValidator()
    >>> predn = torch.tensor([[100, 100, 50, 30, 0.9, 0, 45]])  # One detection: x,y,w,h,conf,cls,angle
    >>> validator.save_one_txt(predn, True, (640, 480), "detection.txt")

Source code in `ultralytics/models/yolo/obb/val.py`

    199
    200
    201
    202
    203
    204
    205
    206
    207
    208
    209
    210
    211
    212
    213
    214
    215
    216
    217
    218
    219
    220
    221
    222
    223
    224
    def save_one_txt(self, predn: dict[str, torch.Tensor], save_conf: bool, shape: tuple[int, int], file: Path) -> None:
        """
        Save YOLO OBB detections to a text file in normalized coordinates.

        Args:
            predn (torch.Tensor): Predicted detections with shape (N, 7) containing bounding boxes, confidence scores,
                class predictions, and angles in format (x, y, w, h, conf, cls, angle).
            save_conf (bool): Whether to save confidence scores in the text file.
            shape (tuple[int, int]): Original image shape in format (height, width).
            file (Path): Output file path to save detections.

        Examples:
            >>> validator = OBBValidator()
            >>> predn = torch.tensor([[100, 100, 50, 30, 0.9, 0, 45]])  # One detection: x,y,w,h,conf,cls,angle
            >>> validator.save_one_txt(predn, True, (640, 480), "detection.txt")
        """
        import numpy as np

        from ultralytics.engine.results import Results

        Results(
            np.zeros((shape[0], shape[1]), dtype=np.uint8),
            path=None,
            names=self.names,
            obb=torch.cat([predn["bboxes"], predn["conf"].unsqueeze(-1), predn["cls"].unsqueeze(-1)], dim=1),
        ).save_txt(file, save_conf=save_conf)

### `` scale_preds

    scale_preds(
        predn: dict[str, Tensor], pbatch: dict[str, Any]
    ) -> dict[str, torch.Tensor]

Scales predictions to the original image size.

Source code in `ultralytics/models/yolo/obb/val.py`

    226
    227
    228
    229
    230
    231
    232
    233
    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 {
            **predn,
            "bboxes": ops.scale_boxes(
                pbatch["imgsz"], predn["bboxes"].clone(), pbatch["ori_shape"], ratio_pad=pbatch["ratio_pad"], xywh=True
            ),
        }

📅 Created 1 year ago ✏️ Updated 1 year ago