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

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

Ghi

Tệp này có sẵn tại https://github.com/ultralytics/ultralytics/blob/main/ultralytics/mô hình/rtdetr/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.rtdetr.predict.RTDETRPredictor

Căn cứ: BasePredictor

RT-DETR (Máy biến áp phát hiện thời gian thực) Predictor mở rộng lớp BasePredictor để đưa ra dự đoán bằng cách sử dụng Baidu của RT-DETR mẫu.

Lớp này tận dụng sức mạnh của Vision Transformers để cung cấp khả năng phát hiện đối tượng theo thời gian thực trong khi vẫn duy trì độ chính xác cao. Nó hỗ trợ các tính năng chính như mã hóa lai hiệu quả và lựa chọn truy vấn nhận biết IoU.

Ví dụ
from ultralytics.utils import ASSETS
from ultralytics.models.rtdetr import RTDETRPredictor

args = dict(model='rtdetr-l.pt', source=ASSETS)
predictor = RTDETRPredictor(overrides=args)
predictor.predict_cli()

Thuộc tính:

Tên Kiểu Sự miêu tả
imgsz int

Kích thước hình ảnh để suy luận (phải vuông và điền tỷ lệ).

args dict

Đối số ghi đè cho công cụ dự đoán.

Mã nguồn trong ultralytics/models/rtdetr/predict.py
class RTDETRPredictor(BasePredictor):
    """
    RT-DETR (Real-Time Detection Transformer) Predictor extending the BasePredictor class for making predictions using
    Baidu's RT-DETR model.

    This class leverages the power of Vision Transformers to provide real-time object detection while maintaining
    high accuracy. It supports key features like efficient hybrid encoding and IoU-aware query selection.

    Example:
        ```python
        from ultralytics.utils import ASSETS
        from ultralytics.models.rtdetr import RTDETRPredictor

        args = dict(model='rtdetr-l.pt', source=ASSETS)
        predictor = RTDETRPredictor(overrides=args)
        predictor.predict_cli()
        ```

    Attributes:
        imgsz (int): Image size for inference (must be square and scale-filled).
        args (dict): Argument overrides for the predictor.
    """

    def postprocess(self, preds, img, orig_imgs):
        """
        Postprocess the raw predictions from the model to generate bounding boxes and confidence scores.

        The method filters detections based on confidence and class if specified in `self.args`.

        Args:
            preds (list): List of [predictions, extra] from the model.
            img (torch.Tensor): Processed input images.
            orig_imgs (list or torch.Tensor): Original, unprocessed images.

        Returns:
            (list[Results]): A list of Results objects containing the post-processed bounding boxes, confidence scores,
                and class labels.
        """
        if not isinstance(preds, (list, tuple)):  # list for PyTorch inference but list[0] Tensor for export inference
            preds = [preds, None]

        nd = preds[0].shape[-1]
        bboxes, scores = preds[0].split((4, nd - 4), dim=-1)

        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 i, bbox in enumerate(bboxes):  # (300, 4)
            bbox = ops.xywh2xyxy(bbox)
            score, cls = scores[i].max(-1, keepdim=True)  # (300, 1)
            idx = score.squeeze(-1) > self.args.conf  # (300, )
            if self.args.classes is not None:
                idx = (cls == torch.tensor(self.args.classes, device=cls.device)).any(1) & idx
            pred = torch.cat([bbox, score, cls], dim=-1)[idx]  # filter
            orig_img = orig_imgs[i]
            oh, ow = orig_img.shape[:2]
            pred[..., [0, 2]] *= ow
            pred[..., [1, 3]] *= oh
            img_path = self.batch[0][i]
            results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
        return results

    def pre_transform(self, im):
        """
        Pre-transforms the input images before feeding them into the model for inference. The input images are
        letterboxed to ensure a square aspect ratio and scale-filled. The size must be square(640) and scaleFilled.

        Args:
            im (list[np.ndarray] |torch.Tensor): Input images of shape (N,3,h,w) for tensor, [(h,w,3) x N] for list.

        Returns:
            (list): List of pre-transformed images ready for model inference.
        """
        letterbox = LetterBox(self.imgsz, auto=False, scaleFill=True)
        return [letterbox(image=x) for x in im]

