fix yolov5 category names (#297)

obss · Dec 6, 2021 · 830638e · 830638e
1 parent 7a2bf29
commit 830638e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sahi/model.py b/sahi/model.py
@@ -483,7 +483,7 @@ def _create_object_prediction_list_from_original_predictions(
             bbox = [x1, y1, x2, y2]
             score = prediction[4]
             category_id = int(prediction[5])
-            category_name = original_predictions.names[category_id]
+            category_name = self.category_mapping[str(category_id)]
 
             # ignore invalid predictions
             if bbox[0] > bbox[2] or bbox[1] > bbox[3] or bbox[0] < 0 or bbox[1] < 0 or bbox[2] < 0 or bbox[3] < 0: