Skip to content

Commit 22d82c3

Browse files
authored
Merge pull request #173 from hardikdava/mmdetection_support
Mmdetection support
2 parents 1b9c65b + 3f78150 commit 22d82c3

File tree

1 file changed

+29
-0
lines changed

1 file changed

+29
-0
lines changed

supervision/detection/core.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,35 @@ def from_yolo_nas(cls, yolo_nas_results) -> Detections:
228228
class_id=yolo_nas_results.prediction.labels.astype(int),
229229
)
230230

231+
@classmethod
232+
def from_mmdetection(cls, mmdet_results) -> Detections:
233+
"""
234+
Creates a Detections instance from a [mmdetection](https://github.com/open-mmlab/mmdetection) inference result.
235+
Also supported for [mmyolo](https://github.com/open-mmlab/mmyolo)
236+
237+
Args:
238+
mmdet_results (mmdet.structures.DetDataSample): The output Results instance from MMDetection
239+
240+
Returns:
241+
Detections: A new Detections object.
242+
243+
Example:
244+
```python
245+
>>> import cv2
246+
>>> import supervision as sv
247+
>>> from mmdet.apis import DetInferencer
248+
249+
>>> inferencer = DetInferencer(model_name, checkpoint, device)
250+
>>> mmdet_result = inferencer(SOURCE_IMAGE_PATH, out_dir='./output', return_datasample=True)["predictions"][0]
251+
>>> detections = sv.Detections.from_mmdet(mmdet_result)
252+
```
253+
"""
254+
return cls(
255+
xyxy=mmdet_results.pred_instances.bboxes.cpu().numpy(),
256+
confidence=mmdet_results.pred_instances.scores.cpu().numpy(),
257+
class_id=mmdet_results.pred_instances.labels.cpu().numpy().astype(int),
258+
)
259+
231260
@classmethod
232261
def from_transformers(cls, transformers_results: dict) -> Detections:
233262
"""

0 commit comments

Comments
 (0)