انتقل إلى المحتوى

مرجع ل ultralytics/models/rtdetr/predict.py

ملاحظه

هذا الملف متاح في https://github.com/ultralytics/ultralytics/ نقطة / الرئيسية /ultralytics/النماذج/rtdetr/predict.py. إذا اكتشفت مشكلة ، فيرجى المساعدة في إصلاحها من خلال المساهمة في طلب 🛠️ سحب. شكرا لك 🙏!



ultralytics.models.rtdetr.predict.RTDETRPredictor

قواعد: BasePredictor

RT-DETR (محول الكشف في الوقت الحقيقي) يقوم المتنبئ بتوسيع فئة BasePredictor لعمل التنبؤات باستخدام بايدو RT-DETR نموذج.

تستفيد هذه الفئة من قوة Vision Transformers لتوفير اكتشاف الأشياء في الوقت الفعلي مع الحفاظ على دقة عالية. وهو يدعم الميزات الرئيسية مثل الترميز المختلط الفعال واختيار الاستعلام المدرك لإنترنت الأشياء.

مثل
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()

سمات:

اسم نوع وصف
imgsz int

حجم الصورة للاستدلال (يجب أن يكون مربعا ومملوءا بالمقياس).

args dict

تتجاوز الوسيطة للمتنبئ.

شفرة المصدر في 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)

قم بمعالجة التنبؤات الأولية من النموذج لإنشاء مربعات محيطة ودرجات ثقة.

تقوم الطريقة بتصفية الاكتشافات بناء على الثقة والفئة إذا تم تحديدها في self.args.

البارامترات:

اسم نوع وصف افتراضي
preds list

قائمة [التنبؤات ، إضافية] من النموذج.

مطلوب
img Tensor

معالجة صور الإدخال.

مطلوب
orig_imgs list or Tensor

الصور الأصلية غير المعالجة.

مطلوب

ارجاع:

نوع وصف
list[Results]

قائمة بكائنات النتائج التي تحتوي على المربعات المحيطة المعالجة اللاحقة ، ودرجات الثقة ، وتسميات الطبقة.

شفرة المصدر في 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)

يحول الصور المدخلة مسبقا قبل تغذيتها في النموذج للاستدلال. الصور المدخلة هي letterboxed لضمان نسبة العرض إلى الارتفاع مربعة وملء المقياس. يجب أن يكون الحجم مربعا (640) ومقياسا.

البارامترات:

اسم نوع وصف افتراضي
im list[ndarray] | Tensor

إدخال صور الشكل (N ، 3 ، h ، w) ل tensor، [(h، w، 3) x N] للقائمة.

مطلوب

ارجاع:

نوع وصف
list

قائمة بالصور المحولة مسبقا الجاهزة للاستدلال على النموذج.

شفرة المصدر في 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]





تم إنشاء 2023-11-12, اخر تحديث 2023-11-25
المؤلفون: جلين جوشر (3)