edgemdt_cl.pytorch

54def multiclass_nms(boxes, scores, score_threshold: float, iou_threshold: float, max_detections: int) -> NMSResults:
55    """
56    Multi-class non-maximum suppression.
57    Detections are returned in descending order of their scores.
58    The output tensors always contain a fixed number of detections, as defined by 'max_detections'.
59    If fewer detections are selected, the output tensors are zero-padded up to 'max_detections'.
60
61    If you also require the input indices of the selected boxes, see `multiclass_nms_with_indices`.
62
63    Args:
64        boxes (Tensor): Input boxes with shape [batch, n_boxes, 4], specified in corner coordinates
65                        (x_min, y_min, x_max, y_max). Agnostic to the x-y axes order.
66        scores (Tensor): Input scores with shape [batch, n_boxes, n_classes].
67        score_threshold (float): The score threshold. Candidates with scores below the threshold are discarded.
68        iou_threshold (float): The Intersection Over Union (IOU) threshold for boxes overlap.
69        max_detections (int): The number of detections to return.
70
71    Returns:
72        'NMSResults' named tuple:
73        - boxes: The selected boxes with shape [batch, max_detections, 4].
74        - scores: The corresponding scores in descending order with shape [batch, max_detections].
75        - labels: The labels for each box with shape [batch, max_detections].
76        - n_valid: The number of valid detections out of 'max_detections' with shape [batch, 1]
77
78    Raises:
79        ValueError: If provided with invalid arguments or input tensors with unexpected or non-matching shapes.
80
81    Example:
82        ```
83        from edgemdt_cl.pytorch import multiclass_nms
84
85        # batch size=1, 1000 boxes, 50 classes
86        boxes = torch.rand(1, 1000, 4)
87        scores = torch.rand(1, 1000, 50)
88        res = multiclass_nms(boxes,
89                             scores,
90                             score_threshold=0.1,
91                             iou_threshold=0.6,
92                             max_detections=300)
93        # res.boxes, res.scores, res.labels, res.n_valid
94        ```
95    """
96    return NMSResults(*torch.ops.edgemdt.multiclass_nms(boxes, scores, score_threshold, iou_threshold, max_detections))

32class NMSResults(NamedTuple):
33    """ Container for non-maximum suppression results """
34    boxes: Tensor
35    scores: Tensor
36    labels: Tensor
37    n_valid: Tensor
38
39    # Note: convenience methods below are replicated in each Results container, since NamedTuple supports neither adding
40    # new fields in derived classes nor multiple inheritance, and we want it to behave like a tuple, so no dataclasses.
41    def detach(self) -> 'NMSResults':
42        """ Detach all tensors and return a new object """
43        return self.apply(lambda t: t.detach())
44
45    def cpu(self) -> 'NMSResults':
46        """ Move all tensors to cpu and return a new object """
47        return self.apply(lambda t: t.cpu())
48
49    def apply(self, f: Callable[[Tensor], Tensor]) -> 'NMSResults':
50        """ Apply any function to all tensors and return a new object """
51        return self.__class__(*[f(t) for t in self])

