Bỏ để qua phần nội dung

Tài liệu tham khảo cho ultralytics/models/yolo/obb/predict.py

Ghi

Tệp này có sẵn tại https://github.com/ultralytics/ultralytics/blob/main/ultralytics/Mô hình/yolo/obb/predict.py. Nếu bạn phát hiện ra một vấn đề, vui lòng giúp khắc phục nó bằng cách đóng góp Yêu cầu 🛠️ kéo. Cảm ơn bạn 🙏 !



ultralytics.models.yolo.obb.predict.OBBPredictor

Căn cứ: DetectionPredictor

Một lớp mở rộng lớp DetectionPredictor để dự đoán dựa trên mô hình Hộp giới hạn định hướng (OBB).

Ví dụ
from ultralytics.utils import ASSETS
from ultralytics.models.yolo.obb import OBBPredictor

args = dict(model='yolov8n-obb.pt', source=ASSETS)
predictor = OBBPredictor(overrides=args)
predictor.predict_cli()
Mã nguồn trong ultralytics/models/yolo/obb/predict.py
class OBBPredictor(DetectionPredictor):
    """
    A class extending the DetectionPredictor class for prediction based on an Oriented Bounding Box (OBB) model.

    Example:
        ```python
        from ultralytics.utils import ASSETS
        from ultralytics.models.yolo.obb import OBBPredictor

        args = dict(model='yolov8n-obb.pt', source=ASSETS)
        predictor = OBBPredictor(overrides=args)
        predictor.predict_cli()
        ```
    """

    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
        """Initializes OBBPredictor with optional model and data configuration overrides."""
        super().__init__(cfg, overrides, _callbacks)
        self.args.task = "obb"

    def postprocess(self, preds, img, orig_imgs):
        """Post-processes predictions and returns a list of Results objects."""
        preds = ops.non_max_suppression(
            preds,
            self.args.conf,
            self.args.iou,
            agnostic=self.args.agnostic_nms,
            max_det=self.args.max_det,
            nc=len(self.model.names),
            classes=self.args.classes,
            rotated=True,
        )

        if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
            orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)

        results = []
        for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):
            rboxes = ops.regularize_rboxes(torch.cat([pred[:, :4], pred[:, -1:]], dim=-1))
            rboxes[:, :4] = ops.scale_boxes(img.shape[2:], rboxes[:, :4], orig_img.shape, xywh=True)
            # xywh, r, conf, cls
            obb = torch.cat([rboxes, pred[:, 4:6]], dim=-1)
            results.append(Results(orig_img, path=img_path, names=self.model.names, obb=obb))
        return results

__init__(cfg=DEFAULT_CFG, overrides=None, _callbacks=None)

Khởi tạo OBBPredictor với ghi đè cấu hình mô hình và dữ liệu tùy chọn.

Mã nguồn trong ultralytics/models/yolo/obb/predict.py
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """Initializes OBBPredictor with optional model and data configuration overrides."""
    super().__init__(cfg, overrides, _callbacks)
    self.args.task = "obb"

postprocess(preds, img, orig_imgs)

Xử lý hậu kỳ các dự đoán và trả về danh sách các đối tượng Kết quả.

Mã nguồn trong ultralytics/models/yolo/obb/predict.py
def postprocess(self, preds, img, orig_imgs):
    """Post-processes predictions and returns a list of Results objects."""
    preds = ops.non_max_suppression(
        preds,
        self.args.conf,
        self.args.iou,
        agnostic=self.args.agnostic_nms,
        max_det=self.args.max_det,
        nc=len(self.model.names),
        classes=self.args.classes,
        rotated=True,
    )

    if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
        orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)

    results = []
    for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):
        rboxes = ops.regularize_rboxes(torch.cat([pred[:, :4], pred[:, -1:]], dim=-1))
        rboxes[:, :4] = ops.scale_boxes(img.shape[2:], rboxes[:, :4], orig_img.shape, xywh=True)
        # xywh, r, conf, cls
        obb = torch.cat([rboxes, pred[:, 4:6]], dim=-1)
        results.append(Results(orig_img, path=img_path, names=self.model.names, obb=obb))
    return results





Created 2024-01-05, Updated 2024-06-02
Authors: glenn-jocher (4), Burhan-Q (1)