111 lines
3.2 KiB
Python
111 lines
3.2 KiB
Python
|
import cv2
|
||
|
|
import json
|
||
|
|
import time
|
||
|
|
import requests
|
||
|
|
import supervision as sv
|
||
|
|
|
||
|
|
from ultralytics import YOLO
|
||
|
|
from flask import Flask,request,Response
|
||
|
|
from multiprocessing import Process, Event
|
||
|
|
|
||
|
|
app = Flask(__name__)
|
||
|
|
|
||
|
|
workers = {}
|
||
|
|
|
||
|
|
@app.route('/start/<id>')
|
||
|
|
def start(id):
|
||
|
|
table = workers.get(id)
|
||
|
|
if table is not None:
|
||
|
|
table.set()
|
||
|
|
source = request.args.get('source')
|
||
|
|
target = request.args.get('target')
|
||
|
|
region = request.args.get('region')
|
||
|
|
if target is None:
|
||
|
|
return "please input target url"
|
||
|
|
return Response(worker(id,source,region,True), mimetype='text/event-stream')
|
||
|
|
|
||
|
|
stop_event = Event()
|
||
|
|
p = Process(target=workerloop, args=(stop_event,id,source,target,region))
|
||
|
|
p.start()
|
||
|
|
workers[id] = stop_event
|
||
|
|
return 'ok'
|
||
|
|
|
||
|
|
@app.route('/stop/<id>')
|
||
|
|
def stop(id):
|
||
|
|
table = workers.get(id)
|
||
|
|
if table is not None:
|
||
|
|
table.set()
|
||
|
|
return 'ok'
|
||
|
|
|
||
|
|
@app.route('/show',methods=['POST'])
|
||
|
|
def show():
|
||
|
|
print(request.data)
|
||
|
|
return 'ok'
|
||
|
|
|
||
|
|
def worker(id,source,region=None,stream=False):
|
||
|
|
if region == None:
|
||
|
|
region = 1000
|
||
|
|
else:
|
||
|
|
region = int(region)
|
||
|
|
model = YOLO('./best.pt')
|
||
|
|
balls = {}
|
||
|
|
timebegin = int(time.time()*1000)
|
||
|
|
print(source)
|
||
|
|
try:
|
||
|
|
cap = cv2.VideoCapture(source)
|
||
|
|
if not cap.isOpened():
|
||
|
|
print("Error opening video stream.")
|
||
|
|
while True:
|
||
|
|
ret, frame = cap.read()
|
||
|
|
if not ret:
|
||
|
|
print("Error read video stream.")
|
||
|
|
break
|
||
|
|
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:
|
||
|
|
detections.tracker_id = result.boxes.id.cpu().numpy().astype(int)
|
||
|
|
detections = detections[(detections.tracker_id != None)]
|
||
|
|
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):
|
||
|
|
continue
|
||
|
|
json_data = json.dumps({"table":id,"balls":balls,"time":timebegin})
|
||
|
|
timebegin += region
|
||
|
|
balls = {}
|
||
|
|
if not stream:
|
||
|
|
yield json_data
|
||
|
|
else:
|
||
|
|
yield f"data: {json_data}\n\n"
|
||
|
|
except GeneratorExit:
|
||
|
|
print("Client disconnected at", time.ctime())
|
||
|
|
|
||
|
|
finally:cap.release()
|
||
|
|
|
||
|
|
def workerloop(stop_event,id,source,target=None,region=None):
|
||
|
|
try:
|
||
|
|
gen = worker(id,source,region)
|
||
|
|
for data in gen:
|
||
|
|
if stop_event.is_set():
|
||
|
|
break
|
||
|
|
json_data = json.loads(data)
|
||
|
|
# 设置请求头
|
||
|
|
headers = {
|
||
|
|
"Content-Type": "application/json"
|
||
|
|
}
|
||
|
|
text = json.dumps({"content":json_data})
|
||
|
|
print(text)
|
||
|
|
response = requests.post(target, data=text, headers=headers)
|
||
|
|
|
||
|
|
print(response.text)
|
||
|
|
finally:
|
||
|
|
gen.close()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
app.run("0.0.0.0",threaded=True)
|