Skip to content

Models API Reference

oriented_det.models

Model registry for oriented detection baselines.

ClassWeightsMixin

Mixin class for models that support class weights.

Provides common functionality for managing class weights in two-stage detectors (RotatedFasterRCNN, OrientedRCNN).

Source code in oriented_det/models/utils.py 
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
class ClassWeightsMixin:
    """Mixin class for models that support class weights.

    Provides common functionality for managing class weights in two-stage
    detectors (RotatedFasterRCNN, OrientedRCNN).
    """

    def __init__(
        self,
        num_classes: int,
        class_weights: Optional[Union[Dict[str, float], torch.Tensor]] = None,
    ):
        """Initialize class weights mixin.

        Args:
            num_classes: Number of object classes (excluding background)
            class_weights: Optional class weights:
                - Dict[str, float]: Mapping from class name to weight
                - torch.Tensor: Tensor of shape [num_classes + 1] (including background)
                - None: Equal weights for all classes
        """
        self.num_classes = num_classes
        self._roi_class_weights = class_weights  # Store raw input
        self.roi_class_weights_tensor: Optional[torch.Tensor] = None  # Will be set when class mapping is known

        # Validate tensor if provided directly
        if class_weights is not None and isinstance(class_weights, torch.Tensor):
            if class_weights.shape[0] != num_classes + 1:
                raise ValueError(
                    f"Class weights tensor must have shape [num_classes + 1] = [{num_classes + 1}], "
                    f"but got shape {class_weights.shape}"
                )
            self.roi_class_weights_tensor = class_weights

    def set_class_weights(
        self,
        class_map: Dict[str, int],
        device: Optional[torch.device] = None,
    ) -> None:
        """Set class weights tensor from class mapping.

        Converts class weights dictionary (class_name -> weight) to tensor format
        using the provided class mapping. Must be called before training if using
        class weights with a dictionary.

        Args:
            class_map: Mapping from class names to class IDs (1-indexed, 0 is background)
            device: Device to create tensor on (defaults to model's device)

        Example:
            >>> class_map = {"plane": 1, "ship": 2, "small-vehicle": 3}
            >>> class_weights = {"plane": 2.0, "ship": 1.5, "small-vehicle": 1.0}
            >>> model = OrientedRCNN(num_classes=3, roi_class_weights=class_weights)
            >>> model.set_class_weights(class_map)
        """
        if torch is None:
            raise RuntimeError("PyTorch is required.")

        if self._roi_class_weights is None:
            return

        if not isinstance(self._roi_class_weights, dict):
            return  # Already a tensor or None

        if device is None:
            # Try to get device from model parameters
            if hasattr(self, 'parameters'):
                device = next(self.parameters()).device if list(self.parameters()) else torch.device("cpu")
            else:
                device = torch.device("cpu")

        # Create weights tensor: [num_classes + 1] (background + object classes)
        weights = torch.ones(self.num_classes + 1, dtype=torch.float32, device=device)

        # Set background weight (index 0) - default to 1.0 if not specified
        if "background" in self._roi_class_weights:
            weights[0] = self._roi_class_weights["background"]
        elif "__background__" in self._roi_class_weights:
            weights[0] = self._roi_class_weights["__background__"]

        # Set class weights (indices 1, 2, ..., num_classes)
        for class_name, class_id in class_map.items():
            if class_id < 1 or class_id > self.num_classes:
                continue  # Skip invalid class IDs
            if class_name in self._roi_class_weights:
                weights[class_id] = self._roi_class_weights[class_name]

        self.roi_class_weights_tensor = weights

    def set_class_weights_tensor(self, weights: "torch.Tensor") -> None:
        """Set ROI class weights directly as a tensor.

        Args:
            weights: Tensor of shape [num_classes + 1] (including background).
        """
        if torch is None:
            raise RuntimeError("PyTorch is required.")
        if not isinstance(weights, torch.Tensor):
            raise TypeError("weights must be a torch.Tensor")
        if weights.ndim != 1 or weights.shape[0] != self.num_classes + 1:
            raise ValueError(
                f"weights must have shape [{self.num_classes + 1}], got {tuple(weights.shape)}"
            )
        self.roi_class_weights_tensor = weights

    @property
    def roi_class_weights(self) -> Optional[torch.Tensor]:
        """Get the current class weights tensor.

        Returns the class weights tensor if it has been set (via `set_class_weights` or
        passed directly as a tensor during initialization). Returns None if no class
        weights are configured.

        Returns:
            Class weights tensor of shape [num_classes + 1] (including background),
            or None if not configured.

        Example:
            >>> model = OrientedRCNN(num_classes=15)
            >>> weights = model.roi_class_weights  # Returns None if not set
            >>> model.set_class_weights(class_map)
            >>> weights = model.roi_class_weights  # Returns tensor after setting
        """
        return self.roi_class_weights_tensor

roi_class_weights property

Get the current class weights tensor.

Returns the class weights tensor if it has been set (via set_class_weights or passed directly as a tensor during initialization). Returns None if no class weights are configured.

Returns:

Type Description
Optional[Tensor]

Class weights tensor of shape [num_classes + 1] (including background),
Optional[Tensor]

or None if not configured.
Example

model = OrientedRCNN(num_classes=15) weights = model.roi_class_weights # Returns None if not set model.set_class_weights(class_map) weights = model.roi_class_weights # Returns tensor after setting

__init__(num_classes, class_weights=None)

Initialize class weights mixin.

Parameters:

Name Type Description Default
num_classes int

Number of object classes (excluding background)

 required
class_weights Optional[Union[Dict[str, float], Tensor]]

Optional class weights: - Dict[str, float]: Mapping from class name to weight - torch.Tensor: Tensor of shape [num_classes + 1] (including background) - None: Equal weights for all classes

 None
Source code in oriented_det/models/utils.py 
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
def __init__(
    self,
    num_classes: int,
    class_weights: Optional[Union[Dict[str, float], torch.Tensor]] = None,
):
    """Initialize class weights mixin.

    Args:
        num_classes: Number of object classes (excluding background)
        class_weights: Optional class weights:
            - Dict[str, float]: Mapping from class name to weight
            - torch.Tensor: Tensor of shape [num_classes + 1] (including background)
            - None: Equal weights for all classes
    """
    self.num_classes = num_classes
    self._roi_class_weights = class_weights  # Store raw input
    self.roi_class_weights_tensor: Optional[torch.Tensor] = None  # Will be set when class mapping is known

    # Validate tensor if provided directly
    if class_weights is not None and isinstance(class_weights, torch.Tensor):
        if class_weights.shape[0] != num_classes + 1:
            raise ValueError(
                f"Class weights tensor must have shape [num_classes + 1] = [{num_classes + 1}], "
                f"but got shape {class_weights.shape}"
            )
        self.roi_class_weights_tensor = class_weights

set_class_weights(class_map, device=None)

Set class weights tensor from class mapping.

Converts class weights dictionary (class_name -> weight) to tensor format using the provided class mapping. Must be called before training if using class weights with a dictionary.

Parameters:

Name Type Description Default
class_map Dict[str, int]

Mapping from class names to class IDs (1-indexed, 0 is background)

 required
device Optional[device]

Device to create tensor on (defaults to model's device)

 None
Example

class_map = {"plane": 1, "ship": 2, "small-vehicle": 3} class_weights = {"plane": 2.0, "ship": 1.5, "small-vehicle": 1.0} model = OrientedRCNN(num_classes=3, roi_class_weights=class_weights) model.set_class_weights(class_map)

Source code in oriented_det/models/utils.py 
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
def set_class_weights(
    self,
    class_map: Dict[str, int],
    device: Optional[torch.device] = None,
) -> None:
    """Set class weights tensor from class mapping.

    Converts class weights dictionary (class_name -> weight) to tensor format
    using the provided class mapping. Must be called before training if using
    class weights with a dictionary.

    Args:
        class_map: Mapping from class names to class IDs (1-indexed, 0 is background)
        device: Device to create tensor on (defaults to model's device)

    Example:
        >>> class_map = {"plane": 1, "ship": 2, "small-vehicle": 3}
        >>> class_weights = {"plane": 2.0, "ship": 1.5, "small-vehicle": 1.0}
        >>> model = OrientedRCNN(num_classes=3, roi_class_weights=class_weights)
        >>> model.set_class_weights(class_map)
    """
    if torch is None:
        raise RuntimeError("PyTorch is required.")

    if self._roi_class_weights is None:
        return

    if not isinstance(self._roi_class_weights, dict):
        return  # Already a tensor or None

    if device is None:
        # Try to get device from model parameters
        if hasattr(self, 'parameters'):
            device = next(self.parameters()).device if list(self.parameters()) else torch.device("cpu")
        else:
            device = torch.device("cpu")

    # Create weights tensor: [num_classes + 1] (background + object classes)
    weights = torch.ones(self.num_classes + 1, dtype=torch.float32, device=device)

    # Set background weight (index 0) - default to 1.0 if not specified
    if "background" in self._roi_class_weights:
        weights[0] = self._roi_class_weights["background"]
    elif "__background__" in self._roi_class_weights:
        weights[0] = self._roi_class_weights["__background__"]

    # Set class weights (indices 1, 2, ..., num_classes)
    for class_name, class_id in class_map.items():
        if class_id < 1 or class_id > self.num_classes:
            continue  # Skip invalid class IDs
        if class_name in self._roi_class_weights:
            weights[class_id] = self._roi_class_weights[class_name]

    self.roi_class_weights_tensor = weights

set_class_weights_tensor(weights)

Set ROI class weights directly as a tensor.

Parameters:

Name Type Description Default
weights 'torch.Tensor'

Tensor of shape [num_classes + 1] (including background).

 required
Source code in oriented_det/models/utils.py 
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
def set_class_weights_tensor(self, weights: "torch.Tensor") -> None:
    """Set ROI class weights directly as a tensor.

    Args:
        weights: Tensor of shape [num_classes + 1] (including background).
    """
    if torch is None:
        raise RuntimeError("PyTorch is required.")
    if not isinstance(weights, torch.Tensor):
        raise TypeError("weights must be a torch.Tensor")
    if weights.ndim != 1 or weights.shape[0] != self.num_classes + 1:
        raise ValueError(
            f"weights must have shape [{self.num_classes + 1}], got {tuple(weights.shape)}"
        )
    self.roi_class_weights_tensor = weights

DeltaXYWHAHBBoxCoder

5 params: dx, dy, dw, dh, da. norm_factor, edge_swap. MMRotate ROI head.

Source code in oriented_det/models/bbox_coder.py 
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
class DeltaXYWHAHBBoxCoder:
    """5 params: dx, dy, dw, dh, da. norm_factor, edge_swap. MMRotate ROI head."""

    def __init__(
        self,
        target_means: Tuple[float, ...] = (0.0, 0.0, 0.0, 0.0, 0.0),
        target_stds: Tuple[float, ...] = (1.0, 1.0, 1.0, 1.0, 1.0),
        norm_factor: Optional[float] = 2.0,
        edge_swap: bool = True,
        angle_range: str = "le90",
    ):
        self.target_means = target_means
        self.target_stds = target_stds
        self.norm_factor = norm_factor
        self.edge_swap = edge_swap
        self.angle_range = angle_range

    def encode(
        self,
        anchors: torch.Tensor,
        gt_boxes: torch.Tensor,
    ) -> torch.Tensor:
        """Encode gt boxes relative to anchors. anchors/gt: [N, 5] (cx, cy, w, h, angle)."""
        return encode_oriented_boxes(
            anchors,
            gt_boxes,
            target_means=self.target_means,
            target_stds=self.target_stds,
            norm_factor=self.norm_factor,
            edge_swap=self.edge_swap,
        )

    def decode(
        self,
        anchors: torch.Tensor,
        deltas: torch.Tensor,
        normalize_le90: bool = True,
    ) -> torch.Tensor:
        """Decode deltas [N, 5] to boxes [N, 5]."""
        return decode_oriented_boxes(
            anchors,
            deltas,
            target_means=self.target_means,
            target_stds=self.target_stds,
            normalize_le90=normalize_le90,
            norm_factor=self.norm_factor,
            edge_swap=self.edge_swap,
        )

decode(anchors, deltas, normalize_le90=True)

Decode deltas [N, 5] to boxes [N, 5].

Source code in oriented_det/models/bbox_coder.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def decode(
    self,
    anchors: torch.Tensor,
    deltas: torch.Tensor,
    normalize_le90: bool = True,
) -> torch.Tensor:
    """Decode deltas [N, 5] to boxes [N, 5]."""
    return decode_oriented_boxes(
        anchors,
        deltas,
        target_means=self.target_means,
        target_stds=self.target_stds,
        normalize_le90=normalize_le90,
        norm_factor=self.norm_factor,
        edge_swap=self.edge_swap,
    )

encode(anchors, gt_boxes)

Encode gt boxes relative to anchors. anchors/gt: [N, 5] (cx, cy, w, h, angle).

Source code in oriented_det/models/bbox_coder.py 
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def encode(
    self,
    anchors: torch.Tensor,
    gt_boxes: torch.Tensor,
) -> torch.Tensor:
    """Encode gt boxes relative to anchors. anchors/gt: [N, 5] (cx, cy, w, h, angle)."""
    return encode_oriented_boxes(
        anchors,
        gt_boxes,
        target_means=self.target_means,
        target_stds=self.target_stds,
        norm_factor=self.norm_factor,
        edge_swap=self.edge_swap,
    )

DeltaXYWHBBoxCoder

4 params: dx, dy, dw, dh. Angle from anchor. MMRotate Rotated Faster R-CNN RPN.

Same as encode_rpn_boxes / decode_rpn_boxes.

Source code in oriented_det/models/bbox_coder.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class DeltaXYWHBBoxCoder:
    """4 params: dx, dy, dw, dh. Angle from anchor. MMRotate Rotated Faster R-CNN RPN.

    Same as encode_rpn_boxes / decode_rpn_boxes.
    """

    def __init__(
        self,
        target_means: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
        target_stds: Tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0),
    ):
        self.target_means = target_means
        self.target_stds = target_stds

    def encode(
        self,
        anchors: torch.Tensor,
        gt_boxes: torch.Tensor,
    ) -> torch.Tensor:
        """Encode gt boxes relative to anchors. anchors/gt: [N, 5] (cx, cy, w, h, angle)."""
        return encode_rpn_boxes(anchors, gt_boxes, self.target_means, self.target_stds)

    def decode(
        self,
        anchors: torch.Tensor,
        deltas: torch.Tensor,
    ) -> torch.Tensor:
        """Decode deltas [N, 4] to boxes [N, 5]. Angle from anchor."""
        return decode_rpn_boxes(anchors, deltas, self.target_means, self.target_stds)

decode(anchors, deltas)

Decode deltas [N, 4] to boxes [N, 5]. Angle from anchor.

Source code in oriented_det/models/bbox_coder.py 
54
55
56
57
58
59
60
def decode(
    self,
    anchors: torch.Tensor,
    deltas: torch.Tensor,
) -> torch.Tensor:
    """Decode deltas [N, 4] to boxes [N, 5]. Angle from anchor."""
    return decode_rpn_boxes(anchors, deltas, self.target_means, self.target_stds)

encode(anchors, gt_boxes)

Encode gt boxes relative to anchors. anchors/gt: [N, 5] (cx, cy, w, h, angle).

Source code in oriented_det/models/bbox_coder.py 
46
47
48
49
50
51
52
def encode(
    self,
    anchors: torch.Tensor,
    gt_boxes: torch.Tensor,
) -> torch.Tensor:
    """Encode gt boxes relative to anchors. anchors/gt: [N, 5] (cx, cy, w, h, angle)."""
    return encode_rpn_boxes(anchors, gt_boxes, self.target_means, self.target_stds)

MidpointOffsetCoder

6 params: dx, dy, dw, dh, da, db. MMRotate Oriented R-CNN RPN (MidpointOffsetCoder).

