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: