billiard/app.py

import cv2
import json
import time
import requests
import threading
import numpy as np
import supervision as sv

from datetime import datetime
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.json)
    data = request.json["content"]
    info(data)
    return 'ok'

def info(data):
    balls = data["balls"]
    table = data["table"]
    text = f"table:{table:>3} "
    for i in range(16):
        ball = "#" if "ball"+str(i) in balls else " "
        text += f"| {i} : {ball} "
    print(datetime.now().strftime('%Y-%m-%d %H:%M:%S'),text)

def worker(id,source,region=None,stream=False):
    if region == None:
        region = 10
    else:
        region = int(region)
    print("start loading model...",id,source,region)
    model = YOLO('./best.pt')
    print("start loaded  model!!!")
    balls = {}
    count = 0
    try:
        cap = cv2.VideoCapture(source)
        if not cap.isOpened():
            print("Error opening video stream.")
        while True:
            use = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            bgn = (time.time()*1000)
            ret, frame = cap.read()
            del(ret)
            del(frame)
            ret, frame = cap.read()
            use += " read:"+str((time.time()*1000) - bgn)
            if not ret:
                print("Error read video stream.")
                cap = cv2.VideoCapture(source)
                if not cap.isOpened():
                    print("Error opening video stream..")
                    cap = cv2.VideoCapture(source)
                    time.sleep(1)
                    if not cap.isOpened():
                        print("Error opening video stream...")
                        time.sleep(3)
                        if not cap.isOpened():
                            print("Error opening video stream....")
                        break
                continue
            result = model.track(frame,show=False,stream=False,persist=True,device=int(id)%4)
            use += " track:"+str((time.time()*1000) - bgn)
            del(ret)
            del(frame)
            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)
            else:
                detections.tracker_id = np.array([])
                detections.conf = np.array([])
                detections.xyxy=np.empty((0, 4), dtype=np.float32)
            # detections = detections[(detections.tracker_id != None)]
            
            count += 1
            names = {}
            for xyxy,_, confidence, class_id, tracker_id in detections:
                name = model.model.names[class_id]
                if name in names:
                    print("Name duplicate",name)
                    continue
                ball =  balls.get(name,{})
                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[name] = ball
                names[name] = True
            names = {}
            use += " names:"+str((time.time()*1000) - bgn)
            if count < region:
                continue
            data = {"table":id,"balls":balls,"time":int(time.time()*1000)}
            info(data)
            json_data = json.dumps(data)
            balls = {}
            count = 0
            use += " dump:"+str((time.time()*1000) - bgn)
            if not stream:
                yield json_data
            else:
                yield f"data: {json_data}\n\n"
            use += " yield:"+str((time.time()*1000) - bgn)
            print("model.track",use)
    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:
            use = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            bgn = (time.time()*1000)
            if stop_event.is_set():
                break
            json_data = json.loads(data)
            use += " loads:"+str((time.time()*1000) - bgn)
            text = json.dumps({"content":json_data})
            use += " dumps:"+str((time.time()*1000) - bgn)

            # 创建并启动线程
            thread = threading.Thread(target=post_request, args=(target, text))
            thread.start()

            use += " post:"+str((time.time()*1000) - bgn)
            print("workerloop",use)
    finally:
        gen.close()

def post_request(url, data):
    use = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    bgn = (time.time()*1000)
    headers = {
        "Content-Type": "application/json"
    }
    response = requests.post(url, data=data,headers=headers)
    use += " post:"+str((time.time()*1000) - bgn ) +"ms"
    print("post_request",use,response.text)

# import pstats
# import cProfile
# def test_function():
#     url = "rtmp://rtmp03open.ys7.com:1935/v3/openlive/L05874022_1_1?expire=1726389273&id=625006970167889920&t=2c7f371063e4ad672cb7a7c34b2a236e9ff37b02052f818c9c4fe7ae9632cd0e&ev=100"
#     gen = worker(0,url,5)
#     count = 0
#     for data in gen:
#         count += 1
#         if count > 10:
#             break
        
if __name__ == '__main__':
    # cProfile.run('test_function()', 'test_function.profile')
    # p = pstats.Stats('test_function.profile')
    # p.sort_stats('cumulative').print_stats(100)  # Top 10 by cumulative time
    app.run("0.0.0.0",threaded=True)