main
luanhailiang 2023-08-14 16:59:16 +08:00
parent f7ba9dae22
commit af4ed177f6
3 changed files with 29 additions and 11 deletions

1
.gitignore vendored 100644
View File

@ -0,0 +1 @@
yolo

27
app.py
View File

@ -44,12 +44,12 @@ def show():
def worker(id,source,region=None,stream=False):
if region == None:
region = 1000
region = 10
else:
region = int(region)
model = YOLO('./best.pt')
balls = {}
timebegin = int(time.time()*1000)
count = 0
print(source)
try:
cap = cv2.VideoCapture(source)
@ -60,23 +60,28 @@ 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)%4)
result = model.track(frame,show=False,stream=False,persist=True)
result = result[0]
detections = sv.Detections.from_yolov8(result)
if result.boxes.id is not None:
detections.tracker_id = result.boxes.id.cpu().numpy().astype(int)
detections = detections[(detections.tracker_id != None)]
count += 1
for xyxy,_, confidence, class_id, tracker_id in detections:
balls[model.model.names[class_id]] = {
"tkid": int(tracker_id),
"conf": round(float(confidence), 2),
"xyxy": [int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3]) ]
}
if timebegin > int(time.time()*1000):
ball = balls.get(model.model.names[class_id],{})
ball["tkid"] = int(tracker_id)
ball["conf"] = round(float(confidence), 2)
ball["xyxy"] = [int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])]
ball["count"] = ball.get("count",0) + 1
if region - count < 3:
ball["final"] = ball.get("final",0) + 1
balls[model.model.names[class_id]] = ball
if count < region:
continue
json_data = json.dumps({"table":id,"balls":balls,"time":timebegin})
timebegin += region
json_data = json.dumps({"table":id,"balls":balls,"time":int(time.time()*1000)})
balls = {}
count = 0
if not stream:
yield json_data
else:

12
venv.bat 100644
View File

@ -0,0 +1,12 @@
python -m venv yolo
./yolo/Scripts/Activate.ps1
python.exe -m pip install --upgrade pip
pip3 install ultralytics
pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
pip3 install supervision
pip3 install flask
pip3 install lapx>=0.5.2
python .\app.py