54def multiclass_nms_with_indices(boxes, scores, score_threshold: float, iou_threshold: float,
55                                max_detections: int) -> NMSWithIndicesResults:
56    """
57    Multi-class non-maximum suppression with indices.
58    Detections are returned in descending order of their scores.
59    The output tensors always contain a fixed number of detections, as defined by 'max_detections'.
60    If fewer detections are selected, the output tensors are zero-padded up to 'max_detections'.
61
62    This operator is identical to `multiclass_nms` except that is also outputs the input indices of the selected boxes.
63
64    Args:
65        boxes (Tensor): Input boxes with shape [batch, n_boxes, 4], specified in corner coordinates
66                        (x_min, y_min, x_max, y_max). Agnostic to the x-y axes order.
67        scores (Tensor): Input scores with shape [batch, n_boxes, n_classes].
68        score_threshold (float): The score threshold. Candidates with scores below the threshold are discarded.
69        iou_threshold (float): The Intersection Over Union (IOU) threshold for boxes overlap.
70        max_detections (int): The number of detections to return.
71
72    Returns:
73        'NMSWithIndicesResults' named tuple:
74        - boxes: The selected boxes with shape [batch, max_detections, 4].
75        - scores: The corresponding scores in descending order with shape [batch, max_detections].
76        - labels: The labels for each box with shape [batch, max_detections].
77        - indices: Indices of the input boxes that have been selected.
78        - n_valid: The number of valid detections out of 'max_detections' with shape [batch, 1]
79
80    Raises:
81        ValueError: If provided with invalid arguments or input tensors with unexpected or non-matching shapes.
82
83    Example:
84        ```
85        from edgemdt_cl.pytorch import multiclass_nms_with_indices
86
87        # batch size=1, 1000 boxes, 50 classes
88        boxes = torch.rand(1, 1000, 4)
89        scores = torch.rand(1, 1000, 50)
90        res = multiclass_nms_with_indices(boxes,
91                                          scores,
92                                          score_threshold=0.1,
93                                          iou_threshold=0.6,
94                                          max_detections=300)
95        # res.boxes, res.scores, res.labels, res.indices, res.n_valid
96        ```
97    """
98    return NMSWithIndicesResults(
99        *torch.ops.edgemdt.multiclass_nms_with_indices(boxes, scores, score_threshold, iou_threshold, max_detections))

31class NMSWithIndicesResults(NamedTuple):
32    """ Container for non-maximum suppression with indices results """
33    boxes: Tensor
34    scores: Tensor
35    labels: Tensor
36    indices: Tensor
37    n_valid: Tensor
38
39    # Note: convenience methods below are replicated in each Results container, since NamedTuple supports neither adding
40    # new fields in derived classes nor multiple inheritance, and we want it to behave like a tuple, so no dataclasses.
41    def detach(self) -> 'NMSWithIndicesResults':
42        """ Detach all tensors and return a new object """
43        return self.apply(lambda t: t.detach())
44
45    def cpu(self) -> 'NMSWithIndicesResults':
46        """ Move all tensors to cpu and return a new object """
47        return self.apply(lambda t: t.cpu())
48
49    def apply(self, f: Callable[[Tensor], Tensor]) -> 'NMSWithIndicesResults':
50        """ Apply any function to all tensors and return a new object """
51        return self.__class__(*[f(t) for t in self])

31class FasterRCNNBoxDecode(nn.Module):
32    """
33    Box decoding as per Faster R-CNN <https://arxiv.org/abs/1506.01497>.
34
35    Args:
36        anchors: Anchors with a shape of (n_boxes, 4) in corner coordinates (y_min, x_min, y_max, x_max).
37        scale_factors: Scaling factors in the format (y, x, height, width).
38        clip_window: Clipping window in the format (y_min, x_min, y_max, x_max).
39
40    Inputs:
41        **rel_codes** (Tensor): Relative codes (encoded offsets) with a shape of (batch, n_boxes, 4) in centroid
42                                coordinates (y_center, x_center, h, w).
43
44    Returns:
45        Decoded boxes with a shape of (batch, n_boxes, 4) in corner coordinates (y_min, x_min, y_max, x_max).
46
47    Raises:
48        ValueError: If provided with invalid arguments or an input tensor with unexpected shape
49
50    Example:
51        ```
52        from edgemdt_cl.pytorch import FasterRCNNBoxDecode
53
54        box_decode = FasterRCNNBoxDecode(anchors,
55                                         scale_factors=(10, 10, 5, 5),
56                                         clip_window=(0, 0, 1, 1))
57        decoded_boxes = box_decode(rel_codes)
58        ```
59    """
60
61    def __init__(self, anchors: torch.Tensor, scale_factors: Sequence[Union[float, int]],
62                 clip_window: Sequence[Union[float, int]]):
63        super().__init__()
64        warnings.warn("FasterRCNNBoxDecode is deprecated and will be removed in a future version.", DeprecationWarning)
65
66        if not (len(anchors.shape) == 2 and anchors.shape[-1] == 4):
67            raise ValueError(f'Invalid anchors shape {anchors.shape}. Expected shape (n_boxes, 4).')
68        self.register_buffer('anchors', anchors)
69
70        if len(scale_factors) != 4:
71            raise ValueError(f'Invalid scale factors {scale_factors}. Expected 4 values for (y, x, height, width).')
72        self.register_buffer('scale_factors', torch.tensor(scale_factors, dtype=torch.float32, device=anchors.device))
73
74        if len(clip_window) != 4:
75            raise ValueError(f'Invalid clip window {clip_window}. Expected 4 values for (y_min, x_min, y_max, x_max).')
76        self.register_buffer('clip_window', torch.tensor(clip_window, dtype=torch.float32, device=anchors.device))
77
78    def forward(self, rel_codes: torch.Tensor) -> torch.Tensor:
79        return torch.ops.edgemdt.faster_rcnn_box_decode(rel_codes, self.anchors, self.scale_factors, self.clip_window)

