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)

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

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