diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..eeabc60 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +yolo \ No newline at end of file diff --git a/app.py b/app.py index a58d6f5..9c31f83 100644 --- a/app.py +++ b/app.py @@ -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: diff --git a/venv.bat b/venv.bat new file mode 100644 index 0000000..3771f56 --- /dev/null +++ b/venv.bat @@ -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 \ No newline at end of file