40def load_custom_ops(ort_session_ops: Optional['ort.SessionOptions'] = None) -> 'ort.SessionOptions':
41    """
42    Registers the custom ops implementation for onnxruntime, and sets up the SessionOptions object for onnxruntime
43    session.
44
45    Args:
46        ort_session_ops: SessionOptions object to register the custom ops library on. If None, creates a new object.
47
48    Returns:
49        SessionOptions object with registered custom ops.
50
51    Example:
52        ```
53        import onnxruntime as ort
54        from edgemdt_cl.pytorch import load_custom_ops
55
56        so = load_custom_ops()
57        session = ort.InferenceSession(model_path, sess_options=so)
58        session.run(...)
59        ```
60        You can also pass your own SessionOptions object upon which to register the custom ops
61        ```
62        load_custom_ops(ort_session_options=so)
63        ```
64    """
65    validate_installed_libraries(required_libraries['torch_ort'])
66
67    # trigger onnxruntime op registration
68    from .nms import nms_ort
69    from .nms_obb import nms_obb_ort
70    from .box_decode import box_decode_ort
71
72    from onnxruntime_extensions import get_library_path
73    from onnxruntime import SessionOptions
74    ort_session_ops = ort_session_ops or SessionOptions()
75    ort_session_ops.register_custom_ops_library(get_library_path())
76    return ort_session_ops

 99class MulticlassNMS(CustomLayer):
100    """
101    A torch.nn.Module for multiclass NMS. See multiclass_nms for additional information.
102
103    Usage example:
104        batch size=1, 1000 boxes, 50 classes
105        boxes = torch.rand(1, 1000, 4)
106        scores = torch.rand(1, 1000, 50)
107        nms = MulticlassNMS(score_threshold=0.1,
108                            iou_threshold=0.6
109                            max_detections=300)
110        res = nms(boxes, scores)
111    """
112
113    def __init__(self, score_threshold: float, iou_threshold: float, max_detections: int):
114        """
115        Args:
116            score_threshold (float): The score threshold. Candidates with scores below the threshold are discarded.
117            iou_threshold (float): The Intersection Over Union (IOU) threshold for boxes overlap.
118            max_detections (int): The number of detections to return.
119        """
120        super(MulticlassNMS, self).__init__()
121        self.score_threshold = score_threshold
122        self.iou_threshold = iou_threshold
123        self.max_detections = max_detections
124
125    def forward(self, boxes: torch.Tensor, scores: torch.Tensor):
126        """
127        Args:
128            boxes (Tensor): Input boxes with shape [batch, n_boxes, 4], specified in corner coordinates
129                            (x_min, y_min, x_max, y_max). Agnostic to the x-y axes order.
130            scores (Tensor): Input scores with shape [batch, n_boxes, n_classes].
131
132        Returns: 'NMSResults' named tuple:
133            - boxes: The selected boxes with shape [batch, max_detections, 4].
134            - scores: The corresponding scores in descending order with shape [batch, max_detections].
135            - labels: The labels for each box with shape [batch, max_detections].
136            - n_valid: The number of valid detections out of 'max_detections' with shape [batch, 1]
137        """
138        nms = multiclass_nms(boxes=boxes,
139                             scores=scores,
140                             score_threshold=self.score_threshold,
141                             iou_threshold=self.iou_threshold,
142                             max_detections=self.max_detections)
143        return nms

102class MulticlassNMSWithIndices(CustomLayer):
103    """
104    A torch.nn.Module for multiclass NMS with indices. See multiclass_nms_with_indices for additional information.
105
106    Usage example:
107        batch size=1, 1000 boxes, 50 classes
108        boxes = torch.rand(1, 1000, 4)
109        scores = torch.rand(1, 1000, 50)
110        nms = MulticlassNMSWithIndices(score_threshold=0.1,
111                                       iou_threshold=0.6
112                                       max_detections=300)
113        res = nms(boxes, scores)
114    """
115
116    def __init__(self, score_threshold: float, iou_threshold: float, max_detections: int):
117        """
118        Args:
119            score_threshold (float): The score threshold. Candidates with scores below the threshold are discarded.
120            iou_threshold (float): The Intersection Over Union (IOU) threshold for boxes overlap.
121            max_detections (int): The number of detections to return.
122        """
123        super(MulticlassNMSWithIndices, self).__init__()
124        self.score_threshold = score_threshold
125        self.iou_threshold = iou_threshold
126        self.max_detections = max_detections
127
128    def forward(self, boxes: torch.Tensor, scores: torch.Tensor):
129        """
130        Args:
131            boxes (Tensor): Input boxes with shape [batch, n_boxes, 4], specified in corner coordinates
132                            (x_min, y_min, x_max, y_max). Agnostic to the x-y axes order.
133            scores (Tensor): Input scores with shape [batch, n_boxes, n_classes].
134
135        Returns: 'NMSWithIndicesResults' named tuple:
136            - boxes: The selected boxes with shape [batch, max_detections, 4].
137            - scores: The corresponding scores in descending order with shape [batch, max_detections].
138            - labels: The labels for each box with shape [batch, max_detections].
139            - indices: Indices of the input boxes that have been selected.
140            - n_valid: The number of valid detections out of 'max_detections' with shape [batch, 1]
141        """
142        nms = multiclass_nms_with_indices(boxes=boxes,
143                                          scores=scores,
144                                          score_threshold=self.score_threshold,
145                                          iou_threshold=self.iou_threshold,
146                                          max_detections=self.max_detections)
147        return nms