postprocess(preds, img, orig_imgs)

Xử lý hậu kỳ các dự đoán thô từ mô hình để tạo ra các hộp giới hạn và điểm tin cậy.

Phương pháp lọc phát hiện dựa trên độ tin cậy và lớp nếu được chỉ định trong self.args.

Thông số:

Tên Kiểu Sự miêu tả Mặc định
preds list

Danh sách [dự đoán, bổ sung] từ mô hình.

bắt buộc
img Tensor

Hình ảnh đầu vào đã xử lý.

bắt buộc
orig_imgs list or Tensor

Hình ảnh gốc, chưa qua xử lý.

bắt buộc

Trở lại:

Kiểu Sự miêu tả
list[Results]

Danh sách các đối tượng Kết quả chứa các hộp giới hạn sau xử lý, điểm tin cậy, và nhãn lớp.

Mã nguồn trong ultralytics/models/rtdetr/predict.py
def postprocess(self, preds, img, orig_imgs):
    """
    Postprocess the raw predictions from the model to generate bounding boxes and confidence scores.

    The method filters detections based on confidence and class if specified in `self.args`.

    Args:
        preds (list): List of [predictions, extra] from the model.
        img (torch.Tensor): Processed input images.
        orig_imgs (list or torch.Tensor): Original, unprocessed images.

    Returns:
        (list[Results]): A list of Results objects containing the post-processed bounding boxes, confidence scores,
            and class labels.
    """
    if not isinstance(preds, (list, tuple)):  # list for PyTorch inference but list[0] Tensor for export inference
        preds = [preds, None]

    nd = preds[0].shape[-1]
    bboxes, scores = preds[0].split((4, nd - 4), dim=-1)

    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 i, bbox in enumerate(bboxes):  # (300, 4)
        bbox = ops.xywh2xyxy(bbox)
        score, cls = scores[i].max(-1, keepdim=True)  # (300, 1)
        idx = score.squeeze(-1) > self.args.conf  # (300, )
        if self.args.classes is not None:
            idx = (cls == torch.tensor(self.args.classes, device=cls.device)).any(1) & idx
        pred = torch.cat([bbox, score, cls], dim=-1)[idx]  # filter
        orig_img = orig_imgs[i]
        oh, ow = orig_img.shape[:2]
        pred[..., [0, 2]] *= ow
        pred[..., [1, 3]] *= oh
        img_path = self.batch[0][i]
        results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
    return results

pre_transform(im)

Chuyển đổi trước các hình ảnh đầu vào trước khi đưa chúng vào mô hình để suy luận. Các hình ảnh đầu vào là hộp thư để đảm bảo tỷ lệ khung hình vuông và được lấp đầy theo tỷ lệ. Kích thước phải là hình vuông (640) và scaleFilled.

Thông số:

Tên Kiểu Sự miêu tả Mặc định
im list[ndarray] | Tensor

Nhập hình ảnh của hình dạng (N, 3, h, w) cho tensor, [(h,w,3) x N] cho danh sách.

bắt buộc

Trở lại:

Kiểu Sự miêu tả
list

Danh sách các hình ảnh được chuyển đổi trước đã sẵn sàng để suy luận mô hình.

Mã nguồn trong ultralytics/models/rtdetr/predict.py
def pre_transform(self, im):
    """
    Pre-transforms the input images before feeding them into the model for inference. The input images are
    letterboxed to ensure a square aspect ratio and scale-filled. The size must be square(640) and scaleFilled.

    Args:
        im (list[np.ndarray] |torch.Tensor): Input images of shape (N,3,h,w) for tensor, [(h,w,3) x N] for list.

    Returns:
        (list): List of pre-transformed images ready for model inference.
    """
    letterbox = LetterBox(self.imgsz, auto=False, scaleFill=True)
    return [letterbox(image=x) for x in im]





Created 2023-11-12, Updated 2024-06-02
Authors: glenn-jocher (5), Burhan-Q (1)