Skip to content

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
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
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
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
300
301
302
303
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 = ops.nms_rotated(b, scores, 0.3)
            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
60
61
62
63
64
65
66
67
68
69
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

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
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
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
 97
 98
 99
100
101
102
103
104
105
106
107
108
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], filename: Union[str, Path]) -> 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
filename str | Path

Path to the image file for which predictions are being processed.

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
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
def pred_to_json(self, predn: Dict[str, torch.Tensor], filename: Union[str, Path]) -> 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.
        filename (str | Path): Path to the image file for which predictions are being processed.

    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.
    """
    stem = Path(filename).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,
                "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
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
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)





📅 Created 1 year ago ✏️ Updated 9 months ago