diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..f8ff2b5 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +*.mp4 filter=lfs diff=lfs merge=lfs -text diff --git a/.gitea/workflows/demo.yaml b/.gitea/workflows/demo.yaml new file mode 100644 index 0000000..c0a79ee --- /dev/null +++ b/.gitea/workflows/demo.yaml @@ -0,0 +1,16 @@ +name: Gitea Actions Demo +run-name: ${{ gitea.actor }} is testing out Gitea Actions 🚀 +on: [push] + +jobs: + Explore-Gitea-Actions: + runs-on: ubuntu-latest + steps: + - name: executing remote ssh commands using password + uses: appleboy/ssh-action@v1.0.0 + with: + host: ${{ secrets.HOST }} + username: ${{ secrets.USERNAME }} + key: ${{ secrets.KEY }} + port: ${{ secrets.PORT }} + script: cd /root/billiard && git pull && docker compose build && docker compose up -d && docker compose restart \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..c5bb83f --- /dev/null +++ b/Dockerfile @@ -0,0 +1,14 @@ +FROM ultralytics/ultralytics:latest + + +RUN pip3 install -U pip +RUN pip3 install supervision +RUN pip3 install flask +RUN pip3 install gunicorn +RUN pip3 install lapx>=0.5.2 +RUN pip3 install -v -e . + +WORKDIR /app + +# 在容器启动时运行 Gunicorn 以运行 Flask 应用 +CMD ["gunicorn", "app:app", "--bind", "0.0.0.0:5000", "--timeout", "0", "--workers", "1"] \ No newline at end of file diff --git a/app.py b/app.py new file mode 100644 index 0000000..a58d6f5 --- /dev/null +++ b/app.py @@ -0,0 +1,110 @@ +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/') +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/') +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) diff --git a/best.pt b/best.pt new file mode 100644 index 0000000..ed390d7 Binary files /dev/null and b/best.pt differ diff --git a/docker-compose.yaml b/docker-compose.yaml new file mode 100644 index 0000000..7d39345 --- /dev/null +++ b/docker-compose.yaml @@ -0,0 +1,19 @@ +version: '3' +services: + billiard: + # command: ["python", "./app.py"] + command: ["gunicorn", "app:app", "--bind", "0.0.0.0:5000", "--timeout", "0", "--workers", "1"] + restart: always + build: . + volumes: + - "./:/app" + ports: + - "8500:5000" + runtime: nvidia + environment: + NVIDIA_VISIBLE_DEVICES: all + logging: + driver: "json-file" + options: + max-size: "100m" + max-file: "3" \ No newline at end of file diff --git a/test.py b/test.py new file mode 100644 index 0000000..77920ff --- /dev/null +++ b/test.py @@ -0,0 +1,5 @@ +from ultralytics import YOLO + +# Configure the tracking parameters and run the tracker +model = YOLO('best.pt') +results = model.track(source="./videos/123.mp4", conf=0.3, iou=0.5, show=True) diff --git a/videos/123.mp4 b/videos/123.mp4 new file mode 100644 index 0000000..619a13d --- /dev/null +++ b/videos/123.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:50f1d450d1e975b5a16160021777e2daf928db5ec5a90c9d7691a9a7cb3c7ad3 +size 15405787