From 162faf83fb346e406b2d99d424bc82f2c6f6249d Mon Sep 17 00:00:00 2001 From: Lena Perry Date: Mon, 8 Dec 2025 09:52:16 -0500 Subject: [PATCH] Fix: map megadetector detection categories to labels (issue #357) --- zamba/images/manager.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/zamba/images/manager.py b/zamba/images/manager.py index 319fe4b7..73e6323c 100644 --- a/zamba/images/manager.py +++ b/zamba/images/manager.py @@ -74,7 +74,12 @@ def predict(config: ImageClassificationPredictConfig) -> None: for detection in results["detections"]: try: bbox = absolute_bbox(image, detection["bbox"], bbox_layout=BboxLayout.XYWH) - detection_category = detection["category"] + if detection["category"] == "1": + detection_category = "animal" + elif detection["category"] == "2": + detection_category = "person" + elif detection["category"] == "3": + detection_category = "vehicle" detection_conf = detection["conf"] img = image.crop(bbox) input_data = image_transforms(img)