Encodes/decodes axis-aligned proposals (xyxy) relative to oriented GT using midpoint offsets. Proposals must be axis-aligned xyxy [N, 4] (x1, y1, x2, y2) so that center (px, py) and size (pw, ph) are uniquely defined; the coder then predicts deltas to transform each horizontal box into an oriented box (da, db = offsets of the oriented box's top/right edge midpoints relative to the horizontal box). Used when the RPN outputs horizontal proposals and a second stage predicts orientation (e.g. Oriented R-CNN).

Source code in oriented_det/models/bbox_coder.py 
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
class MidpointOffsetCoder:
    """6 params: dx, dy, dw, dh, da, db. MMRotate Oriented R-CNN RPN (MidpointOffsetCoder).

    Encodes/decodes axis-aligned proposals (xyxy) relative to oriented GT using midpoint offsets.
    Proposals must be axis-aligned xyxy [N, 4] (x1, y1, x2, y2) so that center (px, py) and
    size (pw, ph) are uniquely defined; the coder then predicts deltas to transform each
    horizontal box into an oriented box (da, db = offsets of the oriented box's top/right
    edge midpoints relative to the horizontal box). Used when the RPN outputs horizontal
    proposals and a second stage predicts orientation (e.g. Oriented R-CNN).
    """

    def __init__(
        self,
        target_means: Tuple[float, ...] = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
        target_stds: Tuple[float, ...] = (1.0, 1.0, 1.0, 1.0, 0.5, 0.5),
        angle_range: str = "le90",
    ):
        self.target_means = target_means
        self.target_stds = target_stds
        self.angle_range = angle_range

    def encode(
        self,
        proposals: torch.Tensor,
        gt_bboxes: torch.Tensor,
    ) -> torch.Tensor:
        """Encode GT obb relative to axis-aligned proposals.

        proposals: [N, 4] xyxy (x1, y1, x2, y2); must be axis-aligned so (px, py, pw, ph)
        are well-defined. gt_bboxes: [N, 5] (cx, cy, w, h, angle). Returns [N, 6] deltas.
        """
        return _midpoint_offset_encode(proposals, gt_bboxes, self.target_means, self.target_stds)

    def decode(
        self,
        proposals: torch.Tensor,
        deltas: torch.Tensor,
        wh_ratio_clip: float = 16 / 1000,
    ) -> torch.Tensor:
        """Decode deltas to oriented boxes.

        proposals: [N, 4] xyxy (axis-aligned); deltas: [N, 6]. Returns [N, 5] obb (cx, cy, w, h, angle).
        """
        return _midpoint_offset_decode(
            proposals, deltas, self.target_means, self.target_stds, wh_ratio_clip
        )

decode(proposals, deltas, wh_ratio_clip=16 / 1000)

Decode deltas to oriented boxes.

proposals: [N, 4] xyxy (axis-aligned); deltas: [N, 6]. Returns [N, 5] obb (cx, cy, w, h, angle).

Source code in oriented_det/models/bbox_coder.py 
222
223
224
225
226
227
228
229
230
231
232
233
234
def decode(
    self,
    proposals: torch.Tensor,
    deltas: torch.Tensor,
    wh_ratio_clip: float = 16 / 1000,
) -> torch.Tensor:
    """Decode deltas to oriented boxes.

    proposals: [N, 4] xyxy (axis-aligned); deltas: [N, 6]. Returns [N, 5] obb (cx, cy, w, h, angle).
    """
    return _midpoint_offset_decode(
        proposals, deltas, self.target_means, self.target_stds, wh_ratio_clip
    )

encode(proposals, gt_bboxes)

Encode GT obb relative to axis-aligned proposals.

proposals: [N, 4] xyxy (x1, y1, x2, y2); must be axis-aligned so (px, py, pw, ph) are well-defined. gt_bboxes: [N, 5] (cx, cy, w, h, angle). Returns [N, 6] deltas.

Source code in oriented_det/models/bbox_coder.py 
210
211
212
213
214
215
216
217
218
219
220
def encode(
    self,
    proposals: torch.Tensor,
    gt_bboxes: torch.Tensor,
) -> torch.Tensor:
    """Encode GT obb relative to axis-aligned proposals.

    proposals: [N, 4] xyxy (x1, y1, x2, y2); must be axis-aligned so (px, py, pw, ph)
    are well-defined. gt_bboxes: [N, 5] (cx, cy, w, h, angle). Returns [N, 6] deltas.
    """
    return _midpoint_offset_encode(proposals, gt_bboxes, self.target_means, self.target_stds)

OrientedRCNN

Bases: ClassWeightsMixin, GroupedCeMixin, Module

Oriented R-CNN (MMRotate-style): horizontal RPN + MidpointOffset (6 params) -> oriented ROI head.

RPN anchor_angles is optional on the constructor only (not in training JSON); default is horizontal [0.0].

Source code in oriented_det/models/oriented_rcnn.py 
 764
 765
 766
 767
 768
 769
 770
 771
 772
 773
 774
 775
 776
 777
 778
 779
 780
 781
 782
 783
 784
 785
 786
 787
 788
 789
 790
 791
 792
 793
 794
 795
 796
 797
 798
 799
 800
 801
 802
 803
 804
 805
 806
 807
 808
 809
 810
 811
 812
 813
 814
 815
 816
 817
 818
 819
 820
 821
 822
 823
 824
 825
 826
 827
 828
 829
 830
 831
 832
 833
 834
 835
 836
 837
 838
 839
 840
 841
 842
 843
 844
 845
 846
 847
 848
 849
 850
 851
 852
 853
 854
 855
 856
 857
 858
 859
 860
 861
 862
 863
 864
 865
 866
 867
 868
 869
 870
 871
 872
 873
 874
 875
 876
 877
 878
 879
 880
 881
 882
 883
 884
 885
 886
 887
 888
 889
 890
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
class OrientedRCNN(ClassWeightsMixin, GroupedCeMixin, nn.Module):
    """Oriented R-CNN (MMRotate-style): horizontal RPN + MidpointOffset (6 params) -> oriented ROI head.

    RPN ``anchor_angles`` is optional on the constructor only (not in training JSON); default is horizontal ``[0.0]``.
    """

    def __init__(
        self,
        num_classes: int,
        *,
        backbone=None,
        backbone_name: str = "resnet50",
        pretrained_backbone: bool = False,
        trainable_layers: int = 5,
        anchor_scales: Optional[List[float]] = None,
        anchor_ratios: Optional[List[float]] = None,
        anchor_angles: Optional[List[float]] = None,
        rpn_pre_nms_top_n: int = 2000,
        rpn_post_nms_top_n: int = 1000,
        max_detections_per_image: int = 100,  # Maximum number of detections per image after NMS
        rpn_nms_threshold: float = 0.7,  # IoU threshold for RPN proposal NMS
        final_nms_iou_threshold: float = 0.5,  # Final detection NMS (after ROI head); not rpn_nms_threshold
        rpn_positive_iou_threshold: float = 0.5,
        rpn_negative_iou_threshold: float = 0.2,
        rpn_fg_bg_sampling_ratio: float = 0.5,
        rpn_batch_size_per_image: int = 256,  # Number of RPN anchors to sample per image (MMRotate: 256)
        rpn_min_pos_iou: float = 0.3,
        rpn_match_low_quality: bool = True,
        roi_match_low_quality: bool = False,
        roi_positive_iou_threshold: float = 0.4,
        roi_negative_iou_threshold: float = 0.3,
        roi_fg_bg_sampling_ratio: float = 0.25,
        roi_batch_size_per_image: int = 512,  # Number of ROI proposals to sample per image for training
        roi_output_size: Tuple[int, int] = (7, 7),
        roi_spatial_scales: Optional[List[float]] = None,
        fpn_strides: Optional[List[int]] = None,
        returned_layers: Optional[List[int]] = None,
        use_checkpoint: bool = False,
        roi_chunk_size: int = 32,
        roi_use_checkpoint: bool = False,
        roi_loss_type: str = "cross_entropy",
        roi_class_weights: Optional[Union[Dict[str, float], torch.Tensor]] = None,
        roi_focal_alpha: float = 1.0,
        roi_focal_gamma: float = 2.0,
        roi_label_smoothing: float = 0.0,
        target_means: Optional[Tuple[float, float, float, float, float]] = None,
        target_stds: Optional[Tuple[float, float, float, float, float]] = None,
        roi_norm_factor: Optional[float] = None,
        roi_edge_swap: bool = True,
        roi_proj_xy: bool = True,
        roi_box_reg_angle_weight: float = 1.0,
        roi_box_reg_angle_schedule_epochs: Optional[List[int]] = None,
        roi_box_reg_angle_schedule_values: Optional[List[float]] = None,
        roi_box_reg_iou_weight: float = 0.0,
        roi_box_reg_iou_loss_type: str = "riou",
        roi_box_reg_kfiou_fun: Optional[str] = None,
        roi_box_reg_probiou_mode: Optional[str] = None,
        roi_box_reg_norm: str = "sampled_all",
        use_hbb_for_matching: bool = False,
        roi_use_hbb_for_matching: bool = False,
        inference_pre_nms_score_threshold: float = 0.05,
        final_nms_iou_schedule_epochs: Optional[List[int]] = None,
        final_nms_iou_schedule_values: Optional[List[float]] = None,
        roi_box_reg_iou_schedule_epochs: Optional[List[int]] = None,
        roi_box_reg_iou_schedule_values: Optional[List[float]] = None,
        nms_class_agnostic: bool = False,
        final_nms_use_cpu: bool = False,
        roi_inference_top_class_only: bool = False,
        add_gt_as_proposals: bool = True,
    ):
        if torch is None or nn is None:
            raise RuntimeError("PyTorch is required for OrientedRCNN.")
        # Initialize ClassWeightsMixin first, then nn.Module
        # ClassWeightsMixin.__init__() sets self.num_classes
        super().__init__(num_classes, roi_class_weights)
        nn.Module.__init__(self)
        self._init_grouped_ce()
        self.target_means = target_means if target_means is not None else (0.0, 0.0, 0.0, 0.0, 0.0)
        self.target_stds = target_stds if target_stds is not None else (1.0, 1.0, 1.0, 1.0, 1.0)
        self.roi_norm_factor = roi_norm_factor
        self.roi_edge_swap = roi_edge_swap
        self.roi_proj_xy = roi_proj_xy
        self._roi_box_reg_angle_schedule_epochs = roi_box_reg_angle_schedule_epochs
        self._roi_box_reg_angle_schedule_values = roi_box_reg_angle_schedule_values
        self._roi_box_reg_angle_weight_default = float(roi_box_reg_angle_weight)
        self.roi_box_reg_angle_weight = float(roi_box_reg_angle_weight)
        self.set_roi_box_reg_angle_weight_for_epoch(0)
        self.roi_box_reg_iou_loss_type = roi_box_reg_iou_loss_type
        self.roi_box_reg_kfiou_fun = roi_box_reg_kfiou_fun
        self.roi_box_reg_probiou_mode = roi_box_reg_probiou_mode
        self.roi_box_reg_norm = roi_box_reg_norm
        self.use_hbb_for_matching = use_hbb_for_matching
        self.roi_use_hbb_for_matching = roi_use_hbb_for_matching
        self.inference_pre_nms_score_threshold = inference_pre_nms_score_threshold
        self.roi_loss_type = roi_loss_type
        self.roi_focal_alpha = roi_focal_alpha
        self.roi_focal_gamma = roi_focal_gamma
        self.roi_label_smoothing = roi_label_smoothing

        # Setup backbone using shared utility (returned_layers=[2,3,4] for MMRotate-style C3–C5 only)
        self.backbone, backbone_channels = setup_backbone(
            backbone=backbone,
            backbone_name=backbone_name,
            pretrained_backbone=pretrained_backbone,
            trainable_layers=trainable_layers,
            returned_layers=returned_layers,
        )
        if fpn_strides is None and returned_layers == [2, 3, 4]:
            fpn_strides = [8, 16, 32, 64]

        self.anchor_scales, self.anchor_ratios, self.anchor_angles, self.num_anchors = setup_anchors(
            anchor_scales=anchor_scales,
            anchor_ratios=anchor_ratios,
            anchor_angles=anchor_angles,
            default_angles=[0.0],
        )
        self.rpn_head = OrientedRPNHead(
            backbone_channels, num_anchors=self.num_anchors, cls_out_channels=1, reg_out_channels=6
        )
        self.roi_align = OrientedROIAlign(
            output_size=roi_output_size,
            spatial_scales=roi_spatial_scales,
            fpn_strides=fpn_strides,
            chunk_size=roi_chunk_size,
            use_checkpoint=roi_use_checkpoint,
        )
        roi_feature_size = roi_output_size[0] * roi_output_size[1] * backbone_channels
        self.roi_class_agnostic_regression = True
        self.roi_head = OrientedROIHead(
            in_channels=roi_feature_size,
            num_classes=num_classes,
            class_agnostic_regression=True,
        )
        self.rpn_pre_nms_top_n = rpn_pre_nms_top_n
        self.rpn_post_nms_top_n = rpn_post_nms_top_n
        self.max_detections_per_image = max_detections_per_image
        self.rpn_nms_threshold = rpn_nms_threshold
        self.final_nms_iou_threshold = final_nms_iou_threshold
        self.nms_class_agnostic = nms_class_agnostic
        self.final_nms_use_cpu = final_nms_use_cpu
        self.roi_inference_top_class_only = roi_inference_top_class_only
        self.rpn_positive_iou_threshold = rpn_positive_iou_threshold
        self.rpn_negative_iou_threshold = rpn_negative_iou_threshold
        self.rpn_fg_bg_sampling_ratio = rpn_fg_bg_sampling_ratio
        self.rpn_batch_size_per_image = rpn_batch_size_per_image
        self.rpn_min_pos_iou = rpn_min_pos_iou
        self.rpn_match_low_quality = rpn_match_low_quality
        self.roi_match_low_quality = roi_match_low_quality
        self.roi_positive_iou_threshold = roi_positive_iou_threshold
        self.roi_negative_iou_threshold = roi_negative_iou_threshold
        self.roi_fg_bg_sampling_ratio = roi_fg_bg_sampling_ratio
        self.roi_batch_size_per_image = roi_batch_size_per_image
        self.fpn_strides = fpn_strides or [4, 8, 16, 32, 64]
        self.roi_spatial_scales = roi_spatial_scales or [1.0 / s for s in self.fpn_strides]
        self.use_checkpoint = use_checkpoint
        self._final_nms_iou_schedule_epochs = final_nms_iou_schedule_epochs
        self._final_nms_iou_schedule_values = final_nms_iou_schedule_values
        self.add_gt_as_proposals = add_gt_as_proposals
        self._roi_box_reg_iou_schedule_epochs = roi_box_reg_iou_schedule_epochs
        self._roi_box_reg_iou_schedule_values = roi_box_reg_iou_schedule_values
        self._roi_box_reg_iou_weight_default = float(roi_box_reg_iou_weight)
        self.roi_box_reg_iou_weight = float(roi_box_reg_iou_weight)
        self.set_roi_box_reg_iou_weight_for_epoch(0)

    def set_final_nms_iou_for_epoch(self, epoch: int) -> None:
        """Update final detection NMS IoU threshold from schedule. Lower = more aggressive suppression."""
        if self._final_nms_iou_schedule_epochs is None or self._final_nms_iou_schedule_values is None:
            return
        if not self._final_nms_iou_schedule_epochs or not self._final_nms_iou_schedule_values:
            return
        idx = 0
        for boundary in self._final_nms_iou_schedule_epochs:
            if epoch < boundary:
                break
            idx += 1
        idx = min(idx, len(self._final_nms_iou_schedule_values) - 1)
        self.final_nms_iou_threshold = self._final_nms_iou_schedule_values[idx]

    def set_roi_box_reg_iou_weight_for_epoch(self, epoch: int) -> None:
        """Update ROI auxiliary IoU loss weight from schedule (0-based epoch index)."""
        from oriented_det.train.piecewise_schedule import resolve_piecewise_schedule

        self.roi_box_reg_iou_weight = resolve_piecewise_schedule(
            epoch,
            self._roi_box_reg_iou_schedule_epochs,
            self._roi_box_reg_iou_schedule_values,
            self._roi_box_reg_iou_weight_default,
        )

    def set_roi_box_reg_angle_weight_for_epoch(self, epoch: int) -> None:
        """Update ROI angle (5th dim) SmoothL1 weight from schedule (0-based epoch index)."""
        from oriented_det.train.piecewise_schedule import resolve_piecewise_schedule

        self.roi_box_reg_angle_weight = resolve_piecewise_schedule(
            epoch,
            self._roi_box_reg_angle_schedule_epochs,
            self._roi_box_reg_angle_schedule_values,
            self._roi_box_reg_angle_weight_default,
        )

    def forward(
        self,
        images: Sequence[torch.Tensor],
        targets: Optional[Sequence[Dict[str, Any]]] = None,
    ) -> Union[Dict[str, torch.Tensor], List[Dict[str, Any]]]:
        if not isinstance(images, (list, tuple)):
            images = [images]
        image_sizes = [(img.shape[-2], img.shape[-1]) for img in images]

        # Extract features using shared utility. include_pool_level keeps torchvision's
        # P6 (stride-64 max-pool level) so the RPN sees 5 levels like MMRotate.
        feature_list = extract_backbone_features(
            self.backbone,
            images,
            use_checkpoint=self.use_checkpoint,
            training=self.training,
            include_pool_level=True,
        )
        feature_map_sizes = [(f.shape[2], f.shape[3]) for f in feature_list]
        fpn_strides_live = derive_fpn_strides_from_grid(image_sizes[0], feature_map_sizes)
        warn_if_fpn_strides_mismatch(self.fpn_strides, fpn_strides_live)
        roi_spatial_scales_live = [1.0 / s for s in fpn_strides_live]

        # Get device from first feature map (needed for device reference)
        images_tensor = torch.stack(images, dim=0)

        if self.training:
            if targets is None:
                raise ValueError("Targets required during training.")
            gt_boxes_list, gt_labels_list, gt_boxes_ignore_list = prepare_targets(targets, device=images_tensor.device)
            obj_logits, bbox_reg = self.rpn_head(feature_list)
            img_h, img_w = image_sizes[0]
            anchors = generate_oriented_anchors(
                image_size=(img_h, img_w),
                feature_map_sizes=feature_map_sizes,
                anchor_scales=self.anchor_scales,
                anchor_ratios=self.anchor_ratios,
                anchor_angles=self.anchor_angles,
                stride_per_level=fpn_strides_live,
            )
            rpn_losses = compute_midpoint_rpn_loss(
                objectness_logits=obj_logits,
                bbox_regression=bbox_reg,
                anchors=anchors,
                gt_boxes=gt_boxes_list,
                gt_boxes_ignore=gt_boxes_ignore_list,
                image_sizes=image_sizes,
                positive_iou_threshold=self.rpn_positive_iou_threshold,
                negative_iou_threshold=self.rpn_negative_iou_threshold,
                box_reg_weight=1.0,
                fg_bg_sampling_ratio=self.rpn_fg_bg_sampling_ratio,
                batch_size_per_image=self.rpn_batch_size_per_image,
                use_hbb_for_matching=self.use_hbb_for_matching,
                min_pos_iou=self.rpn_min_pos_iou,
                match_low_quality=self.rpn_match_low_quality,
                sample_from_all_levels=True,
            )
            proposals_list = generate_midpoint_proposals(
                obj_logits,
                bbox_reg,
                anchors,
                image_sizes,
                score_threshold=0.0,
                nms_threshold=self.rpn_nms_threshold,
                pre_nms_top_n=self.rpn_pre_nms_top_n,
                post_nms_top_n=self.rpn_post_nms_top_n,
                min_size=0.0,
            )
            # Optionally add GT boxes as proposals (MMRotate default).
            if self.add_gt_as_proposals:
                for img_idx in range(len(proposals_list)):
                    gt = gt_boxes_list[img_idx]
                    if gt.numel() > 0:
                        proposals_list[img_idx] = torch.cat([proposals_list[img_idx], gt.to(proposals_list[img_idx].device)], dim=0)
            box_to_img = torch.cat(
                [torch.full((p.shape[0],), i, dtype=torch.long, device=p.device) for i, p in enumerate(proposals_list)]
            ) if any(p.shape[0] > 0 for p in proposals_list) else torch.zeros(0, dtype=torch.long, device=images_tensor.device)
            oriented_proposals = torch.cat(proposals_list, dim=0).detach() if box_to_img.numel() > 0 else torch.zeros((0, 5), dtype=torch.float32, device=images_tensor.device)
            # Clamp proposal w/h to image scale so ROI loss does not explode on degenerate boxes
            min_side = 1.0
            for img_idx in range(len(image_sizes)):
                img_h, img_w = image_sizes[img_idx]
                max_side = float(max(img_w, img_h)) * 2.0
                mask = box_to_img == img_idx
                if mask.any():
                    oriented_proposals[mask, 2] = oriented_proposals[mask, 2].clamp(min=min_side, max=max_side)
                    oriented_proposals[mask, 3] = oriented_proposals[mask, 3].clamp(min=min_side, max=max_side)
            if oriented_proposals.shape[0] == 0:
                roi_losses = {"loss_classifier": torch.tensor(0.0, device=images_tensor.device), "loss_box_reg": torch.tensor(0.0, device=images_tensor.device)}
                return {**rpn_losses, **roi_losses}
            box_to_img = box_to_img.to(device=oriented_proposals.device)
            roi_feats = self.roi_align(
                feature_maps=feature_list,
                boxes=oriented_proposals,
                image_sizes=image_sizes,
                box_to_image=box_to_img,
                fpn_strides_override=fpn_strides_live,
                spatial_scales_override=roi_spatial_scales_live,
            )
            class_logits, box_reg = self.roi_head(roi_feats.view(roi_feats.shape[0], -1))
            roi_losses_list = []
            for img_idx in range(len(images)):
                mask = box_to_img == img_idx
                if mask.sum() == 0:
                    continue
                img_proposals = oriented_proposals[mask]
                img_logits = class_logits[mask]
                img_reg = box_reg[mask]
                img_gt = gt_boxes_list[img_idx]
                img_labels = gt_labels_list[img_idx]
                if len(img_gt) > 0:
                    roi_losses_list.append(
                        compute_oriented_roi_loss(
                            class_logits=img_logits, box_regression=img_reg, proposals=img_proposals,
                            gt_boxes=img_gt, gt_labels=img_labels,
                            gt_boxes_ignore=gt_boxes_ignore_list[img_idx],
                            positive_iou_threshold=self.roi_positive_iou_threshold,
                            negative_iou_threshold=self.roi_negative_iou_threshold,
                            box_reg_weight=1.0,
                            box_reg_angle_weight=self.roi_box_reg_angle_weight,
                            fg_bg_sampling_ratio=self.roi_fg_bg_sampling_ratio,
                            batch_size_per_image=self.roi_batch_size_per_image,
                            num_classes=self.num_classes, loss_type=self.roi_loss_type,
                            class_weights=self.roi_class_weights_tensor,
                            focal_alpha=self.roi_focal_alpha, focal_gamma=self.roi_focal_gamma,
                            target_means=self.target_means, target_stds=self.target_stds,
                            norm_factor=self.roi_norm_factor, edge_swap=self.roi_edge_swap,
                            proj_xy=self.roi_proj_xy,
                            box_reg_iou_weight=self.roi_box_reg_iou_weight,
                            box_reg_iou_loss_type=self.roi_box_reg_iou_loss_type,
                            box_reg_kfiou_fun=self.roi_box_reg_kfiou_fun,
                            box_reg_probiou_mode=self.roi_box_reg_probiou_mode,
                            use_hbb_for_matching=self.roi_use_hbb_for_matching,
                            reg_norm=self.roi_box_reg_norm,
                            match_low_quality=self.roi_match_low_quality,
                            label_smoothing=self.roi_label_smoothing,
                            **self.roi_grouped_ce_kwargs(),
                        )
                    )
            if roi_losses_list:
                roi_losses = {}
                for key in roi_losses_list[0].keys():
                    values = [l[key] for l in roi_losses_list if key in l]
                    if not values:
                        continue
                    if torch.is_tensor(values[0]):
                        roi_losses[key] = torch.stack(values).mean()
                    else:
                        roi_losses[key] = float(sum(values) / len(values))
            else:
                roi_losses = {"loss_classifier": torch.tensor(0.0, device=images_tensor.device), "loss_box_reg": torch.tensor(0.0, device=images_tensor.device)}
            return {**rpn_losses, **roi_losses}

        obj_logits, bbox_reg = self.rpn_head(feature_list)
        img_h, img_w = image_sizes[0]
        anchors = generate_oriented_anchors(
            image_size=(img_h, img_w),
            feature_map_sizes=feature_map_sizes,
            anchor_scales=self.anchor_scales,
            anchor_ratios=self.anchor_ratios,
            anchor_angles=self.anchor_angles,
            stride_per_level=fpn_strides_live,
        )
        proposals_list = generate_midpoint_proposals(
            obj_logits,
            bbox_reg,
            anchors,
            image_sizes,
            score_threshold=0.0,
            nms_threshold=self.rpn_nms_threshold,
            pre_nms_top_n=self.rpn_pre_nms_top_n,
            post_nms_top_n=self.rpn_post_nms_top_n,
            min_size=0.0,
        )
        return_debug = getattr(self, '_return_anchors_proposals', False)
        if return_debug:
            all_anchors = torch.cat(anchors, dim=0).detach().cpu()
        outputs = []
        from ..ops.rotated_ops import rotated_nms
        for img_idx, oriented in enumerate(proposals_list):
            if oriented.shape[0] == 0:
                out = {"rboxes": [], "labels": torch.zeros(0, dtype=torch.int64, device=images_tensor.device), "scores": torch.zeros(0, dtype=torch.float32, device=images_tensor.device)}
                if return_debug:
                    out["anchors"] = all_anchors
                    out["proposals"] = torch.zeros((0, 5), dtype=torch.float32)
                outputs.append(out)
                continue
            # Clamp proposal w/h to image scale (same as training) so refined boxes stay sane
            ih, iw = image_sizes[img_idx]
            max_side = float(max(iw, ih)) * 2.0
            oriented[:, 2] = oriented[:, 2].clamp(min=1.0, max=max_side)
            oriented[:, 3] = oriented[:, 3].clamp(min=1.0, max=max_side)
            box_to_img = torch.full((oriented.shape[0],), img_idx, dtype=torch.long, device=oriented.device)
            roi_feats = self.roi_align(
                feature_maps=feature_list,
                boxes=oriented,
                image_sizes=image_sizes,
                box_to_image=box_to_img,
                fpn_strides_override=fpn_strides_live,
                spatial_scales_override=roi_spatial_scales_live,
            )
            class_logits, box_reg = self.roi_head(roi_feats.view(roi_feats.shape[0], -1))
            class_probs = F.softmax(class_logits, dim=1)
            fg_probs = class_probs[:, 1:]
            score_threshold = self.inference_pre_nms_score_threshold
            if self.roi_inference_top_class_only:
                max_scores, argmax_cls = fg_probs.max(dim=1)
                keep_prop = max_scores > score_threshold
                if keep_prop.any():
                    proposal_indices = keep_prop.nonzero(as_tuple=True)[0]
                    filtered_proposals = oriented[proposal_indices]
                    filtered_scores = max_scores[proposal_indices]
                    filtered_class_indices = argmax_cls[proposal_indices]
                    filtered_labels = filtered_class_indices + 1
                    filtered_box_reg = box_reg[proposal_indices]
                else:
                    filtered_proposals = None
            else:
                candidate_mask = fg_probs > score_threshold  # [N, C]
                if candidate_mask.any():
                    proposal_indices, class_indices = candidate_mask.nonzero(as_tuple=True)
                    filtered_proposals = oriented[proposal_indices]
                    filtered_scores = fg_probs[proposal_indices, class_indices]
                    filtered_labels = class_indices + 1
                    filtered_box_reg = box_reg[proposal_indices]
                    filtered_class_indices = class_indices
                else:
                    filtered_proposals = None
            total_candidates = int(fg_probs.numel())
            from ..utils.logging import logger
            _roi_mode = "top_class" if self.roi_inference_top_class_only else "multiclass"
            logger.trace(
                "ROI inference candidates: mode={}, pre_nms_thr={:.3f}, total={}, kept_after_thr={}",
                _roi_mode,
                score_threshold,
                total_candidates,
                0 if filtered_proposals is None else int(filtered_proposals.shape[0]),
            )
            if filtered_proposals is not None:
                if self.roi_class_agnostic_regression:
                    selected_reg = filtered_box_reg
                else:
                    filtered_box_reg = filtered_box_reg.view(len(filtered_proposals), self.num_classes, 5)
                    selected_reg = filtered_box_reg[
                        torch.arange(len(filtered_proposals), device=filtered_proposals.device),
                        filtered_class_indices
                    ]
                refined = decode_oriented_boxes(
                    filtered_proposals, selected_reg,
                    target_means=self.target_means, target_stds=self.target_stds,
                    norm_factor=self.roi_norm_factor, edge_swap=self.roi_edge_swap,
                    proj_xy=self.roi_proj_xy,
                )
                if self.nms_class_agnostic:
                    keep = rotated_nms(
                        refined,
                        filtered_scores,
                        self.final_nms_iou_threshold,
                        self.max_detections_per_image,
                        force_cpu=self.final_nms_use_cpu,
                    )
                    if len(keep) > 0:
                        k_scores = filtered_scores[keep]
                        _, order = k_scores.sort(descending=True)
                        keep = keep[order]
                    kept_per_class = {"agnostic": int(len(keep))}
                else:
                    # Per-class cap is an optimization; final per-image cap applied after merge.
                    keep_per_class = []
                    kept_per_class = {}
                    for cls in filtered_labels.unique():
                        cls_mask = filtered_labels == cls
                        cls_indices = cls_mask.nonzero(as_tuple=True)[0]
                        cls_keep = rotated_nms(
                            refined[cls_mask],
                            filtered_scores[cls_mask],
                            self.final_nms_iou_threshold,
                            self.max_detections_per_image,
                            force_cpu=self.final_nms_use_cpu,
                        )
                        kept_per_class[int(cls.item())] = int(len(cls_keep))
                        keep_per_class.append(cls_indices[cls_keep])
                    if keep_per_class:
                        keep = torch.cat(keep_per_class, dim=0)
                        keep_scores = filtered_scores[keep]
                        _, keep_order = keep_scores.sort(descending=True)
                        keep = keep[keep_order]
                        if self.max_detections_per_image is not None:
                            keep = keep[:self.max_detections_per_image]
                    else:
                        keep = torch.zeros((0,), dtype=torch.long, device=filtered_scores.device)
                final_boxes = refined[keep]
                final_scores = filtered_scores[keep]
                final_labels = filtered_labels[keep]
                logger.trace(
                    "ROI inference kept after {} NMS: total={}, per_class={}",
                    "class-agnostic" if self.nms_class_agnostic else "class-wise",
                    int(len(keep)),
                    kept_per_class,
                )
                valid = (final_boxes[:, 2] >= 1.0) & (final_boxes[:, 3] >= 1.0)
                final_boxes, final_scores, final_labels = final_boxes[valid], final_scores[valid], final_labels[valid]
                out = {"rboxes": tensor_to_rboxes(final_boxes), "labels": final_labels, "scores": final_scores}
                if return_debug:
                    out["anchors"] = all_anchors
                    out["proposals"] = oriented.detach().cpu()
                outputs.append(out)
            else:
                out = {"rboxes": [], "labels": torch.zeros(0, dtype=torch.int64, device=images_tensor.device), "scores": torch.zeros(0, dtype=torch.float32, device=images_tensor.device)}
                if return_debug:
                    out["anchors"] = all_anchors
                    out["proposals"] = oriented.detach().cpu()
                outputs.append(out)
        return outputs

set_final_nms_iou_for_epoch(epoch)

Update final detection NMS IoU threshold from schedule. Lower = more aggressive suppression.

Source code in oriented_det/models/oriented_rcnn.py 
928
929
930
931
932
933
934
935
936
937
938
939
940
def set_final_nms_iou_for_epoch(self, epoch: int) -> None:
    """Update final detection NMS IoU threshold from schedule. Lower = more aggressive suppression."""
    if self._final_nms_iou_schedule_epochs is None or self._final_nms_iou_schedule_values is None:
        return
    if not self._final_nms_iou_schedule_epochs or not self._final_nms_iou_schedule_values:
        return
    idx = 0
    for boundary in self._final_nms_iou_schedule_epochs:
        if epoch < boundary:
            break
        idx += 1
    idx = min(idx, len(self._final_nms_iou_schedule_values) - 1)
    self.final_nms_iou_threshold = self._final_nms_iou_schedule_values[idx]

set_roi_box_reg_angle_weight_for_epoch(epoch)

Update ROI angle (5th dim) SmoothL1 weight from schedule (0-based epoch index).

Source code in oriented_det/models/oriented_rcnn.py 
953
954
955
956
957
958
959
960
961
962
def set_roi_box_reg_angle_weight_for_epoch(self, epoch: int) -> None:
    """Update ROI angle (5th dim) SmoothL1 weight from schedule (0-based epoch index)."""
    from oriented_det.train.piecewise_schedule import resolve_piecewise_schedule

    self.roi_box_reg_angle_weight = resolve_piecewise_schedule(
        epoch,
        self._roi_box_reg_angle_schedule_epochs,
        self._roi_box_reg_angle_schedule_values,
        self._roi_box_reg_angle_weight_default,
    )

set_roi_box_reg_iou_weight_for_epoch(epoch)

Update ROI auxiliary IoU loss weight from schedule (0-based epoch index).

Source code in oriented_det/models/oriented_rcnn.py 
942
943
944
945
946
947
948
949
950
951
def set_roi_box_reg_iou_weight_for_epoch(self, epoch: int) -> None:
    """Update ROI auxiliary IoU loss weight from schedule (0-based epoch index)."""
    from oriented_det.train.piecewise_schedule import resolve_piecewise_schedule

    self.roi_box_reg_iou_weight = resolve_piecewise_schedule(
        epoch,
        self._roi_box_reg_iou_schedule_epochs,
        self._roi_box_reg_iou_schedule_values,
        self._roi_box_reg_iou_weight_default,
    )

RotatedFasterRCNN

Bases: ClassWeightsMixin, GroupedCeMixin, Module

Rotated Faster R-CNN (MMRotate-style): horizontal RPN + horizontal RoIAlign + rotated ROI head.

Mimics MMRotate's Rotated Faster R-CNN: - Horizontal RPN with 4 reg params (DeltaXYWHBBoxCoder) - Horizontal RoIAlign - ROI head with 5 params (DeltaXYWHTHBBoxCoder)

Parameters:

Name Type Description Default
num_classes int

Number of object classes (excluding background)

 required
backbone

Optional backbone module (if None, creates ResNet+FPN)

 None
backbone_name str

Name of backbone to create ("resnet18", "resnet50", etc.)

 'resnet50'
pretrained_backbone bool

Whether to use pretrained backbone weights

 False
trainable_layers int

Number of backbone layers to keep trainable. For ResNet50, typically 5 layers (conv1 + layer1-4). Set to 5 to train all layers. Default is 5 (all layers trainable).

 5
anchor_scales Optional[List[float]]

List of anchor scales for RPN

 None
anchor_ratios Optional[List[float]]

List of anchor aspect ratios for RPN

 None
rpn_pre_nms_top_n int

Number of top proposals before NMS in RPN

 2000
rpn_post_nms_top_n int

Number of top proposals after NMS in RPN

 2000
rpn_positive_iou_threshold float

IoU threshold for positive RPN anchors

 0.7
rpn_negative_iou_threshold float

IoU threshold for negative RPN anchors

 0.3
roi_positive_iou_threshold float

IoU threshold for positive ROI proposals

 0.5
roi_negative_iou_threshold float

IoU threshold for negative ROI proposals

 0.5
roi_output_size Tuple[int, int]

Output size for ROI align (height, width)

 (7, 7)
roi_spatial_scales Optional[List[float]]

Spatial scales for each FPN level in ROI align (defaults from fpn_strides at init)

 None
fpn_strides Optional[List[int]]

Nominal strides per FPN level for ROI module init; forward pass overrides with strides derived from image size and feature map shapes so RPN anchors and ROI align match the grid.

 None
use_checkpoint bool

Whether to use gradient checkpointing (trades compute for memory)

 False
roi_chunk_size int

Number of ROIs to process in parallel during ROI align. Lower values use less memory but are slower. Recommended: 16-64 for typical GPU memory (8-24GB).

 32
roi_use_checkpoint bool

If True, use gradient checkpointing in ROI align to trade compute for memory (~2x less memory, ~30% slower).

 False
roi_loss_type str

Type of classification loss: - "cross_entropy": Standard cross-entropy loss (default) - "focal": Focal loss for handling class imbalance and hard examples

 'cross_entropy'
roi_class_weights Optional[Union[Dict[str, float], Tensor]]

Optional class weights for loss computation: - Dict[str, float]: Mapping from class name to weight (requires calling set_class_weights() with class_map) - torch.Tensor: Tensor of shape [num_classes + 1] (including background) - None: Equal weights for all classes (default)

 None
roi_focal_alpha float

Alpha parameter for focal loss (default: 1.0)

 1.0
roi_focal_gamma float

Gamma parameter for focal loss (default: 2.0) - gamma=0: equivalent to cross-entropy - gamma>0: focuses on hard examples - gamma=2: standard value used in literature

 2.0
Example

Standard cross-entropy loss

model = OrientedRCNN(num_classes=15)

Class-weighted cross-entropy loss

class_weights = {"small-vehicle": 0.5, "plane": 2.0, ...} model = OrientedRCNN(num_classes=15, roi_class_weights=class_weights) model.set_class_weights(class_map) # Must call before training

Focal loss

model = OrientedRCNN(num_classes=15, roi_loss_type="focal", roi_focal_gamma=2.0)

Focal loss with class weights

model = OrientedRCNN( ... num_classes=15, ... roi_loss_type="focal", ... roi_class_weights=class_weights, ... roi_focal_gamma=2.0 ... ) model.set_class_weights(class_map)

Source code in oriented_det/models/oriented_rcnn.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
class RotatedFasterRCNN(ClassWeightsMixin, GroupedCeMixin, nn.Module):
    """Rotated Faster R-CNN (MMRotate-style): horizontal RPN + horizontal RoIAlign + rotated ROI head.

    Mimics MMRotate's Rotated Faster R-CNN:
    - Horizontal RPN with 4 reg params (DeltaXYWHBBoxCoder)
    - Horizontal RoIAlign
    - ROI head with 5 params (DeltaXYWHTHBBoxCoder)

    Args:
        num_classes: Number of object classes (excluding background)
        backbone: Optional backbone module (if None, creates ResNet+FPN)
        backbone_name: Name of backbone to create ("resnet18", "resnet50", etc.)
        pretrained_backbone: Whether to use pretrained backbone weights
        trainable_layers: Number of backbone layers to keep trainable. For ResNet50,
                         typically 5 layers (conv1 + layer1-4). Set to 5 to train all layers.
                         Default is 5 (all layers trainable).
        anchor_scales: List of anchor scales for RPN
        anchor_ratios: List of anchor aspect ratios for RPN
        rpn_pre_nms_top_n: Number of top proposals before NMS in RPN
        rpn_post_nms_top_n: Number of top proposals after NMS in RPN
        rpn_positive_iou_threshold: IoU threshold for positive RPN anchors
        rpn_negative_iou_threshold: IoU threshold for negative RPN anchors
        roi_positive_iou_threshold: IoU threshold for positive ROI proposals
        roi_negative_iou_threshold: IoU threshold for negative ROI proposals
        roi_output_size: Output size for ROI align (height, width)
        roi_spatial_scales: Spatial scales for each FPN level in ROI align (defaults from ``fpn_strides`` at init)
        fpn_strides: Nominal strides per FPN level for ROI module init; forward pass overrides with strides
            derived from image size and feature map shapes so RPN anchors and ROI align match the grid.
        use_checkpoint: Whether to use gradient checkpointing (trades compute for memory)
        roi_chunk_size: Number of ROIs to process in parallel during ROI align.
                       Lower values use less memory but are slower.
                       Recommended: 16-64 for typical GPU memory (8-24GB).
        roi_use_checkpoint: If True, use gradient checkpointing in ROI align to
                           trade compute for memory (~2x less memory, ~30% slower).
        roi_loss_type: Type of classification loss:
                      - "cross_entropy": Standard cross-entropy loss (default)
                      - "focal": Focal loss for handling class imbalance and hard examples
        roi_class_weights: Optional class weights for loss computation:
                          - Dict[str, float]: Mapping from class name to weight
                            (requires calling set_class_weights() with class_map)
                          - torch.Tensor: Tensor of shape [num_classes + 1] (including background)
                          - None: Equal weights for all classes (default)
        roi_focal_alpha: Alpha parameter for focal loss (default: 1.0)
        roi_focal_gamma: Gamma parameter for focal loss (default: 2.0)
                         - gamma=0: equivalent to cross-entropy
                         - gamma>0: focuses on hard examples
                         - gamma=2: standard value used in literature

    Example:
        >>> # Standard cross-entropy loss
        >>> model = OrientedRCNN(num_classes=15)

        >>> # Class-weighted cross-entropy loss
        >>> class_weights = {"small-vehicle": 0.5, "plane": 2.0, ...}
        >>> model = OrientedRCNN(num_classes=15, roi_class_weights=class_weights)
        >>> model.set_class_weights(class_map)  # Must call before training

        >>> # Focal loss
        >>> model = OrientedRCNN(num_classes=15, roi_loss_type="focal", roi_focal_gamma=2.0)

        >>> # Focal loss with class weights
        >>> model = OrientedRCNN(
        ...     num_classes=15,
        ...     roi_loss_type="focal",
        ...     roi_class_weights=class_weights,
        ...     roi_focal_gamma=2.0
        ... )
        >>> model.set_class_weights(class_map)
    """

    def __init__(
        self,
        num_classes: int,
        *,
        backbone=None,
        backbone_name: str = "resnet50",
        pretrained_backbone: bool = False,
        trainable_layers: int = 5,
        anchor_scales: Optional[List[float]] = None,
        anchor_ratios: Optional[List[float]] = None,
        anchor_angles: Optional[List[float]] = None,
        rpn_pre_nms_top_n: int = 2000,
        rpn_post_nms_top_n: int = 2000,
        rpn_min_size: float = 0.0,  # Minimum proposal side length (pixels); use 2–4 to filter tiny boxes
        max_detections_per_image: int = 2000,  # MMRotate test_cfg rcnn max_per_img
        rpn_nms_threshold: float = 0.7,  # IoU threshold for RPN proposal NMS
        final_nms_iou_threshold: float = 0.1,  # MMRotate rotated NMS iou (DOTA le90 config)
        rpn_positive_iou_threshold: float = 0.7,  # MMRotate RPN MaxIoUAssigner pos_iou_thr
        rpn_negative_iou_threshold: float = 0.3,  # MMRotate RPN neg_iou_thr
        rpn_fg_bg_sampling_ratio: float = 0.5,  # Ratio of foreground to background in sampled RPN anchors
        rpn_batch_size_per_image: int = 256,  # Number of RPN anchors to sample per image (MMRotate: 256)
        rpn_min_pos_iou: float = 0.3,  # Min IoU for best-anchor-per-GT to be positive (MMRotate RPN assigner)
        rpn_match_low_quality: bool = True,  # If True, force best anchor per GT to positive when IoU >= rpn_min_pos_iou
        roi_match_low_quality: bool = False,  # MMRotate RCNN assigner match_low_quality (usually False)
        roi_min_pos_iou: float = 0.5,  # MMRotate RCNN MaxIoUAssigner min_pos_iou
        add_gt_as_proposals: bool = True,  # If True, append GT boxes to RPN proposals in training (MMRotate default)
        roi_positive_iou_threshold: float = 0.5,  # MMRotate RCNN pos_iou_thr
        roi_negative_iou_threshold: float = 0.5,  # MMRotate RCNN neg_iou_thr (same as pos for DOTA config)
        roi_fg_bg_sampling_ratio: float = 0.25,  # Ratio of foreground to background in sampled ROI proposals
        roi_batch_size_per_image: int = 512,  # Number of ROI proposals to sample per image for training
        roi_output_size: Tuple[int, int] = (7, 7),
        roi_spatial_scales: Optional[List[float]] = None,
        fpn_strides: Optional[List[int]] = None,
        returned_layers: Optional[List[int]] = None,
        use_checkpoint: bool = False,
        roi_chunk_size: int = 32,
        roi_use_checkpoint: bool = False,
        # Classification loss configuration
        roi_loss_type: str = "cross_entropy",
        roi_class_weights: Optional[Union[Dict[str, float], torch.Tensor]] = None,
        roi_focal_alpha: float = 1.0,
        roi_focal_gamma: float = 2.0,
        roi_label_smoothing: float = 0.0,
        # Box regression target normalization (like MMRotate: only ROI head uses small stds)
        target_means: Optional[Tuple[float, float, float, float, float]] = None,
        target_stds: Optional[Tuple[float, float, float, float, float]] = None,
        roi_norm_factor: float = 2.0,
        roi_edge_swap: bool = True,
        roi_proj_xy: bool = False,
        roi_box_reg_angle_weight: float = 1.0,
        roi_box_reg_angle_schedule_epochs: Optional[List[int]] = None,
        roi_box_reg_angle_schedule_values: Optional[List[float]] = None,
        roi_box_reg_iou_weight: float = 0.0,
        roi_box_reg_iou_loss_type: str = "riou",
        roi_box_reg_kfiou_fun: Optional[str] = None,
        roi_box_reg_probiou_mode: Optional[str] = None,
        roi_box_reg_main_loss_type: str = "smooth_l1",
        roi_box_reg_norm: str = "sampled_all",
        roi_box_reg_smooth_l1_aux_weight: float = 0.0,
        use_hbb_for_matching: bool = False,
        inference_pre_nms_score_threshold: float = 0.05,
        final_nms_iou_schedule_epochs: Optional[List[int]] = None,
        final_nms_iou_schedule_values: Optional[List[float]] = None,
        roi_box_reg_iou_schedule_epochs: Optional[List[int]] = None,
        roi_box_reg_iou_schedule_values: Optional[List[float]] = None,
        nms_class_agnostic: bool = False,
        final_nms_use_cpu: bool = False,
        roi_inference_top_class_only: bool = False,
    ):
        """Initialize Rotated Faster R-CNN model.

        Args:
            num_classes: Number of object classes (excluding background)
            backbone: Optional backbone module (if None, creates ResNet+FPN)
            backbone_name: Name of backbone to create ("resnet18", "resnet50", etc.)
            pretrained_backbone: Whether to use pretrained backbone weights
            trainable_layers: Number of backbone layers to keep trainable. For ResNet50,
                             typically 5 layers (conv1 + layer1-4). Set to 5 to train all layers.
                             Default is 5 (all layers trainable).
            anchor_scales: List of anchor scales for RPN
            anchor_ratios: List of anchor aspect ratios for RPN
            anchor_angles: Optional RPN angles in radians (Python API only; not in JSON). ``None`` -> ``[0.0]``.
            rpn_pre_nms_top_n: Number of top proposals before NMS in RPN
            rpn_post_nms_top_n: Number of top proposals after NMS in RPN
            rpn_positive_iou_threshold: IoU threshold for positive RPN anchors
            rpn_negative_iou_threshold: IoU threshold for negative RPN anchors
            rpn_fg_bg_sampling_ratio: Ratio of foreground to background in sampled RPN anchors (default: 0.5)
            roi_positive_iou_threshold: IoU threshold for positive ROI proposals
            roi_negative_iou_threshold: IoU threshold for negative ROI proposals
            roi_fg_bg_sampling_ratio: Ratio of foreground to background in sampled ROI proposals (default: 0.25)
            roi_output_size: Output size for ROI align (height, width)
            roi_spatial_scales: Spatial scales for each FPN level in ROI align (defaults from ``fpn_strides`` at init)
            fpn_strides: Nominal strides for ROI module init; forward uses strides derived from the feature grid.
            use_checkpoint: Whether to use gradient checkpointing (trades compute for memory)
            roi_chunk_size: Number of ROIs to process in parallel during ROI align
            roi_use_checkpoint: If True, use gradient checkpointing in ROI align
            roi_loss_type: Type of classification loss:
                          - "cross_entropy": Standard cross-entropy loss (default)
                          - "focal": Focal loss for handling class imbalance
            roi_class_weights: Optional class weights for loss computation.
                              Can be:
                              - Dict[str, float]: Mapping from class name to weight
                              - torch.Tensor: Tensor of shape [num_classes + 1] (including background)
                              - None: Equal weights for all classes (default)
                              Note: If dict is provided, it will be converted to tensor
                                    using the model's class mapping. Background (index 0)
                                    defaults to weight 1.0 if not specified.
            roi_focal_alpha: Alpha parameter for focal loss (default: 1.0)
            roi_focal_gamma: Gamma parameter for focal loss (default: 2.0)
            roi_label_smoothing: Label smoothing for ROI classification (default: 0.0). Use e.g. 0.1 to reduce overconfidence.
            target_means: Optional means for regression target normalization [dx, dy, dw, dh, dangle].
                         Default: (0.0, 0.0, 0.0, 0.0, 0.0) - no centering (like MMRotate default).
                         Can be set to empirical means computed from dataset for better normalization.
            target_stds: Optional stds for regression target normalization [dx, dy, dw, dh, dangle].
                        Default: (0.1, 0.1, 0.2, 0.2, 0.1) (MMRotate DeltaXYWHTHBBoxCoder).
                        Can be set to empirical stds computed from dataset for better normalization.
            roi_norm_factor: Angle scaling for ROI encode/decode (MMRotate uses 2.0).
            roi_edge_swap: Whether to use edge_swap for ROI (MMRotate uses True).
            roi_proj_xy: If True, encode/decode ROI dx/dy in proposal local frame. For
                horizontal xyxy RoIs (angle 0) this is equivalent to global offsets.
            roi_box_reg_angle_weight: Weight for the angle (5th) component in ROI box regression loss.
                Use > 1.0 (e.g. 2.0) to improve orientation alignment when predictions stick to anchor angles.
            roi_box_reg_iou_weight: Weight for auxiliary decoded-box loss when main is Smooth L1.
            roi_box_reg_iou_loss_type: ``riou`` (default), ``kfiou``, or ``probiou`` for decoded aux/main.
            roi_box_reg_kfiou_fun: Optional KFIoU overlap mapping when using ``kfiou``.
            roi_box_reg_probiou_mode: ``l1`` (default) or ``l2`` when using ``probiou``.
            roi_box_reg_main_loss_type: Primary ROI reg loss: ``smooth_l1`` (default) or decoded
                ``probiou`` / ``riou`` / ``kfiou``.
            roi_box_reg_norm: ``sampled_all`` (MMRotate avg_factor) or ``positives_only``.
            roi_box_reg_smooth_l1_aux_weight: Encoded Smooth L1 aux when main is decoded.
            roi_min_pos_iou: Min IoU for low-quality ROI matches when ``roi_match_low_quality`` is True.
            roi_inference_top_class_only: If True, ROI inference uses argmax fg class per proposal
                (see ``oriented_det/models/README.md``); if False, multiclass thresholding.
            RPN uses fixed (0,0,0,0) and (1,1,1,1) per MMRotate DeltaXYWHBBoxCoder.
        """
        if torch is None or nn is None:
            raise RuntimeError("PyTorch is required for RotatedFasterRCNN.")

        # Initialize ClassWeightsMixin first, then nn.Module
        # ClassWeightsMixin.__init__() sets self.num_classes
        super().__init__(num_classes, roi_class_weights)
        nn.Module.__init__(self)
        self._init_grouped_ce()

        # Box regression: ROI head uses target_means/stds/norm_factor; RPN uses identity (MMRotate)
        self.target_means = target_means if target_means is not None else (0.0, 0.0, 0.0, 0.0, 0.0)
        self.target_stds = target_stds if target_stds is not None else (0.1, 0.1, 0.2, 0.2, 0.1)
        self.roi_norm_factor = roi_norm_factor
        self.roi_edge_swap = roi_edge_swap
        self.roi_proj_xy = roi_proj_xy
        self.roi_box_reg_iou_loss_type = roi_box_reg_iou_loss_type
        self.roi_box_reg_kfiou_fun = roi_box_reg_kfiou_fun
        self.roi_box_reg_probiou_mode = roi_box_reg_probiou_mode
        self.roi_box_reg_main_loss_type = roi_box_reg_main_loss_type
        self.roi_box_reg_norm = roi_box_reg_norm
        self.roi_box_reg_smooth_l1_aux_weight = float(roi_box_reg_smooth_l1_aux_weight)
        self.use_hbb_for_matching = use_hbb_for_matching
        self.inference_pre_nms_score_threshold = inference_pre_nms_score_threshold
        self._roi_box_reg_angle_schedule_epochs = roi_box_reg_angle_schedule_epochs
        self._roi_box_reg_angle_schedule_values = roi_box_reg_angle_schedule_values
        self._roi_box_reg_angle_weight_default = float(roi_box_reg_angle_weight)
        self.roi_box_reg_angle_weight = float(roi_box_reg_angle_weight)
        self.set_roi_box_reg_angle_weight_for_epoch(0)

        # Store loss configuration
        self.roi_loss_type = roi_loss_type
        self.roi_focal_alpha = roi_focal_alpha
        self.roi_focal_gamma = roi_focal_gamma
        self.roi_label_smoothing = roi_label_smoothing

        # Setup backbone using shared utility (returned_layers=[2,3,4] for MMRotate-style C3–C5 only)
        self.backbone, backbone_channels = setup_backbone(
            backbone=backbone,
            backbone_name=backbone_name,
            pretrained_backbone=pretrained_backbone,
            trainable_layers=trainable_layers,
            returned_layers=returned_layers,
        )
        # FPN strides: default [4,8,16,32,64] for C2–C5; [8,16,32,64] when returned_layers=[2,3,4]
        if fpn_strides is None and returned_layers == [2, 3, 4]:
            fpn_strides = [8, 16, 32, 64]

        # Default horizontal RPN priors; optional anchor_angles for advanced Python use only.
        self.anchor_scales, self.anchor_ratios, self.anchor_angles, self.num_anchors = setup_anchors(
            anchor_scales=anchor_scales,
            anchor_ratios=anchor_ratios,
            anchor_angles=anchor_angles,
            default_angles=[0.0],
        )

        # RPN head configuration
        # MMRotate format: cls_out_channels=1, reg_out_channels=4 (dx, dy, dw, dh - angle from anchor)
        rpn_cls_out_channels = 1
        rpn_reg_out_channels = 4

        # RPN head
        self.rpn_head = OrientedRPNHead(
            backbone_channels, 
            num_anchors=self.num_anchors,
            cls_out_channels=rpn_cls_out_channels,
            reg_out_channels=rpn_reg_out_channels,
        )

        # ROI components (MMRotate Rotated Faster R-CNN uses horizontal RoIAlign)
        self.roi_align = HorizontalROIAlign(
            output_size=roi_output_size,
            spatial_scales=roi_spatial_scales,
            fpn_strides=fpn_strides,
            chunk_size=roi_chunk_size,
        )

        # ROI head input size: output_size[0] * output_size[1] * backbone_channels
        roi_feature_size = roi_output_size[0] * roi_output_size[1] * backbone_channels
        # Use class-agnostic regression (MMRotate format) - better for initial training when class prediction is weak
        self.roi_class_agnostic_regression = True  # Store flag for inference
        self.roi_head = OrientedROIHead(
            in_channels=roi_feature_size,
            num_classes=num_classes,
            class_agnostic_regression=True,  # MMRotate format: class-agnostic regression
        )

        # Configuration
        self.rpn_pre_nms_top_n = rpn_pre_nms_top_n
        self.rpn_post_nms_top_n = rpn_post_nms_top_n
        self.rpn_min_size = rpn_min_size
        self.max_detections_per_image = max_detections_per_image
        self.rpn_nms_threshold = rpn_nms_threshold
        self.final_nms_iou_threshold = final_nms_iou_threshold
        self.nms_class_agnostic = nms_class_agnostic
        self.final_nms_use_cpu = final_nms_use_cpu
        self.roi_inference_top_class_only = roi_inference_top_class_only
        self.rpn_positive_iou_threshold = rpn_positive_iou_threshold
        self.rpn_negative_iou_threshold = rpn_negative_iou_threshold
        self.rpn_fg_bg_sampling_ratio = rpn_fg_bg_sampling_ratio
        self.rpn_batch_size_per_image = rpn_batch_size_per_image
        self.rpn_min_pos_iou = rpn_min_pos_iou
        self.rpn_match_low_quality = rpn_match_low_quality
        self.roi_match_low_quality = roi_match_low_quality
        self.roi_min_pos_iou = roi_min_pos_iou
        self.add_gt_as_proposals = add_gt_as_proposals
        self.roi_positive_iou_threshold = roi_positive_iou_threshold
        self.roi_negative_iou_threshold = roi_negative_iou_threshold
        self.roi_fg_bg_sampling_ratio = roi_fg_bg_sampling_ratio
        self.roi_batch_size_per_image = roi_batch_size_per_image
        self.fpn_strides = fpn_strides or [4, 8, 16, 32, 64]
        self.roi_spatial_scales = roi_spatial_scales or [1.0 / s for s in self.fpn_strides]

        # Gradient checkpointing (trades computation for memory)
        # When enabled, recomputes activations during backward pass instead of storing them
        # Reduces memory usage by 2-3x at the cost of ~30% slower backward pass
        self.use_checkpoint = use_checkpoint

        # NMS IoU schedule for final detection NMS (optional; essential for RetinaNet, optional for R-CNN)
        self._final_nms_iou_schedule_epochs = final_nms_iou_schedule_epochs
        self._final_nms_iou_schedule_values = final_nms_iou_schedule_values
        self._roi_box_reg_iou_schedule_epochs = roi_box_reg_iou_schedule_epochs
        self._roi_box_reg_iou_schedule_values = roi_box_reg_iou_schedule_values
        self._roi_box_reg_iou_weight_default = float(roi_box_reg_iou_weight)
        self.roi_box_reg_iou_weight = float(roi_box_reg_iou_weight)
        self.set_roi_box_reg_iou_weight_for_epoch(0)

    def set_final_nms_iou_for_epoch(self, epoch: int) -> None:
        """Update final detection NMS IoU threshold from schedule. Lower = more aggressive suppression."""
        if self._final_nms_iou_schedule_epochs is None or self._final_nms_iou_schedule_values is None:
            return
        if not self._final_nms_iou_schedule_epochs or not self._final_nms_iou_schedule_values:
            return
        idx = 0
        for boundary in self._final_nms_iou_schedule_epochs:
            if epoch < boundary:
                break
            idx += 1
        idx = min(idx, len(self._final_nms_iou_schedule_values) - 1)
        self.final_nms_iou_threshold = self._final_nms_iou_schedule_values[idx]

    def set_roi_box_reg_iou_weight_for_epoch(self, epoch: int) -> None:
        """Update ROI auxiliary IoU loss weight from schedule (0-based epoch index)."""
        from oriented_det.train.piecewise_schedule import resolve_piecewise_schedule

        self.roi_box_reg_iou_weight = resolve_piecewise_schedule(
            epoch,
            self._roi_box_reg_iou_schedule_epochs,
            self._roi_box_reg_iou_schedule_values,
            self._roi_box_reg_iou_weight_default,
        )

    def set_roi_box_reg_angle_weight_for_epoch(self, epoch: int) -> None:
        """Update ROI angle (5th dim) SmoothL1 weight from schedule (0-based epoch index)."""
        from oriented_det.train.piecewise_schedule import resolve_piecewise_schedule

        self.roi_box_reg_angle_weight = resolve_piecewise_schedule(
            epoch,
            self._roi_box_reg_angle_schedule_epochs,
            self._roi_box_reg_angle_schedule_values,
            self._roi_box_reg_angle_weight_default,
        )

    def forward(
        self,
        images: Sequence[torch.Tensor],
        targets: Optional[Sequence[Dict[str, Any]]] = None,
    ) -> Union[Dict[str, torch.Tensor], List[Dict[str, Any]]]:
        """Forward pass through oriented R-CNN.

        Args:
            images: List of image tensors (C, H, W) in [0, 1] range
            targets: Optional list of target dicts for training.
                    Each dict must contain:
                    - "rboxes" (List[RBox] or tensor [N, 5] with format [cx, cy, w, h, angle])
                    - "labels" (tensor [N])

        Returns:
            - Training: Dict with loss keys:
                - "loss_objectness": RPN classification loss
                - "loss_rpn_box_reg": RPN box regression loss
                - "loss_classifier": ROI classification loss
                - "loss_box_reg": ROI box regression loss
            - Inference: List of dicts, one per image, containing:
                - "rboxes": List[RBox] with oriented boxes
                - "rboxes": List[RBox] or Tensor [N, 5] with oriented boxes [cx, cy, w, h, angle]
                - "labels": Tensor [N] with class labels
                - "scores": Tensor [N] with confidence scores
        """
        if not isinstance(images, (list, tuple)):
            images = [images]

        # Get image sizes
        image_sizes = [(img.shape[-2], img.shape[-1]) for img in images]
        timing_enabled = bool(
            TRACE_FIRST_TRAIN_FORWARD_TIMING
            and self.training
            and not getattr(self, "_timed_first_train_forward", False)
        )
        if timing_enabled:
            self._timed_first_train_forward = True
            timing_start = time.perf_counter()
            timing_last = timing_start

            def _timing_mark(label: str) -> None:
                nonlocal timing_last
                if torch is not None and torch.cuda.is_available():
                    try:
                        torch.cuda.synchronize()
                    except Exception:
                        pass
                now = time.perf_counter()
                print(f"[first-batch timing] {label}: +{now - timing_last:.3f}s total={now - timing_start:.3f}s", flush=True)
                timing_last = now

            print("[first-batch timing] begin RotatedFasterRCNN.forward", flush=True)
        else:
            def _timing_mark(label: str) -> None:
                return None

        # Extract features using shared utility. include_pool_level keeps torchvision's
        # P6 (stride-64 max-pool level) so the RPN sees 5 levels like MMRotate;
        # horizontal_roi_align still restricts ROI extraction to the first 4 levels.
        feature_list = extract_backbone_features(
            self.backbone,
            images,
            use_checkpoint=self.use_checkpoint,
            training=self.training,
            include_pool_level=True,
        )
        feature_map_sizes = [(f.shape[2], f.shape[3]) for f in feature_list]
        fpn_strides_live = derive_fpn_strides_from_grid(image_sizes[0], feature_map_sizes)
        warn_if_fpn_strides_mismatch(self.fpn_strides, fpn_strides_live)
        roi_spatial_scales_live = [1.0 / s for s in fpn_strides_live]
        _timing_mark("backbone_fpn")

        # Get device from first feature map
        images_tensor = torch.stack(images, dim=0)  # [B, C, H, W] - needed for device reference

        if self.training:
            if targets is None:
                raise ValueError("Targets required during training.")

            # Prepare targets (preserve oriented boxes)
            gt_boxes_list, gt_labels_list, gt_boxes_ignore_list = prepare_targets(targets, device=images_tensor.device)
            _timing_mark("prepare_targets")

            # RPN forward
            objectness_logits, bbox_regression = self.rpn_head(feature_list)
            _timing_mark("rpn_head")

            # Generate anchors (strides from actual grid — see fpn_strides_live)
            img_h, img_w = image_sizes[0]
            anchors = generate_oriented_anchors(
                image_size=(img_h, img_w),
                feature_map_sizes=feature_map_sizes,
                anchor_scales=self.anchor_scales,
                anchor_ratios=self.anchor_ratios,
                anchor_angles=self.anchor_angles,
                stride_per_level=fpn_strides_live,
            )
            _timing_mark("generate_anchors")

            # Compute RPN loss (4 params, no angle - MMRotate style)
            rpn_losses = compute_oriented_rpn_loss(
                objectness_logits=objectness_logits,
                bbox_regression=bbox_regression,
                anchors=anchors,
                gt_boxes=gt_boxes_list,
                gt_boxes_ignore=gt_boxes_ignore_list,
                image_sizes=image_sizes,
                positive_iou_threshold=self.rpn_positive_iou_threshold,
                negative_iou_threshold=self.rpn_negative_iou_threshold,
                box_reg_weight=1.0,
                fg_bg_sampling_ratio=self.rpn_fg_bg_sampling_ratio,
                batch_size_per_image=self.rpn_batch_size_per_image,
                min_pos_iou=self.rpn_min_pos_iou,
                match_low_quality=self.rpn_match_low_quality,
                target_means=(0.0, 0.0, 0.0, 0.0),
                target_stds=(1.0, 1.0, 1.0, 1.0),
                use_horizontal_targets=True,
                use_hbb_for_matching=self.use_hbb_for_matching,
            )
            _timing_mark("rpn_loss")

            # Horizontal proposals (xyxy) — MMRotate semantics. No objectness score
            # threshold: MMRotate's RPN keeps top-k proposals by score only;
            # inference_pre_nms_score_threshold applies to ROI-head scores, not the RPN.
            proposals_xyxy = generate_horizontal_proposals(
                objectness_logits=objectness_logits,
                bbox_regression=bbox_regression,
                anchors=anchors,
                image_sizes=image_sizes,
                score_threshold=0.0,
                nms_threshold=self.rpn_nms_threshold,
                pre_nms_top_n=self.rpn_pre_nms_top_n,
                post_nms_top_n=self.rpn_post_nms_top_n,
                min_size=self.rpn_min_size,
                target_means=(0.0, 0.0, 0.0, 0.0),
                target_stds=(1.0, 1.0, 1.0, 1.0),
            )
            _timing_mark("generate_proposals")
            if timing_enabled:
                proposal_counts = [int(len(p)) for p in proposals_xyxy]
                print(
                    f"[first-batch timing] proposals_after_cap: "
                    f"images={len(proposal_counts)} total={sum(proposal_counts)} "
                    f"min={min(proposal_counts) if proposal_counts else 0} "
                    f"max={max(proposal_counts) if proposal_counts else 0}",
                    flush=True,
                )
            # RPN-only ROI assignment stats for logs (matches inference-time proposals; not inflated by add_gt)
            roi_diag_list: List[Dict[str, float]] = []
            for img_idx in range(len(images)):
                if len(gt_boxes_list[img_idx]) == 0:
                    continue
                roi_diag_list.append(
                    compute_roi_matching_diagnostics(
                        xyxy_to_obb(proposals_xyxy[img_idx]),
                        gt_boxes_list[img_idx],
                        gt_labels_list[img_idx],
                        device=images_tensor.device,
                        positive_iou_threshold=self.roi_positive_iou_threshold,
                        negative_iou_threshold=self.roi_negative_iou_threshold,
                        use_hbb_for_matching=self.use_hbb_for_matching,
                        match_low_quality=self.roi_match_low_quality,
                        roi_min_pos_iou=self.roi_min_pos_iou,
                        fg_bg_sampling_ratio=self.roi_fg_bg_sampling_ratio,
                        batch_size_per_image=self.roi_batch_size_per_image,
                    )
                )
            _timing_mark("roi_matching_diagnostics")
            # Optionally add GT boxes as proposals (MMRotate-style) so RoI head always sees positive samples
            if self.add_gt_as_proposals:
                for img_idx in range(len(proposals_xyxy)):
                    gt_boxes = gt_boxes_list[img_idx]
                    if len(gt_boxes) > 0:
                        gt_xyxy = obb_to_xyxy_gpu(
                            gt_boxes.to(device=proposals_xyxy[img_idx].device, dtype=proposals_xyxy[img_idx].dtype)
                        )
                        proposals_xyxy[img_idx] = torch.cat([proposals_xyxy[img_idx], gt_xyxy], dim=0)
            _timing_mark("add_gt_proposals")

            # ROI forward
            # Track which image each proposal belongs to
            box_to_image = []
            all_proposals = []
            for img_idx, img_proposals in enumerate(proposals_xyxy):
                num_proposals = len(img_proposals)
                box_to_image.extend([img_idx] * num_proposals)
                all_proposals.append(img_proposals)

            if len(all_proposals) > 0:
                all_proposals_tensor = torch.cat(all_proposals, dim=0).detach()  # [total,4] xyxy
                box_to_image_tensor = torch.tensor(box_to_image, dtype=torch.long, device=all_proposals_tensor.device)

                # ROI align
                roi_features = self.roi_align(
                    feature_maps=feature_list,
                    boxes_xyxy=all_proposals_tensor,
                    image_sizes=image_sizes,
                    box_to_image=box_to_image_tensor,
                    fpn_strides_override=fpn_strides_live,
                    spatial_scales_override=roi_spatial_scales_live,
                )  # [total_proposals, C, roi_h, roi_w]
                _timing_mark("roi_align")

                # Flatten ROI features
                roi_features_flat = roi_features.view(roi_features.shape[0], -1)  # [total_proposals, C*roi_h*roi_w]

                # ROI head
                class_logits, box_regression = self.roi_head(roi_features_flat)
                _timing_mark("roi_head")

                # Compute ROI loss
                # Group proposals by image
                roi_losses_list = []
                start_idx = 0
                for img_idx in range(len(images)):
                    # Find proposals for this image
                    img_mask = box_to_image_tensor == img_idx
                    img_proposal_indices = img_mask.nonzero(as_tuple=True)[0]

                    if len(img_proposal_indices) > 0:
                        img_proposals = all_proposals_tensor[img_proposal_indices]
                        img_class_logits = class_logits[img_proposal_indices]
                        img_box_regression = box_regression[img_proposal_indices]
                        img_gt_boxes = gt_boxes_list[img_idx]
                        img_gt_labels = gt_labels_list[img_idx]

                        if len(img_gt_boxes) > 0:
                            roi_losses = compute_horizontal_roi_loss(
                                class_logits=img_class_logits,
                                box_regression=img_box_regression,
                                proposals_xyxy=img_proposals,
                                gt_boxes=img_gt_boxes,
                                gt_labels=img_gt_labels,
                                gt_boxes_ignore=gt_boxes_ignore_list[img_idx],
                                positive_iou_threshold=self.roi_positive_iou_threshold,
                                negative_iou_threshold=self.roi_negative_iou_threshold,
                                box_reg_weight=1.0,
                                box_reg_angle_weight=self.roi_box_reg_angle_weight,
                                main_loss_type=self.roi_box_reg_main_loss_type,
                                reg_norm=self.roi_box_reg_norm,
                                box_reg_iou_weight=self.roi_box_reg_iou_weight,
                                box_reg_iou_loss_type=self.roi_box_reg_iou_loss_type,
                                box_reg_kfiou_fun=self.roi_box_reg_kfiou_fun,
                                box_reg_probiou_mode=self.roi_box_reg_probiou_mode,
                                smooth_l1_aux_weight=self.roi_box_reg_smooth_l1_aux_weight,
                                fg_bg_sampling_ratio=self.roi_fg_bg_sampling_ratio,
                                batch_size_per_image=self.roi_batch_size_per_image,
                                num_classes=self.num_classes,
                                loss_type=self.roi_loss_type,
                                class_weights=self.roi_class_weights_tensor,
                                focal_alpha=self.roi_focal_alpha,
                                focal_gamma=self.roi_focal_gamma,
                                means=self.target_means,
                                stds=self.target_stds,
                                norm_factor=self.roi_norm_factor,
                                edge_swap=self.roi_edge_swap,
                                proj_xy=self.roi_proj_xy,
                                label_smoothing=self.roi_label_smoothing,
                                match_low_quality=self.roi_match_low_quality,
                                roi_min_pos_iou=self.roi_min_pos_iou,
                                include_assignment_diagnostics=False,
                                **self.roi_grouped_ce_kwargs(),
                            )
                            roi_losses_list.append(roi_losses)
                _timing_mark("roi_loss")

                # Aggregate ROI losses
                if roi_losses_list:
                    roi_losses = {}
                    for key in roi_losses_list[0].keys():
                        values = [l[key] for l in roi_losses_list if key in l]
                        if not values:
                            continue
                        if torch.is_tensor(values[0]):
                            roi_losses[key] = torch.stack(values).mean()
                        else:
                            roi_losses[key] = float(sum(values) / len(values))
                else:
                    # Maintain gradient flow: compute zero losses from model outputs
                    # Use ROI head outputs to maintain connection to computation graph
                    if len(all_proposals) > 0 and len(all_proposals[0]) > 0 and class_logits.numel() > 0:
                        # If we have proposals but no matches, use ROI head outputs
                        dummy_loss_cls = (class_logits[0:1] * 0.0).sum()
                        dummy_loss_reg = (box_regression[0:1] * 0.0).sum() if box_regression.numel() > 0 else dummy_loss_cls
                        roi_losses = {
                            "loss_classifier": dummy_loss_cls,
                            "loss_box_reg": dummy_loss_reg,
                        }
                    else:
                        # No proposals at all - use RPN outputs to maintain gradient flow
                        if len(objectness_logits) > 0 and objectness_logits[0].numel() > 0:
                            dummy_loss = (objectness_logits[0] * 0.0).sum()
                        else:
                            dummy_loss = torch.tensor(0.0, device=images_tensor.device, requires_grad=True)
                        roi_losses = {
                            "loss_classifier": dummy_loss,
                            "loss_box_reg": dummy_loss,
                        }
            else:
                # No proposals generated - maintain gradient flow using RPN outputs
                if len(objectness_logits) > 0 and objectness_logits[0].numel() > 0:
                    dummy_loss = (objectness_logits[0] * 0.0).sum()
                else:
                    dummy_loss = torch.tensor(0.0, device=images_tensor.device, requires_grad=True)
                roi_losses = {
                    "loss_classifier": dummy_loss,
                    "loss_box_reg": dummy_loss,
                }

            if roi_diag_list:
                diag_agg = {
                    k: float(sum(d[k] for d in roi_diag_list) / len(roi_diag_list))
                    for k in roi_diag_list[0]
                }
                roi_losses.update(diag_agg)

            # Combine losses
            return {
                **rpn_losses,
                **roi_losses,
            }

        else:
            from .faster_rcnn_inference import faster_rcnn_inference

            return faster_rcnn_inference(self, images)

__init__(num_classes, *, backbone=None, backbone_name='resnet50', pretrained_backbone=False, trainable_layers=5, anchor_scales=None, anchor_ratios=None, anchor_angles=None, rpn_pre_nms_top_n=2000, rpn_post_nms_top_n=2000, rpn_min_size=0.0, max_detections_per_image=2000, rpn_nms_threshold=0.7, final_nms_iou_threshold=0.1, rpn_positive_iou_threshold=0.7, rpn_negative_iou_threshold=0.3, rpn_fg_bg_sampling_ratio=0.5, rpn_batch_size_per_image=256, rpn_min_pos_iou=0.3, rpn_match_low_quality=True, roi_match_low_quality=False, roi_min_pos_iou=0.5, add_gt_as_proposals=True, roi_positive_iou_threshold=0.5, roi_negative_iou_threshold=0.5, roi_fg_bg_sampling_ratio=0.25, roi_batch_size_per_image=512, roi_output_size=(7, 7), roi_spatial_scales=None, fpn_strides=None, returned_layers=None, use_checkpoint=False, roi_chunk_size=32, roi_use_checkpoint=False, roi_loss_type='cross_entropy', roi_class_weights=None, roi_focal_alpha=1.0, roi_focal_gamma=2.0, roi_label_smoothing=0.0, target_means=None, target_stds=None, roi_norm_factor=2.0, roi_edge_swap=True, roi_proj_xy=False, roi_box_reg_angle_weight=1.0, roi_box_reg_angle_schedule_epochs=None, roi_box_reg_angle_schedule_values=None, roi_box_reg_iou_weight=0.0, roi_box_reg_iou_loss_type='riou', roi_box_reg_kfiou_fun=None, roi_box_reg_probiou_mode=None, roi_box_reg_main_loss_type='smooth_l1', roi_box_reg_norm='sampled_all', roi_box_reg_smooth_l1_aux_weight=0.0, use_hbb_for_matching=False, inference_pre_nms_score_threshold=0.05, final_nms_iou_schedule_epochs=None, final_nms_iou_schedule_values=None, roi_box_reg_iou_schedule_epochs=None, roi_box_reg_iou_schedule_values=None, nms_class_agnostic=False, final_nms_use_cpu=False, roi_inference_top_class_only=False)

Initialize Rotated Faster R-CNN model.

Parameters:

Name Type Description Default
num_classes int

Number of object classes (excluding background)

 required
backbone

Optional backbone module (if None, creates ResNet+FPN)

 None
backbone_name str

Name of backbone to create ("resnet18", "resnet50", etc.)

 'resnet50'
pretrained_backbone bool

Whether to use pretrained backbone weights

 False
trainable_layers int

Number of backbone layers to keep trainable. For ResNet50, typically 5 layers (conv1 + layer1-4). Set to 5 to train all layers. Default is 5 (all layers trainable).

 5
anchor_scales Optional[List[float]]

List of anchor scales for RPN

 None
anchor_ratios Optional[List[float]]

List of anchor aspect ratios for RPN

 None
anchor_angles Optional[List[float]]

Optional RPN angles in radians (Python API only; not in JSON). None -> [0.0].

 None
rpn_pre_nms_top_n int

Number of top proposals before NMS in RPN

 2000
rpn_post_nms_top_n int

Number of top proposals after NMS in RPN

 2000
rpn_positive_iou_threshold float

IoU threshold for positive RPN anchors

 0.7
rpn_negative_iou_threshold float

IoU threshold for negative RPN anchors

 0.3
rpn_fg_bg_sampling_ratio float

Ratio of foreground to background in sampled RPN anchors (default: 0.5)

 0.5
roi_positive_iou_threshold float

IoU threshold for positive ROI proposals

 0.5
roi_negative_iou_threshold float

IoU threshold for negative ROI proposals

 0.5
roi_fg_bg_sampling_ratio float

Ratio of foreground to background in sampled ROI proposals (default: 0.25)

 0.25
roi_output_size Tuple[int, int]

Output size for ROI align (height, width)

 (7, 7)
roi_spatial_scales Optional[List[float]]

Spatial scales for each FPN level in ROI align (defaults from fpn_strides at init)

 None
fpn_strides Optional[List[int]]

Nominal strides for ROI module init; forward uses strides derived from the feature grid.

 None
use_checkpoint bool

Whether to use gradient checkpointing (trades compute for memory)

 False
roi_chunk_size int

Number of ROIs to process in parallel during ROI align

 32
roi_use_checkpoint bool

If True, use gradient checkpointing in ROI align

 False
roi_loss_type str

Type of classification loss: - "cross_entropy": Standard cross-entropy loss (default) - "focal": Focal loss for handling class imbalance

 'cross_entropy'
roi_class_weights Optional[Union[Dict[str, float], Tensor]]

Optional class weights for loss computation. Can be: - Dict[str, float]: Mapping from class name to weight - torch.Tensor: Tensor of shape [num_classes + 1] (including background) - None: Equal weights for all classes (default) Note: If dict is provided, it will be converted to tensor using the model's class mapping. Background (index 0) defaults to weight 1.0 if not specified.

 None
roi_focal_alpha float

Alpha parameter for focal loss (default: 1.0)

 1.0
roi_focal_gamma float

Gamma parameter for focal loss (default: 2.0)

 2.0
roi_label_smoothing float

Label smoothing for ROI classification (default: 0.0). Use e.g. 0.1 to reduce overconfidence.

 0.0
target_means Optional[Tuple[float, float, float, float, float]]

Optional means for regression target normalization [dx, dy, dw, dh, dangle]. Default: (0.0, 0.0, 0.0, 0.0, 0.0) - no centering (like MMRotate default). Can be set to empirical means computed from dataset for better normalization.

 None
target_stds Optional[Tuple[float, float, float, float, float]]

Optional stds for regression target normalization [dx, dy, dw, dh, dangle]. Default: (0.1, 0.1, 0.2, 0.2, 0.1) (MMRotate DeltaXYWHTHBBoxCoder). Can be set to empirical stds computed from dataset for better normalization.

 None
roi_norm_factor float

Angle scaling for ROI encode/decode (MMRotate uses 2.0).

 2.0
roi_edge_swap bool

Whether to use edge_swap for ROI (MMRotate uses True).

 True
roi_proj_xy bool

If True, encode/decode ROI dx/dy in proposal local frame. For horizontal xyxy RoIs (angle 0) this is equivalent to global offsets.

 False
roi_box_reg_angle_weight float

Weight for the angle (5th) component in ROI box regression loss. Use > 1.0 (e.g. 2.0) to improve orientation alignment when predictions stick to anchor angles.

 1.0
roi_box_reg_iou_weight float

Weight for auxiliary decoded-box loss when main is Smooth L1.

 0.0
roi_box_reg_iou_loss_type str

riou (default), kfiou, or probiou for decoded aux/main.

 'riou'
roi_box_reg_kfiou_fun Optional[str]

Optional KFIoU overlap mapping when using kfiou.

 None
roi_box_reg_probiou_mode Optional[str]

l1 (default) or l2 when using probiou.

 None
roi_box_reg_main_loss_type str

Primary ROI reg loss: smooth_l1 (default) or decoded probiou / riou / kfiou.

 'smooth_l1'
roi_box_reg_norm str

sampled_all (MMRotate avg_factor) or positives_only.

 'sampled_all'
roi_box_reg_smooth_l1_aux_weight float

Encoded Smooth L1 aux when main is decoded.

 0.0
roi_min_pos_iou float

Min IoU for low-quality ROI matches when roi_match_low_quality is True.

 0.5
roi_inference_top_class_only bool

If True, ROI inference uses argmax fg class per proposal (see oriented_det/models/README.md); if False, multiclass thresholding.

 False
Source code in oriented_det/models/oriented_rcnn.py 
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
def __init__(
    self,
    num_classes: int,
    *,
    backbone=None,
    backbone_name: str = "resnet50",
    pretrained_backbone: bool = False,
    trainable_layers: int = 5,
    anchor_scales: Optional[List[float]] = None,
    anchor_ratios: Optional[List[float]] = None,
    anchor_angles: Optional[List[float]] = None,
    rpn_pre_nms_top_n: int = 2000,
    rpn_post_nms_top_n: int = 2000,
    rpn_min_size: float = 0.0,  # Minimum proposal side length (pixels); use 2–4 to filter tiny boxes
    max_detections_per_image: int = 2000,  # MMRotate test_cfg rcnn max_per_img
    rpn_nms_threshold: float = 0.7,  # IoU threshold for RPN proposal NMS
    final_nms_iou_threshold: float = 0.1,  # MMRotate rotated NMS iou (DOTA le90 config)
    rpn_positive_iou_threshold: float = 0.7,  # MMRotate RPN MaxIoUAssigner pos_iou_thr
    rpn_negative_iou_threshold: float = 0.3,  # MMRotate RPN neg_iou_thr
    rpn_fg_bg_sampling_ratio: float = 0.5,  # Ratio of foreground to background in sampled RPN anchors
    rpn_batch_size_per_image: int = 256,  # Number of RPN anchors to sample per image (MMRotate: 256)
    rpn_min_pos_iou: float = 0.3,  # Min IoU for best-anchor-per-GT to be positive (MMRotate RPN assigner)
    rpn_match_low_quality: bool = True,  # If True, force best anchor per GT to positive when IoU >= rpn_min_pos_iou
    roi_match_low_quality: bool = False,  # MMRotate RCNN assigner match_low_quality (usually False)
    roi_min_pos_iou: float = 0.5,  # MMRotate RCNN MaxIoUAssigner min_pos_iou
    add_gt_as_proposals: bool = True,  # If True, append GT boxes to RPN proposals in training (MMRotate default)
    roi_positive_iou_threshold: float = 0.5,  # MMRotate RCNN pos_iou_thr
    roi_negative_iou_threshold: float = 0.5,  # MMRotate RCNN neg_iou_thr (same as pos for DOTA config)
    roi_fg_bg_sampling_ratio: float = 0.25,  # Ratio of foreground to background in sampled ROI proposals
    roi_batch_size_per_image: int = 512,  # Number of ROI proposals to sample per image for training
    roi_output_size: Tuple[int, int] = (7, 7),
    roi_spatial_scales: Optional[List[float]] = None,
    fpn_strides: Optional[List[int]] = None,
    returned_layers: Optional[List[int]] = None,
    use_checkpoint: bool = False,
    roi_chunk_size: int = 32,
    roi_use_checkpoint: bool = False,
    # Classification loss configuration
    roi_loss_type: str = "cross_entropy",
    roi_class_weights: Optional[Union[Dict[str, float], torch.Tensor]] = None,
    roi_focal_alpha: float = 1.0,
    roi_focal_gamma: float = 2.0,
    roi_label_smoothing: float = 0.0,
    # Box regression target normalization (like MMRotate: only ROI head uses small stds)
    target_means: Optional[Tuple[float, float, float, float, float]] = None,
    target_stds: Optional[Tuple[float, float, float, float, float]] = None,
    roi_norm_factor: float = 2.0,
    roi_edge_swap: bool = True,
    roi_proj_xy: bool = False,
    roi_box_reg_angle_weight: float = 1.0,
    roi_box_reg_angle_schedule_epochs: Optional[List[int]] = None,
    roi_box_reg_angle_schedule_values: Optional[List[float]] = None,
    roi_box_reg_iou_weight: float = 0.0,
    roi_box_reg_iou_loss_type: str = "riou",
    roi_box_reg_kfiou_fun: Optional[str] = None,
    roi_box_reg_probiou_mode: Optional[str] = None,
    roi_box_reg_main_loss_type: str = "smooth_l1",
    roi_box_reg_norm: str = "sampled_all",
    roi_box_reg_smooth_l1_aux_weight: float = 0.0,
    use_hbb_for_matching: bool = False,
    inference_pre_nms_score_threshold: float = 0.05,
    final_nms_iou_schedule_epochs: Optional[List[int]] = None,
    final_nms_iou_schedule_values: Optional[List[float]] = None,
    roi_box_reg_iou_schedule_epochs: Optional[List[int]] = None,
    roi_box_reg_iou_schedule_values: Optional[List[float]] = None,
    nms_class_agnostic: bool = False,
    final_nms_use_cpu: bool = False,
    roi_inference_top_class_only: bool = False,
):
    """Initialize Rotated Faster R-CNN model.

    Args:
        num_classes: Number of object classes (excluding background)
        backbone: Optional backbone module (if None, creates ResNet+FPN)
        backbone_name: Name of backbone to create ("resnet18", "resnet50", etc.)
        pretrained_backbone: Whether to use pretrained backbone weights
        trainable_layers: Number of backbone layers to keep trainable. For ResNet50,
                         typically 5 layers (conv1 + layer1-4). Set to 5 to train all layers.
                         Default is 5 (all layers trainable).
        anchor_scales: List of anchor scales for RPN
        anchor_ratios: List of anchor aspect ratios for RPN
        anchor_angles: Optional RPN angles in radians (Python API only; not in JSON). ``None`` -> ``[0.0]``.
        rpn_pre_nms_top_n: Number of top proposals before NMS in RPN
        rpn_post_nms_top_n: Number of top proposals after NMS in RPN
        rpn_positive_iou_threshold: IoU threshold for positive RPN anchors
        rpn_negative_iou_threshold: IoU threshold for negative RPN anchors
        rpn_fg_bg_sampling_ratio: Ratio of foreground to background in sampled RPN anchors (default: 0.5)
        roi_positive_iou_threshold: IoU threshold for positive ROI proposals
        roi_negative_iou_threshold: IoU threshold for negative ROI proposals
        roi_fg_bg_sampling_ratio: Ratio of foreground to background in sampled ROI proposals (default: 0.25)
        roi_output_size: Output size for ROI align (height, width)
        roi_spatial_scales: Spatial scales for each FPN level in ROI align (defaults from ``fpn_strides`` at init)
        fpn_strides: Nominal strides for ROI module init; forward uses strides derived from the feature grid.
        use_checkpoint: Whether to use gradient checkpointing (trades compute for memory)
        roi_chunk_size: Number of ROIs to process in parallel during ROI align
        roi_use_checkpoint: If True, use gradient checkpointing in ROI align
        roi_loss_type: Type of classification loss:
                      - "cross_entropy": Standard cross-entropy loss (default)
                      - "focal": Focal loss for handling class imbalance
        roi_class_weights: Optional class weights for loss computation.
                          Can be:
                          - Dict[str, float]: Mapping from class name to weight
                          - torch.Tensor: Tensor of shape [num_classes + 1] (including background)
                          - None: Equal weights for all classes (default)
                          Note: If dict is provided, it will be converted to tensor
                                using the model's class mapping. Background (index 0)
                                defaults to weight 1.0 if not specified.
        roi_focal_alpha: Alpha parameter for focal loss (default: 1.0)
        roi_focal_gamma: Gamma parameter for focal loss (default: 2.0)
        roi_label_smoothing: Label smoothing for ROI classification (default: 0.0). Use e.g. 0.1 to reduce overconfidence.
        target_means: Optional means for regression target normalization [dx, dy, dw, dh, dangle].
                     Default: (0.0, 0.0, 0.0, 0.0, 0.0) - no centering (like MMRotate default).
                     Can be set to empirical means computed from dataset for better normalization.
        target_stds: Optional stds for regression target normalization [dx, dy, dw, dh, dangle].
                    Default: (0.1, 0.1, 0.2, 0.2, 0.1) (MMRotate DeltaXYWHTHBBoxCoder).
                    Can be set to empirical stds computed from dataset for better normalization.
        roi_norm_factor: Angle scaling for ROI encode/decode (MMRotate uses 2.0).
        roi_edge_swap: Whether to use edge_swap for ROI (MMRotate uses True).
        roi_proj_xy: If True, encode/decode ROI dx/dy in proposal local frame. For
            horizontal xyxy RoIs (angle 0) this is equivalent to global offsets.
        roi_box_reg_angle_weight: Weight for the angle (5th) component in ROI box regression loss.
            Use > 1.0 (e.g. 2.0) to improve orientation alignment when predictions stick to anchor angles.
        roi_box_reg_iou_weight: Weight for auxiliary decoded-box loss when main is Smooth L1.
        roi_box_reg_iou_loss_type: ``riou`` (default), ``kfiou``, or ``probiou`` for decoded aux/main.
        roi_box_reg_kfiou_fun: Optional KFIoU overlap mapping when using ``kfiou``.
        roi_box_reg_probiou_mode: ``l1`` (default) or ``l2`` when using ``probiou``.
        roi_box_reg_main_loss_type: Primary ROI reg loss: ``smooth_l1`` (default) or decoded
            ``probiou`` / ``riou`` / ``kfiou``.
        roi_box_reg_norm: ``sampled_all`` (MMRotate avg_factor) or ``positives_only``.
        roi_box_reg_smooth_l1_aux_weight: Encoded Smooth L1 aux when main is decoded.
        roi_min_pos_iou: Min IoU for low-quality ROI matches when ``roi_match_low_quality`` is True.
        roi_inference_top_class_only: If True, ROI inference uses argmax fg class per proposal
            (see ``oriented_det/models/README.md``); if False, multiclass thresholding.
        RPN uses fixed (0,0,0,0) and (1,1,1,1) per MMRotate DeltaXYWHBBoxCoder.
    """
    if torch is None or nn is None:
        raise RuntimeError("PyTorch is required for RotatedFasterRCNN.")

    # Initialize ClassWeightsMixin first, then nn.Module
    # ClassWeightsMixin.__init__() sets self.num_classes
    super().__init__(num_classes, roi_class_weights)
    nn.Module.__init__(self)
    self._init_grouped_ce()

    # Box regression: ROI head uses target_means/stds/norm_factor; RPN uses identity (MMRotate)
    self.target_means = target_means if target_means is not None else (0.0, 0.0, 0.0, 0.0, 0.0)
    self.target_stds = target_stds if target_stds is not None else (0.1, 0.1, 0.2, 0.2, 0.1)
    self.roi_norm_factor = roi_norm_factor
    self.roi_edge_swap = roi_edge_swap
    self.roi_proj_xy = roi_proj_xy
    self.roi_box_reg_iou_loss_type = roi_box_reg_iou_loss_type
    self.roi_box_reg_kfiou_fun = roi_box_reg_kfiou_fun
    self.roi_box_reg_probiou_mode = roi_box_reg_probiou_mode
    self.roi_box_reg_main_loss_type = roi_box_reg_main_loss_type
    self.roi_box_reg_norm = roi_box_reg_norm
    self.roi_box_reg_smooth_l1_aux_weight = float(roi_box_reg_smooth_l1_aux_weight)
    self.use_hbb_for_matching = use_hbb_for_matching
    self.inference_pre_nms_score_threshold = inference_pre_nms_score_threshold
    self._roi_box_reg_angle_schedule_epochs = roi_box_reg_angle_schedule_epochs
    self._roi_box_reg_angle_schedule_values = roi_box_reg_angle_schedule_values
    self._roi_box_reg_angle_weight_default = float(roi_box_reg_angle_weight)
    self.roi_box_reg_angle_weight = float(roi_box_reg_angle_weight)
    self.set_roi_box_reg_angle_weight_for_epoch(0)

    # Store loss configuration
    self.roi_loss_type = roi_loss_type
    self.roi_focal_alpha = roi_focal_alpha
    self.roi_focal_gamma = roi_focal_gamma
    self.roi_label_smoothing = roi_label_smoothing

    # Setup backbone using shared utility (returned_layers=[2,3,4] for MMRotate-style C3–C5 only)
    self.backbone, backbone_channels = setup_backbone(
        backbone=backbone,
        backbone_name=backbone_name,
        pretrained_backbone=pretrained_backbone,
        trainable_layers=trainable_layers,
        returned_layers=returned_layers,
    )
    # FPN strides: default [4,8,16,32,64] for C2–C5; [8,16,32,64] when returned_layers=[2,3,4]
    if fpn_strides is None and returned_layers == [2, 3, 4]:
        fpn_strides = [8, 16, 32, 64]

    # Default horizontal RPN priors; optional anchor_angles for advanced Python use only.
    self.anchor_scales, self.anchor_ratios, self.anchor_angles, self.num_anchors = setup_anchors(
        anchor_scales=anchor_scales,
        anchor_ratios=anchor_ratios,
        anchor_angles=anchor_angles,
        default_angles=[0.0],
    )

    # RPN head configuration
    # MMRotate format: cls_out_channels=1, reg_out_channels=4 (dx, dy, dw, dh - angle from anchor)
    rpn_cls_out_channels = 1
    rpn_reg_out_channels = 4

    # RPN head
    self.rpn_head = OrientedRPNHead(
        backbone_channels, 
        num_anchors=self.num_anchors,
        cls_out_channels=rpn_cls_out_channels,
        reg_out_channels=rpn_reg_out_channels,
    )

    # ROI components (MMRotate Rotated Faster R-CNN uses horizontal RoIAlign)
    self.roi_align = HorizontalROIAlign(
        output_size=roi_output_size,
        spatial_scales=roi_spatial_scales,
        fpn_strides=fpn_strides,
        chunk_size=roi_chunk_size,
    )

    # ROI head input size: output_size[0] * output_size[1] * backbone_channels
    roi_feature_size = roi_output_size[0] * roi_output_size[1] * backbone_channels
    # Use class-agnostic regression (MMRotate format) - better for initial training when class prediction is weak
    self.roi_class_agnostic_regression = True  # Store flag for inference
    self.roi_head = OrientedROIHead(
        in_channels=roi_feature_size,
        num_classes=num_classes,
        class_agnostic_regression=True,  # MMRotate format: class-agnostic regression
    )

    # Configuration
    self.rpn_pre_nms_top_n = rpn_pre_nms_top_n
    self.rpn_post_nms_top_n = rpn_post_nms_top_n
    self.rpn_min_size = rpn_min_size
    self.max_detections_per_image = max_detections_per_image
    self.rpn_nms_threshold = rpn_nms_threshold
    self.final_nms_iou_threshold = final_nms_iou_threshold
    self.nms_class_agnostic = nms_class_agnostic
    self.final_nms_use_cpu = final_nms_use_cpu
    self.roi_inference_top_class_only = roi_inference_top_class_only
    self.rpn_positive_iou_threshold = rpn_positive_iou_threshold
    self.rpn_negative_iou_threshold = rpn_negative_iou_threshold
    self.rpn_fg_bg_sampling_ratio = rpn_fg_bg_sampling_ratio
    self.rpn_batch_size_per_image = rpn_batch_size_per_image
    self.rpn_min_pos_iou = rpn_min_pos_iou
    self.rpn_match_low_quality = rpn_match_low_quality
    self.roi_match_low_quality = roi_match_low_quality
    self.roi_min_pos_iou = roi_min_pos_iou
    self.add_gt_as_proposals = add_gt_as_proposals
    self.roi_positive_iou_threshold = roi_positive_iou_threshold
    self.roi_negative_iou_threshold = roi_negative_iou_threshold
    self.roi_fg_bg_sampling_ratio = roi_fg_bg_sampling_ratio
    self.roi_batch_size_per_image = roi_batch_size_per_image
    self.fpn_strides = fpn_strides or [4, 8, 16, 32, 64]
    self.roi_spatial_scales = roi_spatial_scales or [1.0 / s for s in self.fpn_strides]

    # Gradient checkpointing (trades computation for memory)
    # When enabled, recomputes activations during backward pass instead of storing them
    # Reduces memory usage by 2-3x at the cost of ~30% slower backward pass
    self.use_checkpoint = use_checkpoint

    # NMS IoU schedule for final detection NMS (optional; essential for RetinaNet, optional for R-CNN)
    self._final_nms_iou_schedule_epochs = final_nms_iou_schedule_epochs
    self._final_nms_iou_schedule_values = final_nms_iou_schedule_values
    self._roi_box_reg_iou_schedule_epochs = roi_box_reg_iou_schedule_epochs
    self._roi_box_reg_iou_schedule_values = roi_box_reg_iou_schedule_values
    self._roi_box_reg_iou_weight_default = float(roi_box_reg_iou_weight)
    self.roi_box_reg_iou_weight = float(roi_box_reg_iou_weight)
    self.set_roi_box_reg_iou_weight_for_epoch(0)

forward(images, targets=None)

Forward pass through oriented R-CNN.

Parameters:

Name Type Description Default
images Sequence[Tensor]

List of image tensors (C, H, W) in [0, 1] range

 required
targets Optional[Sequence[Dict[str, Any]]]

Optional list of target dicts for training. Each dict must contain: - "rboxes" (List[RBox] or tensor [N, 5] with format [cx, cy, w, h, angle]) - "labels" (tensor [N])

 None

Returns:

Type Description
Union[Dict[str, Tensor], List[Dict[str, Any]]]
  • Training: Dict with loss keys:
  • "loss_objectness": RPN classification loss
  • "loss_rpn_box_reg": RPN box regression loss
  • "loss_classifier": ROI classification loss
  • "loss_box_reg": ROI box regression loss
Union[Dict[str, Tensor], List[Dict[str, Any]]]
  • Inference: List of dicts, one per image, containing:
  • "rboxes": List[RBox] with oriented boxes
  • "rboxes": List[RBox] or Tensor [N, 5] with oriented boxes [cx, cy, w, h, angle]
  • "labels": Tensor [N] with class labels
  • "scores": Tensor [N] with confidence scores
Source code in oriented_det/models/oriented_rcnn.py 
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
def forward(
    self,
    images: Sequence[torch.Tensor],
    targets: Optional[Sequence[Dict[str, Any]]] = None,
) -> Union[Dict[str, torch.Tensor], List[Dict[str, Any]]]:
    """Forward pass through oriented R-CNN.

    Args:
        images: List of image tensors (C, H, W) in [0, 1] range
        targets: Optional list of target dicts for training.
                Each dict must contain:
                - "rboxes" (List[RBox] or tensor [N, 5] with format [cx, cy, w, h, angle])
                - "labels" (tensor [N])

    Returns:
        - Training: Dict with loss keys:
            - "loss_objectness": RPN classification loss
            - "loss_rpn_box_reg": RPN box regression loss
            - "loss_classifier": ROI classification loss
            - "loss_box_reg": ROI box regression loss
        - Inference: List of dicts, one per image, containing:
            - "rboxes": List[RBox] with oriented boxes
            - "rboxes": List[RBox] or Tensor [N, 5] with oriented boxes [cx, cy, w, h, angle]
            - "labels": Tensor [N] with class labels
            - "scores": Tensor [N] with confidence scores
    """
    if not isinstance(images, (list, tuple)):
        images = [images]

    # Get image sizes
    image_sizes = [(img.shape[-2], img.shape[-1]) for img in images]
    timing_enabled = bool(
        TRACE_FIRST_TRAIN_FORWARD_TIMING
        and self.training
        and not getattr(self, "_timed_first_train_forward", False)
    )
    if timing_enabled:
        self._timed_first_train_forward = True
        timing_start = time.perf_counter()
        timing_last = timing_start

        def _timing_mark(label: str) -> None:
            nonlocal timing_last
            if torch is not None and torch.cuda.is_available():
                try:
                    torch.cuda.synchronize()
                except Exception:
                    pass
            now = time.perf_counter()
            print(f"[first-batch timing] {label}: +{now - timing_last:.3f}s total={now - timing_start:.3f}s", flush=True)
            timing_last = now

        print("[first-batch timing] begin RotatedFasterRCNN.forward", flush=True)
    else:
        def _timing_mark(label: str) -> None:
            return None

    # Extract features using shared utility. include_pool_level keeps torchvision's
    # P6 (stride-64 max-pool level) so the RPN sees 5 levels like MMRotate;
    # horizontal_roi_align still restricts ROI extraction to the first 4 levels.
    feature_list = extract_backbone_features(
        self.backbone,
        images,
        use_checkpoint=self.use_checkpoint,
        training=self.training,
        include_pool_level=True,
    )
    feature_map_sizes = [(f.shape[2], f.shape[3]) for f in feature_list]
    fpn_strides_live = derive_fpn_strides_from_grid(image_sizes[0], feature_map_sizes)
    warn_if_fpn_strides_mismatch(self.fpn_strides, fpn_strides_live)
    roi_spatial_scales_live = [1.0 / s for s in fpn_strides_live]
    _timing_mark("backbone_fpn")

    # Get device from first feature map
    images_tensor = torch.stack(images, dim=0)  # [B, C, H, W] - needed for device reference

    if self.training:
        if targets is None:
            raise ValueError("Targets required during training.")

        # Prepare targets (preserve oriented boxes)
        gt_boxes_list, gt_labels_list, gt_boxes_ignore_list = prepare_targets(targets, device=images_tensor.device)
        _timing_mark("prepare_targets")

        # RPN forward
        objectness_logits, bbox_regression = self.rpn_head(feature_list)
        _timing_mark("rpn_head")

        # Generate anchors (strides from actual grid — see fpn_strides_live)
        img_h, img_w = image_sizes[0]
        anchors = generate_oriented_anchors(
            image_size=(img_h, img_w),
            feature_map_sizes=feature_map_sizes,
            anchor_scales=self.anchor_scales,
            anchor_ratios=self.anchor_ratios,
            anchor_angles=self.anchor_angles,
            stride_per_level=fpn_strides_live,
        )
        _timing_mark("generate_anchors")

        # Compute RPN loss (4 params, no angle - MMRotate style)
        rpn_losses = compute_oriented_rpn_loss(
            objectness_logits=objectness_logits,
            bbox_regression=bbox_regression,
            anchors=anchors,
            gt_boxes=gt_boxes_list,
            gt_boxes_ignore=gt_boxes_ignore_list,
            image_sizes=image_sizes,
            positive_iou_threshold=self.rpn_positive_iou_threshold,
            negative_iou_threshold=self.rpn_negative_iou_threshold,
            box_reg_weight=1.0,
            fg_bg_sampling_ratio=self.rpn_fg_bg_sampling_ratio,
            batch_size_per_image=self.rpn_batch_size_per_image,
            min_pos_iou=self.rpn_min_pos_iou,
            match_low_quality=self.rpn_match_low_quality,
            target_means=(0.0, 0.0, 0.0, 0.0),
            target_stds=(1.0, 1.0, 1.0, 1.0),
            use_horizontal_targets=True,
            use_hbb_for_matching=self.use_hbb_for_matching,
        )
        _timing_mark("rpn_loss")

        # Horizontal proposals (xyxy) — MMRotate semantics. No objectness score
        # threshold: MMRotate's RPN keeps top-k proposals by score only;
        # inference_pre_nms_score_threshold applies to ROI-head scores, not the RPN.
        proposals_xyxy = generate_horizontal_proposals(
            objectness_logits=objectness_logits,
            bbox_regression=bbox_regression,
            anchors=anchors,
            image_sizes=image_sizes,
            score_threshold=0.0,
            nms_threshold=self.rpn_nms_threshold,
            pre_nms_top_n=self.rpn_pre_nms_top_n,
            post_nms_top_n=self.rpn_post_nms_top_n,
            min_size=self.rpn_min_size,
            target_means=(0.0, 0.0, 0.0, 0.0),
            target_stds=(1.0, 1.0, 1.0, 1.0),
        )
        _timing_mark("generate_proposals")
        if timing_enabled:
            proposal_counts = [int(len(p)) for p in proposals_xyxy]
            print(
                f"[first-batch timing] proposals_after_cap: "
                f"images={len(proposal_counts)} total={sum(proposal_counts)} "
                f"min={min(proposal_counts) if proposal_counts else 0} "
                f"max={max(proposal_counts) if proposal_counts else 0}",
                flush=True,
            )
        # RPN-only ROI assignment stats for logs (matches inference-time proposals; not inflated by add_gt)
        roi_diag_list: List[Dict[str, float]] = []
        for img_idx in range(len(images)):
            if len(gt_boxes_list[img_idx]) == 0:
                continue
            roi_diag_list.append(
                compute_roi_matching_diagnostics(
                    xyxy_to_obb(proposals_xyxy[img_idx]),
                    gt_boxes_list[img_idx],
                    gt_labels_list[img_idx],
                    device=images_tensor.device,
                    positive_iou_threshold=self.roi_positive_iou_threshold,
                    negative_iou_threshold=self.roi_negative_iou_threshold,
                    use_hbb_for_matching=self.use_hbb_for_matching,
                    match_low_quality=self.roi_match_low_quality,
                    roi_min_pos_iou=self.roi_min_pos_iou,
                    fg_bg_sampling_ratio=self.roi_fg_bg_sampling_ratio,
                    batch_size_per_image=self.roi_batch_size_per_image,
                )
            )
        _timing_mark("roi_matching_diagnostics")
        # Optionally add GT boxes as proposals (MMRotate-style) so RoI head always sees positive samples
        if self.add_gt_as_proposals:
            for img_idx in range(len(proposals_xyxy)):
                gt_boxes = gt_boxes_list[img_idx]
                if len(gt_boxes) > 0:
                    gt_xyxy = obb_to_xyxy_gpu(
                        gt_boxes.to(device=proposals_xyxy[img_idx].device, dtype=proposals_xyxy[img_idx].dtype)
                    )
                    proposals_xyxy[img_idx] = torch.cat([proposals_xyxy[img_idx], gt_xyxy], dim=0)
        _timing_mark("add_gt_proposals")

        # ROI forward
        # Track which image each proposal belongs to
        box_to_image = []
        all_proposals = []
        for img_idx, img_proposals in enumerate(proposals_xyxy):
            num_proposals = len(img_proposals)
            box_to_image.extend([img_idx] * num_proposals)
            all_proposals.append(img_proposals)

        if len(all_proposals) > 0:
            all_proposals_tensor = torch.cat(all_proposals, dim=0).detach()  # [total,4] xyxy
            box_to_image_tensor = torch.tensor(box_to_image, dtype=torch.long, device=all_proposals_tensor.device)

            # ROI align
            roi_features = self.roi_align(
                feature_maps=feature_list,
                boxes_xyxy=all_proposals_tensor,
                image_sizes=image_sizes,
                box_to_image=box_to_image_tensor,
                fpn_strides_override=fpn_strides_live,
                spatial_scales_override=roi_spatial_scales_live,
            )  # [total_proposals, C, roi_h, roi_w]
            _timing_mark("roi_align")

            # Flatten ROI features
            roi_features_flat = roi_features.view(roi_features.shape[0], -1)  # [total_proposals, C*roi_h*roi_w]

            # ROI head
            class_logits, box_regression = self.roi_head(roi_features_flat)
            _timing_mark("roi_head")

            # Compute ROI loss
            # Group proposals by image
            roi_losses_list = []
            start_idx = 0
            for img_idx in range(len(images)):
                # Find proposals for this image
                img_mask = box_to_image_tensor == img_idx
                img_proposal_indices = img_mask.nonzero(as_tuple=True)[0]

                if len(img_proposal_indices) > 0:
                    img_proposals = all_proposals_tensor[img_proposal_indices]
                    img_class_logits = class_logits[img_proposal_indices]
                    img_box_regression = box_regression[img_proposal_indices]
                    img_gt_boxes = gt_boxes_list[img_idx]
                    img_gt_labels = gt_labels_list[img_idx]

                    if len(img_gt_boxes) > 0:
                        roi_losses = compute_horizontal_roi_loss(
                            class_logits=img_class_logits,
                            box_regression=img_box_regression,
                            proposals_xyxy=img_proposals,
                            gt_boxes=img_gt_boxes,
                            gt_labels=img_gt_labels,
                            gt_boxes_ignore=gt_boxes_ignore_list[img_idx],
                            positive_iou_threshold=self.roi_positive_iou_threshold,
                            negative_iou_threshold=self.roi_negative_iou_threshold,
                            box_reg_weight=1.0,
                            box_reg_angle_weight=self.roi_box_reg_angle_weight,
                            main_loss_type=self.roi_box_reg_main_loss_type,
                            reg_norm=self.roi_box_reg_norm,
                            box_reg_iou_weight=self.roi_box_reg_iou_weight,
                            box_reg_iou_loss_type=self.roi_box_reg_iou_loss_type,
                            box_reg_kfiou_fun=self.roi_box_reg_kfiou_fun,
                            box_reg_probiou_mode=self.roi_box_reg_probiou_mode,
                            smooth_l1_aux_weight=self.roi_box_reg_smooth_l1_aux_weight,
                            fg_bg_sampling_ratio=self.roi_fg_bg_sampling_ratio,
                            batch_size_per_image=self.roi_batch_size_per_image,
                            num_classes=self.num_classes,
                            loss_type=self.roi_loss_type,
                            class_weights=self.roi_class_weights_tensor,
                            focal_alpha=self.roi_focal_alpha,
                            focal_gamma=self.roi_focal_gamma,
                            means=self.target_means,
                            stds=self.target_stds,
                            norm_factor=self.roi_norm_factor,
                            edge_swap=self.roi_edge_swap,
                            proj_xy=self.roi_proj_xy,
                            label_smoothing=self.roi_label_smoothing,
                            match_low_quality=self.roi_match_low_quality,
                            roi_min_pos_iou=self.roi_min_pos_iou,
                            include_assignment_diagnostics=False,
                            **self.roi_grouped_ce_kwargs(),
                        )
                        roi_losses_list.append(roi_losses)
            _timing_mark("roi_loss")

            # Aggregate ROI losses
            if roi_losses_list:
                roi_losses = {}
                for key in roi_losses_list[0].keys():
                    values = [l[key] for l in roi_losses_list if key in l]
                    if not values:
                        continue
                    if torch.is_tensor(values[0]):
                        roi_losses[key] = torch.stack(values).mean()
                    else:
                        roi_losses[key] = float(sum(values) / len(values))
            else:
                # Maintain gradient flow: compute zero losses from model outputs
                # Use ROI head outputs to maintain connection to computation graph
                if len(all_proposals) > 0 and len(all_proposals[0]) > 0 and class_logits.numel() > 0:
                    # If we have proposals but no matches, use ROI head outputs
                    dummy_loss_cls = (class_logits[0:1] * 0.0).sum()
                    dummy_loss_reg = (box_regression[0:1] * 0.0).sum() if box_regression.numel() > 0 else dummy_loss_cls
                    roi_losses = {
                        "loss_classifier": dummy_loss_cls,
                        "loss_box_reg": dummy_loss_reg,
                    }
                else:
                    # No proposals at all - use RPN outputs to maintain gradient flow
                    if len(objectness_logits) > 0 and objectness_logits[0].numel() > 0:
                        dummy_loss = (objectness_logits[0] * 0.0).sum()
                    else:
                        dummy_loss = torch.tensor(0.0, device=images_tensor.device, requires_grad=True)
                    roi_losses = {
                        "loss_classifier": dummy_loss,
                        "loss_box_reg": dummy_loss,
                    }
        else:
            # No proposals generated - maintain gradient flow using RPN outputs
            if len(objectness_logits) > 0 and objectness_logits[0].numel() > 0:
                dummy_loss = (objectness_logits[0] * 0.0).sum()
            else:
                dummy_loss = torch.tensor(0.0, device=images_tensor.device, requires_grad=True)
            roi_losses = {
                "loss_classifier": dummy_loss,
                "loss_box_reg": dummy_loss,
            }

        if roi_diag_list:
            diag_agg = {
                k: float(sum(d[k] for d in roi_diag_list) / len(roi_diag_list))
                for k in roi_diag_list[0]
            }
            roi_losses.update(diag_agg)

        # Combine losses
        return {
            **rpn_losses,
            **roi_losses,
        }

    else:
        from .faster_rcnn_inference import faster_rcnn_inference

        return faster_rcnn_inference(self, images)

set_final_nms_iou_for_epoch(epoch)

Update final detection NMS IoU threshold from schedule. Lower = more aggressive suppression.

Source code in oriented_det/models/oriented_rcnn.py 
399
400
401
402
403
404
405
406
407
408
409
410
411
def set_final_nms_iou_for_epoch(self, epoch: int) -> None:
    """Update final detection NMS IoU threshold from schedule. Lower = more aggressive suppression."""
    if self._final_nms_iou_schedule_epochs is None or self._final_nms_iou_schedule_values is None:
        return
    if not self._final_nms_iou_schedule_epochs or not self._final_nms_iou_schedule_values:
        return
    idx = 0
    for boundary in self._final_nms_iou_schedule_epochs:
        if epoch < boundary:
            break
        idx += 1
    idx = min(idx, len(self._final_nms_iou_schedule_values) - 1)
    self.final_nms_iou_threshold = self._final_nms_iou_schedule_values[idx]

set_roi_box_reg_angle_weight_for_epoch(epoch)

Update ROI angle (5th dim) SmoothL1 weight from schedule (0-based epoch index).

Source code in oriented_det/models/oriented_rcnn.py 
424
425
426
427
428
429
430
431
432
433
def set_roi_box_reg_angle_weight_for_epoch(self, epoch: int) -> None:
    """Update ROI angle (5th dim) SmoothL1 weight from schedule (0-based epoch index)."""
    from oriented_det.train.piecewise_schedule import resolve_piecewise_schedule

    self.roi_box_reg_angle_weight = resolve_piecewise_schedule(
        epoch,
        self._roi_box_reg_angle_schedule_epochs,
        self._roi_box_reg_angle_schedule_values,
        self._roi_box_reg_angle_weight_default,
    )

set_roi_box_reg_iou_weight_for_epoch(epoch)

Update ROI auxiliary IoU loss weight from schedule (0-based epoch index).

Source code in oriented_det/models/oriented_rcnn.py 
413
414
415
416
417
418
419
420
421
422
def set_roi_box_reg_iou_weight_for_epoch(self, epoch: int) -> None:
    """Update ROI auxiliary IoU loss weight from schedule (0-based epoch index)."""
    from oriented_det.train.piecewise_schedule import resolve_piecewise_schedule

    self.roi_box_reg_iou_weight = resolve_piecewise_schedule(
        epoch,
        self._roi_box_reg_iou_schedule_epochs,
        self._roi_box_reg_iou_schedule_values,
        self._roi_box_reg_iou_weight_default,
    )

RotatedRetinaNet

Bases: Module

Complete Rotated RetinaNet model for true oriented object detection.

This model implements a full single-stage oriented detector that: - Uses MMRotate-style anchor priors: horizontal boxes as (cx, cy, w, h, theta) with fixed theta=0 at init - Predicts oriented bounding boxes with 5 parameters (cx, cy, w, h, angle) - Uses oriented IoU for matching and oriented NMS for post-processing - Preserves angle information throughout training and inference - Uses sigmoid focal loss for classification (MMRotate FocalLoss use_sigmoid=True)

Parameters:

Name Type Description Default
num_classes int

Number of object classes (excluding background)

 required
backbone

Optional backbone module (if None, creates ResNet+FPN)

 None
backbone_name str

Name of backbone to create ("resnet18", "resnet50", etc.)

 'resnet50'
pretrained_backbone bool

Whether to use pretrained backbone weights

 False
trainable_layers int

Number of backbone layers to keep trainable

 5
anchor_scales Optional[List[float]]

List of anchor scales for Rotated RetinaNet

 None
anchor_ratios Optional[List[float]]

List of anchor aspect ratios for Rotated RetinaNet

 None
anchor_angles Optional[List[float]]

Optional RPN anchor angles in radians (advanced, not in ModelConfig / JSON). None uses horizontal priors [0.0] (recommended). Non-default multi-angle banks can hurt accuracy.

 None
positive_iou_threshold float

IoU threshold for positive anchors

 0.5
negative_iou_threshold float

IoU threshold for negative anchors

 0.4
focal_alpha float

Alpha parameter for focal loss (default: 0.25, Rotated RetinaNet standard)

 0.25
focal_gamma float

Gamma parameter for focal loss (default: 2.0, Rotated RetinaNet standard)

 2.0
box_reg_weight float

Weight for box regression loss (default: 1.0)

 1.0
score_threshold float

Score threshold for inference (default: 0.05)

 0.05
final_nms_iou_threshold float

IoU threshold for final oriented NMS (default: 0.5); not RPN NMS

 0.5
max_detections_per_image int

Maximum number of detections per image (default: 100)

 100
final_nms_iou_schedule_epochs Optional[List[int]]

Optional epoch boundaries for final NMS IoU schedule (e.g. [50, 150, 250])

 None
final_nms_iou_schedule_values Optional[List[float]]

Final NMS IoU per segment (e.g. [0.6, 0.45, 0.35, 0.25]); lower = more suppression

 None
target_means Optional[Tuple[float, float, float, float, float]]

Optional means for target normalization (MMRotate compatibility)

 None
target_stds Optional[Tuple[float, float, float, float, float]]

Optional stds for target normalization (MMRotate compatibility)

 None
norm_factor Optional[float]

Optional angle scaling for encode/decode (MMRotate Rotated RetinaNet uses None)

 None
edge_swap bool

Whether to use edge_swap in bbox coder (MMRotate uses True)

 True
use_hbb_for_matching bool

If True, use HBB IoU for anchor-GT matching (optional; priors use a single reference angle).

 False
final_nms_use_cpu bool

If True, skip GPU sampling NMS and run final class-aware NMS with exact polygon IoU on CPU.

 False
Example

model = RotatedRetinaNet(num_classes=15, backbone_name="resnet50") model.train() losses = model(images, targets) model.eval() predictions = model(images)

Source code in oriented_det/models/rotated_retinanet.py 
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
class RotatedRetinaNet(nn.Module):
    """Complete Rotated RetinaNet model for true oriented object detection.

    This model implements a full single-stage oriented detector that:
    - Uses MMRotate-style anchor priors: horizontal boxes as (cx, cy, w, h, theta) with fixed theta=0 at init
    - Predicts oriented bounding boxes with 5 parameters (cx, cy, w, h, angle)
    - Uses oriented IoU for matching and oriented NMS for post-processing
    - Preserves angle information throughout training and inference
    - Uses sigmoid focal loss for classification (MMRotate FocalLoss use_sigmoid=True)

    Args:
        num_classes: Number of object classes (excluding background)
        backbone: Optional backbone module (if None, creates ResNet+FPN)
        backbone_name: Name of backbone to create ("resnet18", "resnet50", etc.)
        pretrained_backbone: Whether to use pretrained backbone weights
        trainable_layers: Number of backbone layers to keep trainable
        anchor_scales: List of anchor scales for Rotated RetinaNet
        anchor_ratios: List of anchor aspect ratios for Rotated RetinaNet
        anchor_angles: Optional RPN anchor angles in radians (advanced, **not** in `ModelConfig` / JSON).
            ``None`` uses horizontal priors ``[0.0]`` (recommended). Non-default multi-angle banks can hurt accuracy.
        positive_iou_threshold: IoU threshold for positive anchors
        negative_iou_threshold: IoU threshold for negative anchors
        focal_alpha: Alpha parameter for focal loss (default: 0.25, Rotated RetinaNet standard)
        focal_gamma: Gamma parameter for focal loss (default: 2.0, Rotated RetinaNet standard)
        box_reg_weight: Weight for box regression loss (default: 1.0)
        score_threshold: Score threshold for inference (default: 0.05)
        final_nms_iou_threshold: IoU threshold for final oriented NMS (default: 0.5); not RPN NMS
        max_detections_per_image: Maximum number of detections per image (default: 100)
        final_nms_iou_schedule_epochs: Optional epoch boundaries for final NMS IoU schedule (e.g. [50, 150, 250])
        final_nms_iou_schedule_values: Final NMS IoU per segment (e.g. [0.6, 0.45, 0.35, 0.25]); lower = more suppression
        target_means: Optional means for target normalization (MMRotate compatibility)
        target_stds: Optional stds for target normalization (MMRotate compatibility)
        norm_factor: Optional angle scaling for encode/decode (MMRotate Rotated RetinaNet uses None)
        edge_swap: Whether to use edge_swap in bbox coder (MMRotate uses True)
        use_hbb_for_matching: If True, use HBB IoU for anchor-GT matching (optional; priors use a single reference angle).
        final_nms_use_cpu: If True, skip GPU sampling NMS and run final class-aware NMS with exact polygon IoU on CPU.

    Example:
        >>> model = RotatedRetinaNet(num_classes=15, backbone_name="resnet50")
        >>> model.train()
        >>> losses = model(images, targets)
        >>> model.eval()
        >>> predictions = model(images)
    """

    def __init__(
        self,
        num_classes: int,
        *,
        backbone=None,
        backbone_name: str = "resnet50",
        pretrained_backbone: bool = False,
        trainable_layers: int = 5,
        anchor_scales: Optional[List[float]] = None,
        anchor_ratios: Optional[List[float]] = None,
        anchor_angles: Optional[List[float]] = None,
        positive_iou_threshold: float = 0.5,
        negative_iou_threshold: float = 0.4,
        focal_alpha: float = 0.25,
        focal_gamma: float = 2.0,
        box_reg_weight: float = 1.0,
        score_threshold: float = 0.05,
        final_nms_iou_threshold: float = 0.5,
        max_detections_per_image: int = 100,
        final_nms_iou_schedule_epochs: Optional[List[int]] = None,
        final_nms_iou_schedule_values: Optional[List[float]] = None,
        roi_box_reg_iou_schedule_epochs: Optional[List[int]] = None,
        roi_box_reg_iou_schedule_values: Optional[List[float]] = None,
        target_means: Optional[Tuple[float, float, float, float, float]] = None,
        target_stds: Optional[Tuple[float, float, float, float, float]] = None,
        norm_factor: Optional[float] = None,
        edge_swap: bool = True,
        box_reg_iou_weight: float = 0.0,
        box_reg_iou_loss_type: str = "riou",
        box_reg_kfiou_fun: Optional[str] = None,
        box_reg_probiou_mode: Optional[str] = None,
        use_hbb_for_matching: bool = False,
        final_nms_use_cpu: bool = False,
        returned_layers: Optional[List[int]] = None,
        fpn_strides: Optional[List[int]] = None,
        fpn_extra_level: bool = False,
        octave_base_scale: Optional[float] = None,
        scales_per_octave: Optional[int] = None,
        stacked_convs: int = 1,
        box_reg_loss_type: str = "smooth_l1",
    ) -> None:
        from .backbones.utils import require_torch
        require_torch()
        super().__init__()

        self.num_classes = num_classes

        # Bbox coder options (MMRotate: norm_factor=None, edge_swap=True for Rotated RetinaNet)
        self.norm_factor = norm_factor
        self.edge_swap = edge_swap
        self.box_reg_iou_loss_type = box_reg_iou_loss_type
        self.box_reg_kfiou_fun = box_reg_kfiou_fun
        self.box_reg_probiou_mode = box_reg_probiou_mode
        self.use_hbb_for_matching = use_hbb_for_matching
        self.final_nms_use_cpu = final_nms_use_cpu
        self.octave_base_scale = octave_base_scale
        self.scales_per_octave = scales_per_octave
        self.box_reg_loss_type = box_reg_loss_type

        # Setup backbone (P6/P7 convs on C5 when fpn_extra_level, MMRotate on_input)
        self.backbone, backbone_channels = setup_backbone(
            backbone=backbone,
            backbone_name=backbone_name,
            pretrained_backbone=pretrained_backbone,
            trainable_layers=trainable_layers,
            returned_layers=returned_layers,
            use_p6p7_extra_levels=fpn_extra_level,
        )

        # Default: horizontal priors (theta=0), MMRotate-style. Optional anchor_angles is Python-only (not in JSON).
        self.anchor_scales, self.anchor_ratios, self.anchor_angles, self.num_anchors = setup_anchors(
            anchor_scales=anchor_scales,
            anchor_ratios=anchor_ratios,
            anchor_angles=anchor_angles,
            default_angles=[0.0],
            octave_base_scale=octave_base_scale,
            scales_per_octave=scales_per_octave,
        )

        # FPN strides (nominal; forward uses grid-derived strides)
        if fpn_strides is not None:
            self.fpn_strides = list(fpn_strides)
        elif returned_layers == [2, 3, 4]:
            self.fpn_strides = [8, 16, 32, 64, 128] if fpn_extra_level else [8, 16, 32, 64]
        else:
            self.fpn_strides = [4, 8, 16, 32, 64]

        self.fpn_extra_level = fpn_extra_level

        # Create Rotated RetinaNet head
        self.head = OrientedRetinaNetHead(
            in_channels=backbone_channels,
            num_classes=num_classes,
            num_anchors=self.num_anchors,
            stacked_convs=stacked_convs,
        )

        # Loss and inference parameters
        self.positive_iou_threshold = positive_iou_threshold
        self.negative_iou_threshold = negative_iou_threshold
        self.focal_alpha = focal_alpha
        self.focal_gamma = focal_gamma
        self.box_reg_weight = box_reg_weight
        self.score_threshold = score_threshold
        self.final_nms_iou_threshold = final_nms_iou_threshold
        self.max_detections_per_image = max_detections_per_image
        self._final_nms_iou_schedule_epochs = final_nms_iou_schedule_epochs
        self._final_nms_iou_schedule_values = final_nms_iou_schedule_values
        self._roi_box_reg_iou_schedule_epochs = roi_box_reg_iou_schedule_epochs
        self._roi_box_reg_iou_schedule_values = roi_box_reg_iou_schedule_values
        self._box_reg_iou_weight_default = float(box_reg_iou_weight)
        self.box_reg_iou_weight = float(box_reg_iou_weight)
        self.set_box_reg_iou_weight_for_epoch(0)

        # Target normalization (MMRotate compatibility)
        self.target_means = target_means
        self.target_stds = target_stds

    def set_final_nms_iou_for_epoch(self, epoch: int) -> None:
        """Update final detection NMS IoU threshold from schedule for the given epoch.
        Lower threshold = more aggressive suppression of overlapping boxes.
        """
        if self._final_nms_iou_schedule_epochs is None or self._final_nms_iou_schedule_values is None:
            return
        if not self._final_nms_iou_schedule_epochs or not self._final_nms_iou_schedule_values:
            return
        idx = 0
        for boundary in self._final_nms_iou_schedule_epochs:
            if epoch < boundary:
                break
            idx += 1
        idx = min(idx, len(self._final_nms_iou_schedule_values) - 1)
        self.final_nms_iou_threshold = self._final_nms_iou_schedule_values[idx]

    def set_box_reg_iou_weight_for_epoch(self, epoch: int) -> None:
        """Update anchor auxiliary IoU loss weight from schedule (0-based epoch index)."""
        from oriented_det.train.piecewise_schedule import resolve_piecewise_schedule

        self.box_reg_iou_weight = resolve_piecewise_schedule(
            epoch,
            self._roi_box_reg_iou_schedule_epochs,
            self._roi_box_reg_iou_schedule_values,
            self._box_reg_iou_weight_default,
        )

    def forward(
        self,
        images: Sequence[torch.Tensor],
        targets: Optional[Sequence[Dict[str, Any]]] = None,
    ) -> Union[Dict[str, torch.Tensor], List[Dict[str, Any]]]:
        """Forward pass through Rotated RetinaNet.

        Args:
            images: List of image tensors (C, H, W) in [0, 1] range
            targets: Optional list of target dicts for training.
                    Each dict must contain:
                    - "rboxes" (List[RBox] or tensor [N, 5] with format [cx, cy, w, h, angle])
                    - "labels" (tensor [N], 1-indexed)

        Returns:
            - Training: Dict with loss keys:
                - "loss_classifier": Classification loss (focal loss)
                - "loss_box_reg": Box regression loss
            - Inference: List of dicts, one per image, containing:
                - "rboxes": List[RBox] with oriented boxes (with predicted angles)
                - "labels": Tensor [N] with class labels (1-indexed)
                - "scores": Tensor [N] with confidence scores
        """
        if not isinstance(images, (list, tuple)):
            images = [images]

        # Get image sizes
        image_sizes = [(img.shape[-2], img.shape[-1]) for img in images]

        # Extract features using shared utility
        feature_list = extract_backbone_features(
            self.backbone,
            images,
            use_checkpoint=False,  # Rotated RetinaNet doesn't use checkpointing currently
            training=self.training,
            include_pool_level=not self.fpn_extra_level,
        )
        feature_map_sizes = [(f.shape[2], f.shape[3]) for f in feature_list]
        fpn_strides_live = derive_fpn_strides_from_grid(image_sizes[0], feature_map_sizes)
        warn_if_fpn_strides_mismatch(self.fpn_strides, fpn_strides_live)

        # Get device from first feature map (needed for device reference)
        images_tensor = torch.stack(images, dim=0)

        # Forward through Rotated RetinaNet head
        classification_logits, bbox_regression = self.head(feature_list)

        if self.training:
            if targets is None:
                raise ValueError("Targets required during training.")

            # Prepare targets using shared utility
            gt_boxes_list, gt_labels_list, gt_boxes_ignore_list = prepare_targets(targets, device=images_tensor.device)

            # Generate anchors for all FPN levels
            img_h, img_w = image_sizes[0]
            anchors = generate_oriented_anchors(
                image_size=(img_h, img_w),
                feature_map_sizes=feature_map_sizes,
                anchor_scales=self.anchor_scales,
                anchor_ratios=self.anchor_ratios,
                anchor_angles=self.anchor_angles,
                stride_per_level=fpn_strides_live,
                octave_base_scale=self.octave_base_scale,
                scales_per_octave=self.scales_per_octave,
            )

            # Compute losses
            losses = compute_oriented_retinanet_loss(
                classification_logits=classification_logits,
                bbox_regression=bbox_regression,
                anchors=anchors,
                gt_boxes=gt_boxes_list,
                gt_labels=gt_labels_list,
                image_sizes=image_sizes,
                gt_boxes_ignore=gt_boxes_ignore_list,
                num_classes=self.num_classes,
                positive_iou_threshold=self.positive_iou_threshold,
                negative_iou_threshold=self.negative_iou_threshold,
                focal_alpha=self.focal_alpha,
                focal_gamma=self.focal_gamma,
                box_reg_weight=self.box_reg_weight,
                target_means=self.target_means,
                target_stds=self.target_stds,
                target_norm_factor=self.norm_factor,
                norm_factor=self.norm_factor,
                edge_swap=self.edge_swap,
                box_reg_iou_weight=self.box_reg_iou_weight,
                box_reg_iou_loss_type=self.box_reg_iou_loss_type,
                box_reg_kfiou_fun=self.box_reg_kfiou_fun,
                box_reg_probiou_mode=self.box_reg_probiou_mode,
                use_hbb_for_matching=self.use_hbb_for_matching,
                box_reg_loss_type=self.box_reg_loss_type,
            )

            return losses

        else:
            # Inference
            # Generate anchors
            img_h, img_w = image_sizes[0]
            anchors = generate_oriented_anchors(
                image_size=(img_h, img_w),
                feature_map_sizes=feature_map_sizes,
                anchor_scales=self.anchor_scales,
                anchor_ratios=self.anchor_ratios,
                anchor_angles=self.anchor_angles,
                stride_per_level=fpn_strides_live,
                octave_base_scale=self.octave_base_scale,
                scales_per_octave=self.scales_per_octave,
            )
            # Optional debug: expose anchors and decoded boxes (pre-NMS) for TensorBoard
            return_debug = getattr(self, '_return_anchors_proposals', False)
            if return_debug:
                all_anchors_cat = torch.cat(anchors, dim=0).detach().cpu()
            # Process each image
            outputs = []
            for img_idx in range(len(images)):
                # Collect predictions from all levels for this image
                all_boxes = []
                all_scores = []
                all_labels = []

                num_classes = self.num_classes  # K sigmoid outputs per anchor (no background channel)
                for level_idx in range(len(feature_list)):
                    B, C_cls, H, W = classification_logits[level_idx].shape
                    num_anchors = self.num_anchors

                    # Get predictions for this image at this level
                    cls_logits = classification_logits[level_idx][img_idx]  # [num_anchors*K, H, W]
                    bbox_pred = bbox_regression[level_idx][img_idx]  # [num_anchors*5, H, W]

                    # Reshape: [num_anchors*K, H, W] -> [H*W*num_anchors, num_classes]
                    cls_logits_flat = cls_logits.view(num_anchors, num_classes, H, W)
                    cls_logits_flat = cls_logits_flat.permute(2, 3, 0, 1).contiguous().view(-1, num_classes)

                    # Reshape: [num_anchors*5, H, W] -> [H*W*num_anchors, 5]
                    bbox_pred_flat = bbox_pred.view(num_anchors, 5, H, W)
                    bbox_pred_flat = bbox_pred_flat.permute(2, 3, 0, 1).contiguous().view(-1, 5)

                    # Get anchors for this level
                    level_anchors = anchors[level_idx]
                    if isinstance(level_anchors, torch.Tensor):
                        level_anchors = level_anchors.to(cls_logits.device)
                    else:
                        level_anchors = torch.tensor(
                            level_anchors, dtype=torch.float32, device=cls_logits.device
                        )

                    # Decode boxes
                    decoded_boxes = decode_oriented_boxes(
                        level_anchors,
                        bbox_pred_flat,
                        target_means=self.target_means,
                        target_stds=self.target_stds,
                        normalize_le90=True,
                        norm_factor=self.norm_factor,
                        edge_swap=self.edge_swap,
                    )

                    # Sigmoid per class (MMRotate use_sigmoid=True); take best foreground class per anchor
                    class_scores = torch.sigmoid(cls_logits_flat)  # [N, num_classes]
                    max_scores, fg_indices = class_scores.max(dim=1)  # [N], [N] in 0..K-1
                    # 1-indexed class label = 1 + fg_indices (downstream pipeline uses 1-indexed labels)
                    class_indices_1idx = fg_indices + 1  # [N], values 1..num_classes

                    # Pre-NMS score filter (MMRotate nms_pre). Keeps ~10³ candidates instead of
                    # ~2×10⁵ decoded anchors per image; eval still applies evaluation.score_threshold.
                    if self.training:
                        score_thresh = 0.0
                    else:
                        score_thresh = self.score_threshold

                    # Filter by score threshold
                    valid_mask = max_scores >= score_thresh
                    if valid_mask.any():
                        all_boxes.append(decoded_boxes[valid_mask])
                        all_scores.append(max_scores[valid_mask])
                        all_labels.append(class_indices_1idx[valid_mask])

                if not all_boxes:
                    out = {
                        "rboxes": [],
                        "labels": torch.zeros((0,), dtype=torch.int64, device=images_tensor.device),
                        "scores": torch.zeros((0,), dtype=torch.float32, device=images_tensor.device),
                    }
                    if return_debug:
                        out["anchors"] = all_anchors_cat
                        out["proposals"] = torch.zeros((0, 5), dtype=torch.float32)
                    outputs.append(out)
                    continue

                # Concatenate predictions from all levels
                all_boxes = torch.cat(all_boxes, dim=0)
                all_scores = torch.cat(all_scores, dim=0)
                all_labels = torch.cat(all_labels, dim=0)

                # Filter out degenerate boxes (zero or near-zero width/height) before RBox conversion
                min_size = 1.0
                valid = (all_boxes[:, 2] >= min_size) & (all_boxes[:, 3] >= min_size) & torch.isfinite(all_boxes).all(dim=1)
                all_boxes = all_boxes[valid]
                all_scores = all_scores[valid]
                all_labels = all_labels[valid]

                # Save decoded boxes before NMS for debug logging (RetinaNet "proposals" analogue)
                proposals_for_debug = all_boxes.detach().cpu() if return_debug else None

                # Apply oriented NMS (class-aware)
                # Use GPU-accelerated NMS when available (much faster)
                # Use GPU NMS on CUDA or MPS (Apple Silicon)
                device_type = images_tensor.device.type
                use_gpu_nms = (
                    not self.final_nms_use_cpu
                    and torch is not None
                    and len(all_boxes) > 0
                    and (
                        (device_type == "cuda" and torch.cuda.is_available())
                        or (
                            device_type == "mps"
                            and getattr(torch.backends, "mps", None) is not None
                            and torch.backends.mps.is_available()
                        )
                    )
                )

                keep_indices_tensor = None
                if use_gpu_nms:
                    # GPU-accelerated NMS: handle class-aware NMS by grouping by class
                    # Group boxes by class and run NMS per class
                    unique_labels = torch.unique(all_labels)
                    all_keep_indices = []

                    for label in unique_labels:
                        class_mask = all_labels == label
                        class_boxes = all_boxes[class_mask]
                        class_scores = all_scores[class_mask]
                        class_indices = torch.where(class_mask)[0]

                        if len(class_boxes) > 0:
                            # Pre-filter: limit to top boxes per class to avoid memory issues
                            # Use a reasonable limit (e.g., 2000 boxes) before NMS to prevent OOM
                            # This is much larger than the final global limit but prevents huge IoU matrices
                            max_boxes_per_class = 2000
                            if len(class_boxes) > max_boxes_per_class:
                                # Sort by score and take top N
                                _, top_indices = class_scores.sort(descending=True)
                                top_indices = top_indices[:max_boxes_per_class]
                                class_boxes = class_boxes[top_indices]
                                class_scores = class_scores[top_indices]
                                class_indices = class_indices[top_indices]

                            # Run GPU NMS for this class
                            # Pass None to let NMS keep all results, we'll apply global limit later
                            class_keep = rotated_nms(
                                class_boxes,
                                class_scores,
                                iou_threshold=self.final_nms_iou_threshold,
                                max_detections=None,  # No per-class limit - apply globally later
                            )
                            # Map back to original indices
                            all_keep_indices.append(class_indices[class_keep])

                    if all_keep_indices:
                        # Filter out empty tensors before concatenating
                        non_empty_indices = [idx for idx in all_keep_indices if len(idx) > 0]
                        if non_empty_indices:
                            keep_indices_tensor = torch.cat(non_empty_indices)
                            # Sort by score to maintain order
                            keep_scores = all_scores[keep_indices_tensor]
                            _, sort_order = keep_scores.sort(descending=True)
                            keep_indices_tensor = keep_indices_tensor[sort_order]

                            # Apply max_detections limit across all classes
                            if self.max_detections_per_image is not None:
                                keep_indices_tensor = keep_indices_tensor[:self.max_detections_per_image]
                        else:
                            # All classes had empty results after NMS
                            keep_indices_tensor = None
                else:
                    # Fall back to CPU NMS (class-aware). Pre-filter per class like the GPU path.
                    keep_indices: list[int] = []
                    unique_labels = torch.unique(all_labels)
                    max_boxes_per_class = 2000
                    for label in unique_labels:
                        class_mask = all_labels == label
                        class_boxes = all_boxes[class_mask]
                        class_scores = all_scores[class_mask]
                        class_indices = torch.where(class_mask)[0]
                        if len(class_boxes) == 0:
                            continue
                        if len(class_boxes) > max_boxes_per_class:
                            _, top_indices = class_scores.sort(descending=True)
                            top_indices = top_indices[:max_boxes_per_class]
                            class_boxes = class_boxes[top_indices]
                            class_scores = class_scores[top_indices]
                            class_indices = class_indices[top_indices]
                        rboxes = tensor_to_rboxes(class_boxes)
                        class_keep = nms.oriented_nms(
                            boxes=rboxes,
                            scores=class_scores.cpu().tolist(),
                            labels=[int(label.item())] * len(rboxes),
                            iou_threshold=self.final_nms_iou_threshold,
                            max_detections=None,
                        )
                        keep_indices.extend(class_indices[class_keep].tolist())
                    if keep_indices:
                        keep_tensor = torch.tensor(keep_indices, dtype=torch.long, device=all_boxes.device)
                        keep_scores = all_scores[keep_tensor]
                        _, sort_order = keep_scores.sort(descending=True)
                        keep_indices = keep_tensor[sort_order].tolist()
                        if self.max_detections_per_image is not None:
                            keep_indices = keep_indices[: self.max_detections_per_image]
                    else:
                        keep_indices = []

                # Filter results
                if use_gpu_nms and keep_indices_tensor is not None and len(keep_indices_tensor) > 0:
                    # GPU path: use tensor indices directly
                    output_rboxes = tensor_to_rboxes(all_boxes[keep_indices_tensor])
                    output_labels = all_labels[keep_indices_tensor]
                    output_scores = all_scores[keep_indices_tensor]
                elif not use_gpu_nms and keep_indices:
                    # CPU path: convert from list indices
                    rboxes = tensor_to_rboxes(all_boxes)
                    output_rboxes = [rboxes[i] for i in keep_indices]
                    output_labels = all_labels[keep_indices]
                    output_scores = all_scores[keep_indices]
                else:
                    output_rboxes = []
                    output_labels = torch.zeros((0,), dtype=torch.int64, device=images_tensor.device)
                    output_scores = torch.zeros((0,), dtype=torch.float32, device=images_tensor.device)

                out = {
                    "rboxes": output_rboxes,
                    "labels": output_labels,
                    "scores": output_scores,
                }
                if return_debug:
                    out["anchors"] = all_anchors_cat
                    out["proposals"] = proposals_for_debug if proposals_for_debug is not None else torch.zeros((0, 5), dtype=torch.float32)
                outputs.append(out)

            return outputs

forward(images, targets=None)

Forward pass through Rotated RetinaNet.

Parameters:

Name Type Description Default
images Sequence[Tensor]

List of image tensors (C, H, W) in [0, 1] range

 required
targets Optional[Sequence[Dict[str, Any]]]

Optional list of target dicts for training. Each dict must contain: - "rboxes" (List[RBox] or tensor [N, 5] with format [cx, cy, w, h, angle]) - "labels" (tensor [N], 1-indexed)

 None

Returns:

Type Description
Union[Dict[str, Tensor], List[Dict[str, Any]]]
  • Training: Dict with loss keys:
  • "loss_classifier": Classification loss (focal loss)
  • "loss_box_reg": Box regression loss
Union[Dict[str, Tensor], List[Dict[str, Any]]]
  • Inference: List of dicts, one per image, containing:
  • "rboxes": List[RBox] with oriented boxes (with predicted angles)
  • "labels": Tensor [N] with class labels (1-indexed)
  • "scores": Tensor [N] with confidence scores
Source code in oriented_det/models/rotated_retinanet.py 
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
def forward(
    self,
    images: Sequence[torch.Tensor],
    targets: Optional[Sequence[Dict[str, Any]]] = None,
) -> Union[Dict[str, torch.Tensor], List[Dict[str, Any]]]:
    """Forward pass through Rotated RetinaNet.

    Args:
        images: List of image tensors (C, H, W) in [0, 1] range
        targets: Optional list of target dicts for training.
                Each dict must contain:
                - "rboxes" (List[RBox] or tensor [N, 5] with format [cx, cy, w, h, angle])
                - "labels" (tensor [N], 1-indexed)

    Returns:
        - Training: Dict with loss keys:
            - "loss_classifier": Classification loss (focal loss)
            - "loss_box_reg": Box regression loss
        - Inference: List of dicts, one per image, containing:
            - "rboxes": List[RBox] with oriented boxes (with predicted angles)
            - "labels": Tensor [N] with class labels (1-indexed)
            - "scores": Tensor [N] with confidence scores
    """
    if not isinstance(images, (list, tuple)):
        images = [images]

    # Get image sizes
    image_sizes = [(img.shape[-2], img.shape[-1]) for img in images]

    # Extract features using shared utility
    feature_list = extract_backbone_features(
        self.backbone,
        images,
        use_checkpoint=False,  # Rotated RetinaNet doesn't use checkpointing currently
        training=self.training,
        include_pool_level=not self.fpn_extra_level,
    )
    feature_map_sizes = [(f.shape[2], f.shape[3]) for f in feature_list]
    fpn_strides_live = derive_fpn_strides_from_grid(image_sizes[0], feature_map_sizes)
    warn_if_fpn_strides_mismatch(self.fpn_strides, fpn_strides_live)

    # Get device from first feature map (needed for device reference)
    images_tensor = torch.stack(images, dim=0)

    # Forward through Rotated RetinaNet head
    classification_logits, bbox_regression = self.head(feature_list)

    if self.training:
        if targets is None:
            raise ValueError("Targets required during training.")

        # Prepare targets using shared utility
        gt_boxes_list, gt_labels_list, gt_boxes_ignore_list = prepare_targets(targets, device=images_tensor.device)

        # Generate anchors for all FPN levels
        img_h, img_w = image_sizes[0]
        anchors = generate_oriented_anchors(
            image_size=(img_h, img_w),
            feature_map_sizes=feature_map_sizes,
            anchor_scales=self.anchor_scales,
            anchor_ratios=self.anchor_ratios,
            anchor_angles=self.anchor_angles,
            stride_per_level=fpn_strides_live,
            octave_base_scale=self.octave_base_scale,
            scales_per_octave=self.scales_per_octave,
        )

        # Compute losses
        losses = compute_oriented_retinanet_loss(
            classification_logits=classification_logits,
            bbox_regression=bbox_regression,
            anchors=anchors,
            gt_boxes=gt_boxes_list,
            gt_labels=gt_labels_list,
            image_sizes=image_sizes,
            gt_boxes_ignore=gt_boxes_ignore_list,
            num_classes=self.num_classes,
            positive_iou_threshold=self.positive_iou_threshold,
            negative_iou_threshold=self.negative_iou_threshold,
            focal_alpha=self.focal_alpha,
            focal_gamma=self.focal_gamma,
            box_reg_weight=self.box_reg_weight,
            target_means=self.target_means,
            target_stds=self.target_stds,
            target_norm_factor=self.norm_factor,
            norm_factor=self.norm_factor,
            edge_swap=self.edge_swap,
            box_reg_iou_weight=self.box_reg_iou_weight,
            box_reg_iou_loss_type=self.box_reg_iou_loss_type,
            box_reg_kfiou_fun=self.box_reg_kfiou_fun,
            box_reg_probiou_mode=self.box_reg_probiou_mode,
            use_hbb_for_matching=self.use_hbb_for_matching,
            box_reg_loss_type=self.box_reg_loss_type,
        )

        return losses

    else:
        # Inference
        # Generate anchors
        img_h, img_w = image_sizes[0]
        anchors = generate_oriented_anchors(
            image_size=(img_h, img_w),
            feature_map_sizes=feature_map_sizes,
            anchor_scales=self.anchor_scales,
            anchor_ratios=self.anchor_ratios,
            anchor_angles=self.anchor_angles,
            stride_per_level=fpn_strides_live,
            octave_base_scale=self.octave_base_scale,
            scales_per_octave=self.scales_per_octave,
        )
        # Optional debug: expose anchors and decoded boxes (pre-NMS) for TensorBoard
        return_debug = getattr(self, '_return_anchors_proposals', False)
        if return_debug:
            all_anchors_cat = torch.cat(anchors, dim=0).detach().cpu()
        # Process each image
        outputs = []
        for img_idx in range(len(images)):
            # Collect predictions from all levels for this image
            all_boxes = []
            all_scores = []
            all_labels = []

            num_classes = self.num_classes  # K sigmoid outputs per anchor (no background channel)
            for level_idx in range(len(feature_list)):
                B, C_cls, H, W = classification_logits[level_idx].shape
                num_anchors = self.num_anchors

                # Get predictions for this image at this level
                cls_logits = classification_logits[level_idx][img_idx]  # [num_anchors*K, H, W]
                bbox_pred = bbox_regression[level_idx][img_idx]  # [num_anchors*5, H, W]

                # Reshape: [num_anchors*K, H, W] -> [H*W*num_anchors, num_classes]
                cls_logits_flat = cls_logits.view(num_anchors, num_classes, H, W)
                cls_logits_flat = cls_logits_flat.permute(2, 3, 0, 1).contiguous().view(-1, num_classes)

                # Reshape: [num_anchors*5, H, W] -> [H*W*num_anchors, 5]
                bbox_pred_flat = bbox_pred.view(num_anchors, 5, H, W)
                bbox_pred_flat = bbox_pred_flat.permute(2, 3, 0, 1).contiguous().view(-1, 5)

                # Get anchors for this level
                level_anchors = anchors[level_idx]
                if isinstance(level_anchors, torch.Tensor):
                    level_anchors = level_anchors.to(cls_logits.device)
                else:
                    level_anchors = torch.tensor(
                        level_anchors, dtype=torch.float32, device=cls_logits.device
                    )

                # Decode boxes
                decoded_boxes = decode_oriented_boxes(
                    level_anchors,
                    bbox_pred_flat,
                    target_means=self.target_means,
                    target_stds=self.target_stds,
                    normalize_le90=True,
                    norm_factor=self.norm_factor,
                    edge_swap=self.edge_swap,
                )

                # Sigmoid per class (MMRotate use_sigmoid=True); take best foreground class per anchor
                class_scores = torch.sigmoid(cls_logits_flat)  # [N, num_classes]
                max_scores, fg_indices = class_scores.max(dim=1)  # [N], [N] in 0..K-1
                # 1-indexed class label = 1 + fg_indices (downstream pipeline uses 1-indexed labels)
                class_indices_1idx = fg_indices + 1  # [N], values 1..num_classes

                # Pre-NMS score filter (MMRotate nms_pre). Keeps ~10³ candidates instead of
                # ~2×10⁵ decoded anchors per image; eval still applies evaluation.score_threshold.
                if self.training:
                    score_thresh = 0.0
                else:
                    score_thresh = self.score_threshold

                # Filter by score threshold
                valid_mask = max_scores >= score_thresh
                if valid_mask.any():
                    all_boxes.append(decoded_boxes[valid_mask])
                    all_scores.append(max_scores[valid_mask])
                    all_labels.append(class_indices_1idx[valid_mask])

            if not all_boxes:
                out = {
                    "rboxes": [],
                    "labels": torch.zeros((0,), dtype=torch.int64, device=images_tensor.device),
                    "scores": torch.zeros((0,), dtype=torch.float32, device=images_tensor.device),
                }
                if return_debug:
                    out["anchors"] = all_anchors_cat
                    out["proposals"] = torch.zeros((0, 5), dtype=torch.float32)
                outputs.append(out)
                continue

            # Concatenate predictions from all levels
            all_boxes = torch.cat(all_boxes, dim=0)
            all_scores = torch.cat(all_scores, dim=0)
            all_labels = torch.cat(all_labels, dim=0)

            # Filter out degenerate boxes (zero or near-zero width/height) before RBox conversion
            min_size = 1.0
            valid = (all_boxes[:, 2] >= min_size) & (all_boxes[:, 3] >= min_size) & torch.isfinite(all_boxes).all(dim=1)
            all_boxes = all_boxes[valid]
            all_scores = all_scores[valid]
            all_labels = all_labels[valid]

            # Save decoded boxes before NMS for debug logging (RetinaNet "proposals" analogue)
            proposals_for_debug = all_boxes.detach().cpu() if return_debug else None

            # Apply oriented NMS (class-aware)
            # Use GPU-accelerated NMS when available (much faster)
            # Use GPU NMS on CUDA or MPS (Apple Silicon)
            device_type = images_tensor.device.type
            use_gpu_nms = (
                not self.final_nms_use_cpu
                and torch is not None
                and len(all_boxes) > 0
                and (
                    (device_type == "cuda" and torch.cuda.is_available())
                    or (
                        device_type == "mps"
                        and getattr(torch.backends, "mps", None) is not None
                        and torch.backends.mps.is_available()
                    )
                )
            )

            keep_indices_tensor = None
            if use_gpu_nms:
                # GPU-accelerated NMS: handle class-aware NMS by grouping by class
                # Group boxes by class and run NMS per class
                unique_labels = torch.unique(all_labels)
                all_keep_indices = []

                for label in unique_labels:
                    class_mask = all_labels == label
                    class_boxes = all_boxes[class_mask]
                    class_scores = all_scores[class_mask]
                    class_indices = torch.where(class_mask)[0]

                    if len(class_boxes) > 0:
                        # Pre-filter: limit to top boxes per class to avoid memory issues
                        # Use a reasonable limit (e.g., 2000 boxes) before NMS to prevent OOM
                        # This is much larger than the final global limit but prevents huge IoU matrices
                        max_boxes_per_class = 2000
                        if len(class_boxes) > max_boxes_per_class:
                            # Sort by score and take top N
                            _, top_indices = class_scores.sort(descending=True)
                            top_indices = top_indices[:max_boxes_per_class]
                            class_boxes = class_boxes[top_indices]
                            class_scores = class_scores[top_indices]
                            class_indices = class_indices[top_indices]

                        # Run GPU NMS for this class
                        # Pass None to let NMS keep all results, we'll apply global limit later
                        class_keep = rotated_nms(
                            class_boxes,
                            class_scores,
                            iou_threshold=self.final_nms_iou_threshold,
                            max_detections=None,  # No per-class limit - apply globally later
                        )
                        # Map back to original indices
                        all_keep_indices.append(class_indices[class_keep])

                if all_keep_indices:
                    # Filter out empty tensors before concatenating
                    non_empty_indices = [idx for idx in all_keep_indices if len(idx) > 0]
                    if non_empty_indices:
                        keep_indices_tensor = torch.cat(non_empty_indices)
                        # Sort by score to maintain order
                        keep_scores = all_scores[keep_indices_tensor]
                        _, sort_order = keep_scores.sort(descending=True)
                        keep_indices_tensor = keep_indices_tensor[sort_order]

                        # Apply max_detections limit across all classes
                        if self.max_detections_per_image is not None:
                            keep_indices_tensor = keep_indices_tensor[:self.max_detections_per_image]
                    else:
                        # All classes had empty results after NMS
                        keep_indices_tensor = None
            else:
                # Fall back to CPU NMS (class-aware). Pre-filter per class like the GPU path.
                keep_indices: list[int] = []
                unique_labels = torch.unique(all_labels)
                max_boxes_per_class = 2000
                for label in unique_labels:
                    class_mask = all_labels == label
                    class_boxes = all_boxes[class_mask]
                    class_scores = all_scores[class_mask]
                    class_indices = torch.where(class_mask)[0]
                    if len(class_boxes) == 0:
                        continue
                    if len(class_boxes) > max_boxes_per_class:
                        _, top_indices = class_scores.sort(descending=True)
                        top_indices = top_indices[:max_boxes_per_class]
                        class_boxes = class_boxes[top_indices]
                        class_scores = class_scores[top_indices]
                        class_indices = class_indices[top_indices]
                    rboxes = tensor_to_rboxes(class_boxes)
                    class_keep = nms.oriented_nms(
                        boxes=rboxes,
                        scores=class_scores.cpu().tolist(),
                        labels=[int(label.item())] * len(rboxes),
                        iou_threshold=self.final_nms_iou_threshold,
                        max_detections=None,
                    )
                    keep_indices.extend(class_indices[class_keep].tolist())
                if keep_indices:
                    keep_tensor = torch.tensor(keep_indices, dtype=torch.long, device=all_boxes.device)
                    keep_scores = all_scores[keep_tensor]
                    _, sort_order = keep_scores.sort(descending=True)
                    keep_indices = keep_tensor[sort_order].tolist()
                    if self.max_detections_per_image is not None:
                        keep_indices = keep_indices[: self.max_detections_per_image]
                else:
                    keep_indices = []

            # Filter results
            if use_gpu_nms and keep_indices_tensor is not None and len(keep_indices_tensor) > 0:
                # GPU path: use tensor indices directly
                output_rboxes = tensor_to_rboxes(all_boxes[keep_indices_tensor])
                output_labels = all_labels[keep_indices_tensor]
                output_scores = all_scores[keep_indices_tensor]
            elif not use_gpu_nms and keep_indices:
                # CPU path: convert from list indices
                rboxes = tensor_to_rboxes(all_boxes)
                output_rboxes = [rboxes[i] for i in keep_indices]
                output_labels = all_labels[keep_indices]
                output_scores = all_scores[keep_indices]
            else:
                output_rboxes = []
                output_labels = torch.zeros((0,), dtype=torch.int64, device=images_tensor.device)
                output_scores = torch.zeros((0,), dtype=torch.float32, device=images_tensor.device)

            out = {
                "rboxes": output_rboxes,
                "labels": output_labels,
                "scores": output_scores,
            }
            if return_debug:
                out["anchors"] = all_anchors_cat
                out["proposals"] = proposals_for_debug if proposals_for_debug is not None else torch.zeros((0, 5), dtype=torch.float32)
            outputs.append(out)

        return outputs

set_box_reg_iou_weight_for_epoch(epoch)

Update anchor auxiliary IoU loss weight from schedule (0-based epoch index).

Source code in oriented_det/models/rotated_retinanet.py 
605
606
607
608
609
610
611
612
613
614
def set_box_reg_iou_weight_for_epoch(self, epoch: int) -> None:
    """Update anchor auxiliary IoU loss weight from schedule (0-based epoch index)."""
    from oriented_det.train.piecewise_schedule import resolve_piecewise_schedule

    self.box_reg_iou_weight = resolve_piecewise_schedule(
        epoch,
        self._roi_box_reg_iou_schedule_epochs,
        self._roi_box_reg_iou_schedule_values,
        self._box_reg_iou_weight_default,
    )

set_final_nms_iou_for_epoch(epoch)

Update final detection NMS IoU threshold from schedule for the given epoch. Lower threshold = more aggressive suppression of overlapping boxes.

Source code in oriented_det/models/rotated_retinanet.py 
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
def set_final_nms_iou_for_epoch(self, epoch: int) -> None:
    """Update final detection NMS IoU threshold from schedule for the given epoch.
    Lower threshold = more aggressive suppression of overlapping boxes.
    """
    if self._final_nms_iou_schedule_epochs is None or self._final_nms_iou_schedule_values is None:
        return
    if not self._final_nms_iou_schedule_epochs or not self._final_nms_iou_schedule_values:
        return
    idx = 0
    for boundary in self._final_nms_iou_schedule_epochs:
        if epoch < boundary:
            break
        idx += 1
    idx = min(idx, len(self._final_nms_iou_schedule_values) - 1)
    self.final_nms_iou_threshold = self._final_nms_iou_schedule_values[idx]

build_resnet_fpn_backbone(backbone_name='resnet50', *, pretrained=False, trainable_layers=3, norm_layer=None, returned_layers=None, use_p6p7_extra_levels=False)

Create a ResNet+FPN backbone or a minimal fallback if torchvision is missing.

Parameters:

Name Type Description Default
backbone_name str

ResNet variant ("resnet50", etc.).

 'resnet50'
pretrained bool

Whether to load ImageNet backbone weights.

 False
trainable_layers int

Number of backbone stages to train.

 3
norm_layer

Norm layer override. Default None keeps torchvision's FrozenBatchNorm2d (frozen BN statistics), matching MMRotate's norm_eval=True detection recipe. Pass torch.nn.BatchNorm2d explicitly to train with live batch statistics.

 None
returned_layers Optional[List[int]]

ResNet stage indices to use for FPN (1–4). Default None uses torchvision default [1,2,3,4] (C2–C5, 5 levels with P6). Use [2,3,4] for MMRotate-style FPN (C3–C5 only, 4 levels with P6).

 None
use_p6p7_extra_levels bool

If True, attach LastLevelP6P7 (MMRotate RetinaNet add_extra_convs='on_input') instead of LastLevelMaxPool.

 False
Source code in oriented_det/models/backbones/resnet_fpn.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def build_resnet_fpn_backbone(
    backbone_name: str = "resnet50",
    *,
    pretrained: bool = False,
    trainable_layers: int = 3,
    norm_layer=None,
    returned_layers: Optional[List[int]] = None,
    use_p6p7_extra_levels: bool = False,
) -> object:
    """Create a ResNet+FPN backbone or a minimal fallback if torchvision is missing.

    Args:
        backbone_name: ResNet variant ("resnet50", etc.).
        pretrained: Whether to load ImageNet backbone weights.
        trainable_layers: Number of backbone stages to train.
        norm_layer: Norm layer override. Default None keeps torchvision's
            ``FrozenBatchNorm2d`` (frozen BN statistics), matching MMRotate's
            ``norm_eval=True`` detection recipe. Pass ``torch.nn.BatchNorm2d``
            explicitly to train with live batch statistics.
        returned_layers: ResNet stage indices to use for FPN (1–4). Default None
            uses torchvision default [1,2,3,4] (C2–C5, 5 levels with P6).
            Use [2,3,4] for MMRotate-style FPN (C3–C5 only, 4 levels with P6).
        use_p6p7_extra_levels: If True, attach ``LastLevelP6P7`` (MMRotate RetinaNet
            ``add_extra_convs='on_input'``) instead of ``LastLevelMaxPool``.
    """
    require_torch()
    if tv_resnet_fpn_backbone is None:
        return make_single_feature_backbone()

    # Do not pass norm_layer=None through: torchvision would then fall back to live
    # BatchNorm2d. Omitting it keeps torchvision's FrozenBatchNorm2d default.
    kwargs = dict(trainable_layers=trainable_layers)
    if norm_layer is not None:
        kwargs["norm_layer"] = norm_layer
    signature = inspect.signature(tv_resnet_fpn_backbone)
    if "weights" in signature.parameters:
        kwargs["weights"] = None if not pretrained else "DEFAULT"
    if "weights_backbone" in signature.parameters:
        kwargs["weights_backbone"] = "DEFAULT" if pretrained else None
    if "returned_layers" in signature.parameters and returned_layers is not None:
        kwargs["returned_layers"] = returned_layers
    if use_p6p7_extra_levels and "extra_blocks" in signature.parameters:
        from torchvision.ops.feature_pyramid_network import LastLevelP6P7

        c5_channels = 512 if any(x in backbone_name.lower() for x in ("18", "34")) else 2048
        kwargs["extra_blocks"] = LastLevelP6P7(c5_channels, 256)

    try:
        backbone = tv_resnet_fpn_backbone(backbone_name=backbone_name, **kwargs)
        return backbone
    except Exception as exc:  # pragma: no cover
        # Offline / restricted environments (CI, sandboxes) may block weight downloads.
        # Fall back to random init while still returning a functional backbone.
        if pretrained:
            import warnings

            warnings.warn(
                f"Failed to load pretrained backbone weights for {backbone_name!r} ({exc}). "
                "Falling back to random initialization.",
                RuntimeWarning,
            )
            kwargs2 = dict(kwargs)
            if "weights" in kwargs2:
                kwargs2["weights"] = None
            if "weights_backbone" in kwargs2:
                kwargs2["weights_backbone"] = None
            backbone = tv_resnet_fpn_backbone(backbone_name=backbone_name, **kwargs2)
            return backbone
        raise

derive_fpn_strides_from_grid(image_size, feature_map_sizes)

Compute FPN strides from the real input size and backbone feature map shapes.

For each level, stride is image_size / feature_map_size (H and W must agree). This keeps anchors, RPN, and ROI align consistent with the actual grid even if fpn_strides in config is wrong or stale.

Parameters:

Name Type Description Default
image_size Tuple[int, int]

(H, W) of the image tensor fed to the backbone (same as anchor generation).

 required
feature_map_sizes Sequence[Tuple[int, int]]

(h, w) per FPN level, from feature.shape[2:].

 required

Returns:

Type Description
List[int]

One integer stride per level.

Raises:

Type Description
ValueError

If sizes are invalid or H/W imply different strides.
Source code in oriented_det/models/utils.py 
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
def derive_fpn_strides_from_grid(
    image_size: Tuple[int, int],
    feature_map_sizes: Sequence[Tuple[int, int]],
) -> List[int]:
    """Compute FPN strides from the real input size and backbone feature map shapes.

    For each level, stride is ``image_size / feature_map_size`` (H and W must agree).
    This keeps anchors, RPN, and ROI align consistent with the actual grid even if
    ``fpn_strides`` in config is wrong or stale.

    Args:
        image_size: ``(H, W)`` of the image tensor fed to the backbone (same as anchor generation).
        feature_map_sizes: ``(h, w)`` per FPN level, from ``feature.shape[2:]``.

    Returns:
        One integer stride per level.

    Raises:
        ValueError: If sizes are invalid or H/W imply different strides.
    """
    img_h, img_w = int(image_size[0]), int(image_size[1])
    if img_h <= 0 or img_w <= 0:
        raise ValueError(f"image_size must be positive, got {image_size}")
    strides: List[int] = []
    for fh, fw in feature_map_sizes:
        fh, fw = int(fh), int(fw)
        if fh <= 0 or fw <= 0:
            raise ValueError(f"feature_map_sizes entries must be positive, got {(fh, fw)}")
        sh = img_h / fh
        sw = img_w / fw
        tol = max(1e-2, 1e-3 * max(sh, sw))
        if abs(sh - sw) > tol:
            raise ValueError(
                f"Anisotropic FPN stride for image {img_h}x{img_w} and feature map {fh}x{fw}: "
                f"stride_h={sh:.6f}, stride_w={sw:.6f}"
            )
        s = int(round((sh + sw) / 2.0))
        if s <= 0:
            raise ValueError(f"Non-positive derived stride {s} for image {image_size}, feature {(fh, fw)}")
        strides.append(s)
    return strides

extract_backbone_features(backbone, images, use_checkpoint=False, training=False, include_pool_level=False)

Extract features from backbone and convert to list format.

Handles: - Stacking images into batch - Optional gradient checkpointing - Converting OrderedDict/dict outputs to list format

Parameters:

Name Type Description Default
backbone Module

Backbone module

 required
images Sequence[Tensor]

Sequence of image tensors (C, H, W) in [0, 1] range

 required
use_checkpoint bool

Whether to use gradient checkpointing

 False
training bool

Whether model is in training mode

 False
include_pool_level bool

If True, also keep torchvision FPN's LastLevelMaxPool output (dict key "pool") as the last pyramid level. RetinaNet needs it (P6, stride 2x the last conv level); the two-stage models ignore it.

 False

Returns:

Type Description
List[Tensor]

List of feature maps from FPN, each of shape [B, C, H, W]
Source code in oriented_det/models/utils.py 
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
def extract_backbone_features(
    backbone: nn.Module,
    images: Sequence[torch.Tensor],
    use_checkpoint: bool = False,
    training: bool = False,
    include_pool_level: bool = False,
) -> List[torch.Tensor]:
    """Extract features from backbone and convert to list format.

    Handles:
    - Stacking images into batch
    - Optional gradient checkpointing
    - Converting OrderedDict/dict outputs to list format

    Args:
        backbone: Backbone module
        images: Sequence of image tensors (C, H, W) in [0, 1] range
        use_checkpoint: Whether to use gradient checkpointing
        training: Whether model is in training mode
        include_pool_level: If True, also keep torchvision FPN's ``LastLevelMaxPool``
            output (dict key ``"pool"``) as the last pyramid level. RetinaNet needs it
            (P6, stride 2x the last conv level); the two-stage models ignore it.

    Returns:
        List of feature maps from FPN, each of shape [B, C, H, W]
    """
    if torch is None:
        raise RuntimeError("PyTorch is required.")

    # Stack images into batch
    images_tensor = torch.stack(images, dim=0)  # [B, C, H, W]

    # Extract features using backbone
    # OPTIONAL: Use gradient checkpointing to reduce memory usage
    # This trades computation (recomputes activations) for memory (2-3x reduction)
    if use_checkpoint and training:
        try:
            from torch.utils.checkpoint import checkpoint
            # Wrap backbone forward in checkpoint - will recompute during backward
            features = checkpoint(backbone, images_tensor, use_reentrant=False)
        except Exception:
            # Fallback if checkpointing fails
            features = backbone(images_tensor)
    else:
        features = backbone(images_tensor)

    # Convert OrderedDict/dict to list (FPN outputs)
    if isinstance(features, dict):
        # FPN returns OrderedDict with keys like "0", "1", "2", etc.
        # or sometimes "fpn0", "fpn1", etc.
        feature_list = []
        for key in sorted(features.keys()):
            if key.isdigit() or key.startswith("fpn"):
                feature_list.append(features[key])
        # torchvision's LastLevelMaxPool level (key "pool") is the coarsest level;
        # keep it last when requested (RetinaNet P6).
        if include_pool_level and "pool" in features:
            feature_list.append(features["pool"])
        # If no numeric keys, just use all values
        if not feature_list:
            feature_list = list(features.values())
    else:
        # Single feature map
        feature_list = [features]

    return feature_list

prepare_targets(targets, device=None)

Prepare targets for training, preserving oriented boxes.

Parameters:

Name Type Description Default
targets Sequence[Dict[str, Any]]

List of target dicts, each containing: - "rboxes" (List[RBox] or tensor [N, 5] with format [cx, cy, w, h, angle]) - "labels" (tensor [N] or list, 1-indexed)

 required
device Optional[device]

Optional device to move tensors to

 None

Returns:

Type Description
List[Tensor]

Tuple of (gt_boxes_list, gt_labels_list, gt_boxes_ignore_list):
List[Tensor]
  • gt_boxes_list: List of tensors [M, 5] with oriented boxes
List[Tensor]
  • gt_labels_list: List of tensors [M] with class labels (1-indexed)
Tuple[List[Tensor], List[Tensor], List[Tensor]]
  • gt_boxes_ignore_list: List of tensors [I, 5] with ignored GT boxes (e.g. DOTA difficult)
Source code in oriented_det/models/utils.py 
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
def prepare_targets(
    targets: Sequence[Dict[str, Any]],
    device: Optional[torch.device] = None,
) -> Tuple[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor]]:
    """Prepare targets for training, preserving oriented boxes.

    Args:
        targets: List of target dicts, each containing:
            - "rboxes" (List[RBox] or tensor [N, 5] with format [cx, cy, w, h, angle])
            - "labels" (tensor [N] or list, 1-indexed)
        device: Optional device to move tensors to

    Returns:
        Tuple of (gt_boxes_list, gt_labels_list, gt_boxes_ignore_list):
        - gt_boxes_list: List of tensors [M, 5] with oriented boxes
        - gt_labels_list: List of tensors [M] with class labels (1-indexed)
        - gt_boxes_ignore_list: List of tensors [I, 5] with ignored GT boxes (e.g. DOTA difficult)
    """
    if torch is None:
        raise RuntimeError("PyTorch is required.")

    gt_boxes_list = []
    gt_labels_list = []
    gt_boxes_ignore_list: List[torch.Tensor] = []

    for target in targets:
        target = dict(target)

        # Get oriented boxes
        if "rboxes" in target:
            rboxes = target["rboxes"]
            if isinstance(rboxes, torch.Tensor):
                gt_boxes = rboxes
                if device is not None:
                    gt_boxes = gt_boxes.to(device)
            else:
                # rboxes is a list of RBox objects
                gt_boxes = rboxes_to_tensor(rboxes, device=device)
        else:
            raise ValueError("Targets must include 'rboxes'.")

        # Get labels
        labels = target.get("labels")
        if labels is None:
            raise ValueError("Targets must include 'labels'.")
        if not torch.is_tensor(labels):
            labels = torch.tensor(labels, dtype=torch.int64, device=device)
        elif device is not None:
            labels = labels.to(device)

        gt_boxes_list.append(gt_boxes)
        gt_labels_list.append(labels)

        # Optional ignore GTs (MMDet/MMRotate style): used to mark anchors/proposals as ignore
        # when they overlap with difficult / don't-care regions.
        if "rboxes_ignore" in target:
            rboxes_ign = target["rboxes_ignore"]
            if isinstance(rboxes_ign, torch.Tensor):
                gt_ign = rboxes_ign.to(device) if device is not None else rboxes_ign
            else:
                gt_ign = rboxes_to_tensor(rboxes_ign, device=device)
        else:
            gt_ign = torch.zeros((0, 5), dtype=torch.float32, device=device)
        gt_boxes_ignore_list.append(gt_ign)

    return gt_boxes_list, gt_labels_list, gt_boxes_ignore_list

rboxes_to_tensor(rboxes, device=None)

Convert RBoxes to tensor format [N, 5] with [cx, cy, w, h, angle].

Parameters:

Name Type Description Default
rboxes Sequence[RBox]

Sequence of RBox objects

 required
device Optional[device]

Optional device to create tensor on

 None

Returns:

Type Description
Tensor

Tensor of shape [N, 5] with format [cx, cy, w, h, angle]
Source code in oriented_det/models/utils.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def rboxes_to_tensor(
    rboxes: Sequence[RBox], 
    device: Optional[torch.device] = None
) -> torch.Tensor:
    """Convert RBoxes to tensor format [N, 5] with [cx, cy, w, h, angle].

    Args:
        rboxes: Sequence of RBox objects
        device: Optional device to create tensor on

    Returns:
        Tensor of shape [N, 5] with format [cx, cy, w, h, angle]
    """
    if torch is None:
        raise RuntimeError("PyTorch is required.")

    if not rboxes:
        return torch.zeros((0, 5), dtype=torch.float32, device=device)

    data = []
    for rbox in rboxes:
        data.append([rbox.cx, rbox.cy, rbox.width, rbox.height, rbox.angle])

    return torch.tensor(data, dtype=torch.float32, device=device)

setup_anchors(anchor_scales, anchor_ratios, anchor_angles, default_angles, octave_base_scale=None, scales_per_octave=None)

Setup anchor configuration and return scales, ratios, angles, and num_anchors.

Parameters:

Name Type Description Default
anchor_scales Optional[List[float]]

Optional list of anchor scales (default: [8])

 required
anchor_ratios Optional[List[float]]

Optional list of anchor ratios (default: [0.5, 1.0, 2.0])

 required
anchor_angles Optional[List[float]]

Optional list of anchor angles (default: provided)

 required
default_angles List[float]

Default angles to use if anchor_angles is None

 required
octave_base_scale Optional[float]

MMRotate RetinaNet octave_base_scale (with scales_per_octave)

 None
scales_per_octave Optional[int]

MMRotate RetinaNet scales_per_octave (e.g. 3)

 None

Returns:

Type Description
List[float]

Tuple of (scales, ratios, angles, num_anchors):
List[float]
  • scales: List of anchor scales
List[float]
  • ratios: List of anchor ratios
int
  • angles: List of anchor angles
Tuple[List[float], List[float], List[float], int]
  • num_anchors: Total number of anchors per location
Source code in oriented_det/models/utils.py 
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
def setup_anchors(
    anchor_scales: Optional[List[float]],
    anchor_ratios: Optional[List[float]],
    anchor_angles: Optional[List[float]],
    default_angles: List[float],
    octave_base_scale: Optional[float] = None,
    scales_per_octave: Optional[int] = None,
) -> Tuple[List[float], List[float], List[float], int]:
    """Setup anchor configuration and return scales, ratios, angles, and num_anchors.

    Args:
        anchor_scales: Optional list of anchor scales (default: [8])
        anchor_ratios: Optional list of anchor ratios (default: [0.5, 1.0, 2.0])
        anchor_angles: Optional list of anchor angles (default: provided)
        default_angles: Default angles to use if anchor_angles is None
        octave_base_scale: MMRotate RetinaNet ``octave_base_scale`` (with ``scales_per_octave``)
        scales_per_octave: MMRotate RetinaNet ``scales_per_octave`` (e.g. 3)

    Returns:
        Tuple of (scales, ratios, angles, num_anchors):
        - scales: List of anchor scales
        - ratios: List of anchor ratios
        - angles: List of anchor angles
        - num_anchors: Total number of anchors per location
    """
    scales = anchor_scales or [8]  # MMRotate standard: single scale
    ratios = anchor_ratios or [0.5, 1.0, 2.0]  # MMRotate standard
    # Handle empty list case: use default_angles if anchor_angles is None or empty
    if anchor_angles is None or (isinstance(anchor_angles, list) and len(anchor_angles) == 0):
        angles = default_angles
    else:
        angles = anchor_angles
    per_loc = len(ratios) * len(angles)
    if octave_base_scale is not None and scales_per_octave is not None:
        num_anchors = per_loc * int(scales_per_octave)
    else:
        num_anchors = per_loc

    return scales, ratios, angles, num_anchors

setup_backbone(backbone, backbone_name='resnet50', pretrained_backbone=False, trainable_layers=5, returned_layers=None, use_p6p7_extra_levels=False)

Setup backbone and return backbone module and output channels.

Parameters:

Name Type Description Default
backbone Optional[Module]

Optional backbone module (if None, creates ResNet+FPN)

 required
backbone_name str

Name of backbone to create ("resnet18", "resnet50", etc.)

 'resnet50'
pretrained_backbone bool

Whether to use pretrained backbone weights

 False
trainable_layers int

Number of backbone layers to keep trainable

 5
returned_layers Optional[List[int]]

ResNet stages for FPN (e.g. [2,3,4] for MMRotate C3–C5 only)

 None
use_p6p7_extra_levels bool

RetinaNet-style P6/P7 convs on C5 (MMRotate on_input).

 False

Returns:

Type Description
Module

Tuple of (backbone, output_channels):
int
  • backbone: The backbone module
Tuple[Module, int]
  • output_channels: Number of output channels (typically 256 for FPN)
Source code in oriented_det/models/utils.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
def setup_backbone(
    backbone: Optional[nn.Module],
    backbone_name: str = "resnet50",
    pretrained_backbone: bool = False,
    trainable_layers: int = 5,
    returned_layers: Optional[List[int]] = None,
    use_p6p7_extra_levels: bool = False,
) -> Tuple[nn.Module, int]:
    """Setup backbone and return backbone module and output channels.

    Args:
        backbone: Optional backbone module (if None, creates ResNet+FPN)
        backbone_name: Name of backbone to create ("resnet18", "resnet50", etc.)
        pretrained_backbone: Whether to use pretrained backbone weights
        trainable_layers: Number of backbone layers to keep trainable
        returned_layers: ResNet stages for FPN (e.g. [2,3,4] for MMRotate C3–C5 only)
        use_p6p7_extra_levels: RetinaNet-style P6/P7 convs on C5 (MMRotate on_input).

    Returns:
        Tuple of (backbone, output_channels):
        - backbone: The backbone module
        - output_channels: Number of output channels (typically 256 for FPN)
    """
    if nn is None:
        raise RuntimeError("PyTorch is required.")

    if backbone is None:
        backbone = build_resnet_fpn_backbone(
            backbone_name,
            pretrained=pretrained_backbone,
            trainable_layers=trainable_layers,
            returned_layers=returned_layers,
            use_p6p7_extra_levels=use_p6p7_extra_levels,
        )

    # Get backbone output channels (typically 256 for FPN)
    if hasattr(backbone, 'out_channels'):
        backbone_channels = backbone.out_channels
    else:
        # Try to infer from FPN if available
        if hasattr(backbone, 'fpn') and hasattr(backbone.fpn, 'out_channels'):
            backbone_channels = backbone.fpn.out_channels
        else:
            # Default for FPN
            backbone_channels = 256

    return backbone, backbone_channels

