# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

import torch

from ultralytics.models.yolo.detect.predict import DetectionPredictor
from ultralytics.utils import ops


class NASPredictor(DetectionPredictor):
    """
    Ultralytics YOLO NAS Predictor for object detection.

    This class extends the `DetectionPredictor` from Ultralytics engine and is responsible for post-processing the
    raw predictions generated by the YOLO NAS models. It applies operations like non-maximum suppression and
    scaling the bounding boxes to fit the original image dimensions.

    Attributes:
        args (Namespace): Namespace containing various configurations for post-processing including confidence threshold,
            IoU threshold, agnostic NMS flag, maximum detections, and class filtering options.
        model (torch.nn.Module): The YOLO NAS model used for inference.
        batch (list): Batch of inputs for processing.

    Examples:
        >>> from ultralytics import NAS
        >>> model = NAS("yolo_nas_s")
        >>> predictor = model.predictor

        Assume that raw_preds, img, orig_imgs are available
        >>> results = predictor.postprocess(raw_preds, img, orig_imgs)

    Notes:
        Typically, this class is not instantiated directly. It is used internally within the `NAS` class.
    """

    def postprocess(self, preds_in, img, orig_imgs):
        """Postprocess predictions and returns a list of Results objects."""
        # Convert boxes from xyxy to xywh format and concatenate with class scores
        boxes = ops.xyxy2xywh(preds_in[0][0])
        preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1)
        return super().postprocess(preds, img, orig_imgs)