103class MulticlassNMSOBB(CustomLayer):
104    """
105    A torch.nn.Module for multiclass NMS for oriented bounding box. See multiclass_nms_obb for additional information.
106
107    Usage example:
108        batch size=1, 1000 boxes, 50 classes
109        boxes = torch.rand(1, 1000, 4)
110        scores = torch.rand(1, 1000, 50)
111        angles = torch.rand(1, 1000, 1)
112        nms_obb = MulticlassNMSOBB(score_threshold=0.1,
113                                   iou_threshold=0.6
114                                   max_detections=300)
115        res = nms_obb(boxes, scores, angles)
116    """
117
118    def __init__(self, score_threshold: float, iou_threshold: float, max_detections: int):
119        """
120        Args:
121            score_threshold (float): The score threshold. Candidates with scores below the threshold are discarded.
122            iou_threshold (float): The Intersection Over Union (IOU) threshold for boxes overlap.
123            max_detections (int): The number of detections to return.
124        """
125        super(MulticlassNMSOBB, self).__init__()
126        self.score_threshold = score_threshold
127        self.iou_threshold = iou_threshold
128        self.max_detections = max_detections
129
130    def forward(self, boxes: torch.Tensor, scores: torch.Tensor, angles: torch.Tensor):
131        """
132        Args:
133            boxes (Tensor): Input boxes with shape [batch, n_boxes, 4], specified in center coordinates,
134                            width and height (x_center, y_center, w, h).
135            scores (Tensor): Input scores with shape [batch, n_boxes, n_classes].
136            angles (Tensor): Input angles with shape [batch, n_boxes, 1].
137
138        Returns: 'NMSOBBResults' named tuple:
139            - boxes: The selected boxes with shape [batch, max_detections, 4].
140            - scores: The corresponding scores in descending order with shape [batch, max_detections].
141            - labels: The labels for each box with shape [batch, max_detections].
142            - angles: The corresponding angles for each box with shape [batch, max_detections].
143            - n_valid: The number of valid detections out of 'max_detections' with shape [batch, 1].
144        """
145        nms = multiclass_nms_obb(boxes=boxes,
146                                 scores=scores,
147                                 angles=angles,
148                                 score_threshold=self.score_threshold,
149                                 iou_threshold=self.iou_threshold,
150                                 max_detections=self.max_detections)
151        return nms

 54def multiclass_nms_obb(boxes, scores, angles, score_threshold: float, iou_threshold: float,
 55                       max_detections: int) -> NMSOBBResults:
 56    """
 57    Multi-class non-maximum suppression for oriented bounding box.
 58    Detections are returned in descending order of their scores.
 59    The output tensors always contain a fixed number of detections, as defined by 'max_detections'.
 60    If fewer detections are selected, the output tensors are zero-padded up to 'max_detections'.
 61
 62    Args:
 63        boxes (Tensor): Input boxes with shape [batch, n_boxes, 4], specified in center coordinates,
 64                        width and height (x_center, y_center, w, h).
 65        scores (Tensor): Input scores with shape [batch, n_boxes, n_classes].
 66        angles (Tensor): Input angles with shape [batch, n_boxes, 1].
 67        score_threshold (float): The score threshold. Candidates with scores below the threshold are discarded.
 68        iou_threshold (float): The Intersection Over Union (IOU) threshold for boxes overlap.
 69        max_detections (int): The number of detections to return.
 70
 71    Returns:
 72        'NMSOBBResults' named tuple:
 73        - boxes: The selected boxes with shape [batch, max_detections, 4].
 74        - scores: The corresponding scores in descending order with shape [batch, max_detections].
 75        - labels: The labels for each box with shape [batch, max_detections].
 76        - angles: The corresponding angles for each box with shape [batch, max_detections].
 77        - n_valid: The number of valid detections out of 'max_detections' with shape [batch, 1].
 78
 79    Raises:
 80        ValueError: If provided with invalid arguments or input tensors with unexpected or non-matching shapes.
 81
 82    Example:
 83        ```
 84        from edgemdt_cl.pytorch import multiclass_nms_obb
 85
 86        # batch size=1, 1000 boxes, 50 classes
 87        boxes = torch.rand(1, 1000, 4)
 88        scores = torch.rand(1, 1000, 50)
 89        angles = torch.rand(1, 1000, 1)
 90        res = multiclass_nms_obb(boxes,
 91                                 scores,
 92                                 angles,
 93                                 score_threshold=0.1,
 94                                 iou_threshold=0.6,
 95                                 max_detections=300)
 96        # res.boxes, res.scores, res.labels, res.angles, res.n_valid
 97        ```
 98    """
 99    return NMSOBBResults(
100        *torch.ops.edgemdt.multiclass_nms_obb(boxes, scores, angles, score_threshold, iou_threshold, max_detections))

31class NMSOBBResults(NamedTuple):
32    """ Container for non-maximum suppression for oriented bounding box results """
33    boxes: Tensor
34    scores: Tensor
35    labels: Tensor
36    angles: Tensor
37    n_valid: Tensor
38
39    # Note: convenience methods below are replicated in each Results container, since NamedTuple supports neither adding
40    # new fields in derived classes nor multiple inheritance, and we want it to behave like a tuple, so no dataclasses.
41    def detach(self) -> 'NMSOBBResults':
42        """ Detach all tensors and return a new object """
43        return self.apply(lambda t: t.detach())
44
45    def cpu(self) -> 'NMSOBBResults':
46        """ Move all tensors to cpu and return a new object """
47        return self.apply(lambda t: t.cpu())
48
49    def apply(self, f: Callable[[Tensor], Tensor]) -> 'NMSOBBResults':
50        """ Apply any function to all tensors and return a new object """
51        return self.__class__(*[f(t) for t in self])

22class CustomLayer(torch.nn.Module, ABC):
23    pass