From 951b2f362a669e93eaecb345ec0108ab4a2ea712 Mon Sep 17 00:00:00 2001
From: luanhailiang <hi@luanhailiang.cn>
Date: Thu, 21 Sep 2023 20:38:58 +0800
Subject: [PATCH] 4

---
 app.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/app.py b/app.py
index 469ea00..c6f40b7 100644
--- a/app.py
+++ b/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: