Update app.py

app.py

@@ -25,7 +25,7 @@ MODEL = "./best.pt"
 
 TARGET_VIDEO_PATH = "test.mp4"
 
-CLASS_ID = [0,1,2,3]
+CLASS_ID = [0,1,2,3,4,5,6]
 
 video_examples = [['example.mp4']]
 
@@ -103,7 +103,7 @@ def ObjectDetection(video_path):
                 class_id=results[0].boxes.cls.cpu().numpy().astype(int)
             )
             # filtering out detections with unwanted classes
-            detections = detections[np.isin(detections.class_id, [0,1,2,3])]
+            detections = detections[np.isin(detections.class_id, CLASS_ID)]
             # tracking detections
             tracks = byte_tracker.update(
                 output_results=detections2boxes(detections = detections),