Skip to content

Reference for ultralytics/models/nas/predict.py

Improvements

This page is sourced from https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/nas/predict.py. Have an improvement or example to add? Open a Pull Request — thank you! 🙏


Summary

class ultralytics.models.nas.predict.NASPredictor

NASPredictor()

Bases: 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

NameTypeDescription
argsNamespaceNamespace containing various configurations for post-processing including confidence threshold, IoU threshold, agnostic NMS flag, maximum detections, and class filtering options.
modeltorch.nn.ModuleThe YOLO NAS model used for inference.
batchlistBatch of inputs for processing.

Methods

NameDescription
postprocessPostprocess NAS model predictions to generate final detection results.

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.

Source code in ultralytics/models/nas/predict.pyView on GitHub
class NASPredictor(DetectionPredictor):


method ultralytics.models.nas.predict.NASPredictor.postprocess

def postprocess(self, preds_in, img, orig_imgs)

Postprocess NAS model predictions to generate final detection results.

This method takes raw predictions from a YOLO NAS model, converts bounding box formats, and applies post-processing operations to generate the final detection results compatible with Ultralytics result visualization and analysis tools.

Args

NameTypeDescriptionDefault
preds_inlistRaw predictions from the NAS model, typically containing bounding boxes and class scores.required
imgtorch.TensorInput image tensor that was fed to the model, with shape (B, C, H, W).required
orig_imgslist | torch.Tensor | np.ndarrayOriginal images before preprocessing, used for scaling coordinates back to original dimensions.required

Returns

TypeDescription
listList of Results objects containing the processed predictions for each image in the batch.

Examples

>>> predictor = NAS("yolo_nas_s").predictor
>>> results = predictor.postprocess(raw_preds, img, orig_imgs)
Source code in ultralytics/models/nas/predict.pyView on GitHub
def postprocess(self, preds_in, img, orig_imgs):
    """Postprocess NAS model predictions to generate final detection results.

    This method takes raw predictions from a YOLO NAS model, converts bounding box formats, and applies
    post-processing operations to generate the final detection results compatible with Ultralytics result
    visualization and analysis tools.

    Args:
        preds_in (list): Raw predictions from the NAS model, typically containing bounding boxes and class scores.
        img (torch.Tensor): Input image tensor that was fed to the model, with shape (B, C, H, W).
        orig_imgs (list | torch.Tensor | np.ndarray): Original images before preprocessing, used for scaling
            coordinates back to original dimensions.

    Returns:
        (list): List of Results objects containing the processed predictions for each image in the batch.

    Examples:
        >>> predictor = NAS("yolo_nas_s").predictor
        >>> results = predictor.postprocess(raw_preds, img, orig_imgs)
    """
    boxes = ops.xyxy2xywh(preds_in[0][0])  # Convert bounding boxes from xyxy to xywh format
    preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1)  # Concatenate boxes with class scores
    return super().postprocess(preds, img, orig_imgs)





📅 Created 2 years ago ✏️ Updated 18 days ago
glenn-jocherjk4eBurhan-Q