feat: initial HSAP platform
Huaxu Sentinel Active Safety Platform with embedded algorithm code, Docker Compose setup, and vendored dataset scaffolds for clone-and-run. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
56
algorithms/dms_yolo/code/ultralytics/models/nas/predict.py
Normal file
56
algorithms/dms_yolo/code/ultralytics/models/nas/predict.py
Normal file
@@ -0,0 +1,56 @@
|
||||
# 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 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)
|
||||
Reference in New Issue
Block a user