From 9b66319bdee1867147c588b8e0a48f94852055e6 Mon Sep 17 00:00:00 2001 From: Eugene Khvedchenya Date: Mon, 27 May 2024 15:16:07 +0300 Subject: [PATCH] Make YoloNASRDFLHead inherit a base YoloNASDFLHead --- .../module_interfaces/exportable_obb_detector.py | 4 +++- .../training/datasets/data_formats/obb/cxcywhr.py | 7 +++++++ .../detection_models/yolo_nas_r/yolo_nas_r_ndfl_heads.py | 2 ++ .../detection_models/yolo_nas_r/yolo_nas_r_variants.py | 3 +++ 4 files changed, 15 insertions(+), 1 deletion(-) diff --git a/src/super_gradients/module_interfaces/exportable_obb_detector.py b/src/super_gradients/module_interfaces/exportable_obb_detector.py index 90071df05f..628ae2cec6 100644 --- a/src/super_gradients/module_interfaces/exportable_obb_detector.py +++ b/src/super_gradients/module_interfaces/exportable_obb_detector.py @@ -51,7 +51,9 @@ def forward(self, predictions: Any) -> Tuple[Tensor, Tensor]: :return: Implementation of this method must return a tuple of two tensors (boxes, scores) with the following semantics: - - boxes - [B, N, 4] + - boxes - [B, N, 5] in (CX, CY, W, H, R) format. + R represents the angle oriented box in radians, counter-clockwise and following + OpenCV convention of angle of rotation for cv2.minAreaRect/cv2.boxPoints. - scores - [B, N, C] Where N is the maximum number of predictions per image (see self.get_num_pre_nms_predictions()), and C is the number of classes. diff --git a/src/super_gradients/training/datasets/data_formats/obb/cxcywhr.py b/src/super_gradients/training/datasets/data_formats/obb/cxcywhr.py index e3af66c0ad..56a522211a 100644 --- a/src/super_gradients/training/datasets/data_formats/obb/cxcywhr.py +++ b/src/super_gradients/training/datasets/data_formats/obb/cxcywhr.py @@ -24,6 +24,13 @@ def cxcywhr_to_poly(boxes: np.ndarray) -> np.ndarray: def poly_to_cxcywhr(poly: np.ndarray) -> np.ndarray: + """ + Convert a polygon format to oriented bounding boxes in CX-CY-W-H-R format + :param poly: Input polygons in [N,...,4, 2] format + :return: [N,..., 5] Oriented bounding boxes in CX-CY-W-H-R format + R represents the angle oriented box in radians, counter-clockwise and following + OpenCV convention of angle of rotation for cv2.minAreaRect/cv2.boxPoints. + """ shape = poly.shape if shape[-2:] != (4, 2): raise ValueError(f"Expected last two dimensions to be (4, 2), got {shape[-2:]}") diff --git a/src/super_gradients/training/models/detection_models/yolo_nas_r/yolo_nas_r_ndfl_heads.py b/src/super_gradients/training/models/detection_models/yolo_nas_r/yolo_nas_r_ndfl_heads.py index 3f55d871dc..4fc3ea3802 100644 --- a/src/super_gradients/training/models/detection_models/yolo_nas_r/yolo_nas_r_ndfl_heads.py +++ b/src/super_gradients/training/models/detection_models/yolo_nas_r/yolo_nas_r_ndfl_heads.py @@ -37,6 +37,8 @@ class YoloNASRLogits: :param size_reduced: Tensor of shape [B, Anchors, 2] with predicted size distribution. None-multiplied by stride. :param angles: Tensor of shape [B, Anchors, 1] with predicted angles (in radians). + Angle is in range [-3*pi/4, pi/4]. Rotation is counter-clockwise and follows + OpenCV convention of angle of rotation for cv2.minAreaRect/cv2.boxPoints. :param offsets: Tensor of shape [B, Anchors, 2] with predicted offsets. Non-multiplied by stride. :param anchor_points: Tensor of shape [Anchors, 2] with anchor points. diff --git a/src/super_gradients/training/models/detection_models/yolo_nas_r/yolo_nas_r_variants.py b/src/super_gradients/training/models/detection_models/yolo_nas_r/yolo_nas_r_variants.py index 16eb5b2a85..104242bb37 100644 --- a/src/super_gradients/training/models/detection_models/yolo_nas_r/yolo_nas_r_variants.py +++ b/src/super_gradients/training/models/detection_models/yolo_nas_r/yolo_nas_r_variants.py @@ -89,6 +89,9 @@ def forward(self, inputs: Union[Tuple[Tensor, Tensor], YoloNASRLogits]): class YoloNASR(ExportableOBBDetectionModel, CustomizableDetector, SupportsInputShapeCheck): """ YoloNAS-R model for Oriented Bounding Box (OBB) Detection. + The model is based on the YOLO-NAS architecture with minor modifications to support oriented bounding boxes. + + Model supports export to ONNX and TensorRT via model.export(...) syntax. """ def __init__(