tensor_to_rboxes(boxes)

Convert tensor [N, 5] with [cx, cy, w, h, angle] to RBoxes.

Parameters:

Name Type Description Default
boxes Tensor

Tensor of shape [N, 5] with format [cx, cy, w, h, angle]

 required

Returns:

Type Description
List[RBox]

List of RBox objects
Source code in oriented_det/models/utils.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def tensor_to_rboxes(boxes: torch.Tensor) -> List[RBox]:
    """Convert tensor [N, 5] with [cx, cy, w, h, angle] to RBoxes.

    Args:
        boxes: Tensor of shape [N, 5] with format [cx, cy, w, h, angle]

    Returns:
        List of RBox objects
    """
    rboxes = []
    for box in boxes.detach().cpu():
        cx, cy, w, h, angle = box.tolist()
        # Validate coordinates are reasonable (not NaN/inf and within expected range)
        if not (math.isfinite(cx) and math.isfinite(cy) and math.isfinite(w) and 
                math.isfinite(h) and math.isfinite(angle)):
            continue  # Skip invalid boxes
        rboxes.append(RBox(cx, cy, w, h, angle))
    return rboxes

warn_if_fpn_strides_mismatch(configured, derived)

Emit a single warning if config strides disagree with grid-derived strides.

Source code in oriented_det/models/utils.py 
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
def warn_if_fpn_strides_mismatch(
    configured: Optional[Sequence[int]],
    derived: Sequence[int],
) -> None:
    """Emit a single warning if config strides disagree with grid-derived strides."""
    if not configured:
        return
    cfg_full = [int(x) for x in configured]
    der = [int(x) for x in derived]
    if len(cfg_full) != len(der):
        warnings.warn(
            f"fpn_strides from config has {len(cfg_full)} levels {cfg_full} but the model "
            f"produced {len(der)} feature levels (derived strides {der}); using derived strides. "
            "Check FPN level configuration (returned_layers / fpn_extra_level).",
            UserWarning,
            stacklevel=3,
        )
        return
    if cfg_full == der:
        return
    warnings.warn(
        f"fpn_strides from config {cfg_full} does not match grid-derived strides {der}; "
        "using derived strides. Fix fpn_strides or training image size to remove this warning.",
        UserWarning,
        stacklevel=3,
    )