luanhailiang
/
billiard
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

4

main
luanhailiang 2023-09-21 20:38:58 +08:00
parent c96c2a0b6f
commit 951b2f362a
1 changed files with 2 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
app.py
View File

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