Updated SAHI inference for custom trained models

sovit-123 · Aug 28, 2024 · f35a78f · f35a78f
1 parent 2a26ce6
commit f35a78f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/sahi_inference.py b/sahi_inference.py
@@ -220,7 +220,9 @@ def main(args):
         model_type='torchvision',
         model=model,
         confidence_threshold=args['threshold'],
-        device=args['device']
+        device=args['device'],
+        category_mapping={str(i): CLASSES[i] for i in range(1, len(CLASSES))},
+        # category_remapping={CLASSES[i]: i for i in range(1, len(CLASSES))}
     )
 
     COLORS = np.random.uniform(0, 255, size=(len(CLASSES), 3))