Update app.py

app.py

@@ -33,7 +33,7 @@ def yolov9_inference(model_id, img_path=None, vid_path=None, tracking_algorithm
         img.save(img_path)
         input_path = img_path
 
-        output_path, df, frame_counts_df = run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='cpu', hide_conf= True)
+        output_path, df, frame_counts_df = run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='0', hide_conf= True)
     elif vid_path is not None:
         vid_name = 'output.mp4'
 
@@ -70,12 +70,12 @@ def yolov9_inference(model_id, img_path=None, vid_path=None, tracking_algorithm
         out.release()
         input_path = vid_name
         if tracking_algorithm == 'deep_sort':
-            output_path, df, frame_counts_df = run_deepsort(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='cpu', draw_trails=True)
+            output_path, df, frame_counts_df = run_deepsort(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='0', draw_trails=True)
         elif tracking_algorithm == 'strong_sort':
-            device_strongsort = torch.device('cpu')
+            device_strongsort = torch.device('cuda:0')
             output_path, df, frame_counts_df = run_strongsort(yolo_weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device=device_strongsort, strong_sort_weights = "osnet_x0_25_msmt17.pt", hide_conf= True)
         else: 
-            output_path, df, frame_counts_df =  run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='cpu', hide_conf= True)
+            output_path, df, frame_counts_df =  run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='0', hide_conf= True)
         # Assuming output_path is the path to the output file
     _, output_extension = os.path.splitext(output_path)
     palette = {"Bus": "red", "Bike": "blue", "Car": "green", "Pedestrian": "yellow", "Truck": "purple"}