From f35a78ff1d2fb57561db4b347a2bde3926813788 Mon Sep 17 00:00:00 2001 From: sovit-123 Date: Wed, 28 Aug 2024 21:45:36 +0530 Subject: [PATCH] Updated SAHI inference for custom trained models --- sahi_inference.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/sahi_inference.py b/sahi_inference.py index f3542d7e..207eba72 100644 --- 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))