Skip to content

Référence pour ultralytics/models/rtdetr/predict.py

Note

Ce fichier est disponible à l'adresse https://github.com/ultralytics/ ultralytics/blob/main/ ultralytics/models/rtdetr/predict .py. Si tu repères un problème, aide à le corriger en contribuant à une Pull Request 🛠️. Merci 🙏 !



ultralytics.models.rtdetr.predict.RTDETRPredictor

Bases : BasePredictor

RT-DETR (Real-Time Detection Transformer) Prédicteur étendant la classe BasePredictor pour faire des prédictions en utilisant le modèle de Baidu. le modèle RT-DETR de Baidu.

Cette classe tire parti de la puissance des transformateurs de vision pour assurer la détection d'objets en temps réel tout en maintenant une grande précision. une grande précision. Elle prend en charge des fonctions clés telles que l'encodage hybride efficace et la sélection de requêtes tenant compte de l'interface utilisateur.

Exemple
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()

Attributs :

Nom Type Description
imgsz int

Taille de l'image pour l'inférence (doit être carrée et remplie à l'échelle).

args dict

Remplace les arguments du prédicteur.

Code source dans 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)

Post-traite les prédictions brutes du modèle pour générer des boîtes de délimitation et des scores de confiance.

La méthode filtre les détections en fonction de la confiance et de la classe si elles sont spécifiées dans la rubrique self.args.

Paramètres :

Nom Type Description DĂ©faut
preds list

Liste des [prédictions, supplémentaires] du modèle.

requis
img Tensor

Traite les images d'entrée.

requis
orig_imgs list or Tensor

Images originales, non traitées.

requis

Retourne :

Type Description
list[Results]

Une liste d'objets Résultats contenant les boîtes de délimitation post-traitées, les scores de confiance, et les étiquettes de classe.

Code source dans 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)

Prétransforme les images d'entrée avant de les introduire dans le modèle pour l'inférence. Les images d'entrée sont letterboxed pour garantir un rapport d'aspect carré et sont remplies à l'échelle. La taille doit être square(640) et scaleFilled.

Paramètres :

Nom Type Description DĂ©faut
im list[ndarray] | Tensor

Images d'entrée de forme (N,3,h,w) pour tensor, [(h,w,3) x N] pour la liste.

requis

Retourne :

Type Description
list

Liste d'images prétransformées prêtes pour l'inférence du modèle.

Code source dans 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]





Création : 2023-11-12, mise à jour : 2024-06-02
Auteur(s) : glenn-jocher (5), Burhan-Q (1)