4

app.py

@@ -9,8 +9,6 @@ from ultralytics import YOLO
 from flask import Flask,request,Response
 from multiprocessing import Process, Event
 
-import torch
-
 app = Flask(__name__)
 
 workers = {}
@@ -61,8 +59,7 @@ def worker(id,source,region=None,stream=False):
     model = YOLO('./best.pt')
     balls = {}
     count = 0
-    device_count = torch.cuda.device_count()
-    print(device_count,source)
+    print(source)
     try:
         cap = cv2.VideoCapture(source)
         if not cap.isOpened():
@@ -72,7 +69,7 @@ def worker(id,source,region=None,stream=False):
             if not ret:
                 print("Error read video stream.")
                 break
-            result = model.track(frame,show=False,stream=False,persist=True,device=int(id)%device_count)
+            result = model.track(frame,show=False,stream=False,persist=True,device=int(id)%4)
             result = result[0]
             detections = sv.Detections.from_yolov8(result)
             if result.boxes.id is not None: