Chengfang Lu
e72bc061c5
Major changes: - New frontend (platform/web/): Vite + React 18 + TypeScript + Tailwind - 4-module navigation: 数据送标 / 模型管理 / 车队管理 / 系统管理 - Data catalog with charts (DMS/ADAS/Lane 3-tab view) - Quality review workflow (标注质检): Good/Fine/Bad scoring with auto-advance - Audit enhancements: batch operations, rejection categories, Feishu notifications - Operation audit log (操作日志) - World model simulation studio (仿真工坊) - Dataset version management with snapshots and diff - ADAS 7-class dataset integration (138K images organized + compressed) - User management with Feishu integration and pagination - CRUD/search/filter on all pages, card layout redesign - PIL-optimized image overlay rendering - Auto-snapshot on build, in_review workflow stage - Removed embedded algorithm code (now in workspace)
57 lines
2.6 KiB
Python
57 lines
2.6 KiB
Python
# 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)
|