tarto2 commited on
Commit
c1e6f92
·
verified ·
1 Parent(s): 4a97138

Upload folder using huggingface_hub

Browse files
.gitattributes CHANGED
@@ -1,35 +1,38 @@
1
- *.7z filter=lfs diff=lfs merge=lfs -text
2
- *.arrow filter=lfs diff=lfs merge=lfs -text
3
- *.bin filter=lfs diff=lfs merge=lfs -text
4
- *.bz2 filter=lfs diff=lfs merge=lfs -text
5
- *.ckpt filter=lfs diff=lfs merge=lfs -text
6
- *.ftz filter=lfs diff=lfs merge=lfs -text
7
- *.gz filter=lfs diff=lfs merge=lfs -text
8
- *.h5 filter=lfs diff=lfs merge=lfs -text
9
- *.joblib filter=lfs diff=lfs merge=lfs -text
10
- *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
- *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
- *.model filter=lfs diff=lfs merge=lfs -text
13
- *.msgpack filter=lfs diff=lfs merge=lfs -text
14
- *.npy filter=lfs diff=lfs merge=lfs -text
15
- *.npz filter=lfs diff=lfs merge=lfs -text
16
- *.onnx filter=lfs diff=lfs merge=lfs -text
17
- *.ot filter=lfs diff=lfs merge=lfs -text
18
- *.parquet filter=lfs diff=lfs merge=lfs -text
19
- *.pb filter=lfs diff=lfs merge=lfs -text
20
- *.pickle filter=lfs diff=lfs merge=lfs -text
21
- *.pkl filter=lfs diff=lfs merge=lfs -text
22
- *.pt filter=lfs diff=lfs merge=lfs -text
23
- *.pth filter=lfs diff=lfs merge=lfs -text
24
- *.rar filter=lfs diff=lfs merge=lfs -text
25
- *.safetensors filter=lfs diff=lfs merge=lfs -text
26
- saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
- *.tar.* filter=lfs diff=lfs merge=lfs -text
28
- *.tar filter=lfs diff=lfs merge=lfs -text
29
- *.tflite filter=lfs diff=lfs merge=lfs -text
30
- *.tgz filter=lfs diff=lfs merge=lfs -text
31
- *.wasm filter=lfs diff=lfs merge=lfs -text
32
- *.xz filter=lfs diff=lfs merge=lfs -text
33
- *.zip filter=lfs diff=lfs merge=lfs -text
34
- *.zst filter=lfs diff=lfs merge=lfs -text
35
- *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ SV_kp.engine filter=lfs diff=lfs merge=lfs -text
37
+ osnet_model.pth.tar-100 filter=lfs diff=lfs merge=lfs -text
38
+ keypoint filter=lfs diff=lfs merge=lfs -text
.gitignore ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ venv
2
+ outputs
3
+ test_predict_batch.py
4
+ *.mp4
5
+ inspect_yolo_model.py
20251029-detection.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8bbacfcb38e38b1b8816788e9e6e845160533719a0b87b693d58b932380d0d28
3
+ size 152961687
20251029-keypoint.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6dd10dba85895c92760cdb5a99c5cfca899c68f361a66c5448f38a187280ee1f
3
+ size 6849672
README.md ADDED
@@ -0,0 +1,132 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ 🚀 Example Chute for Turbovision 🪂
2
+
3
+ This repository demonstrates how to deploy a Chute via the Turbovision CLI, hosted on Hugging Face Hub. It serves as a minimal example showcasing the required structure and workflow for integrating machine learning models, preprocessing, and orchestration into a reproducible Chute environment.
4
+
5
+ ## Repository Structure
6
+
7
+ The following two files must be present (in their current locations) for a successful deployment — their content can be modified as needed:
8
+
9
+ | File | Purpose |
10
+ |------|---------|
11
+ | `miner.py` | Defines the ML model type(s), orchestration, and all pre/postprocessing logic. |
12
+ | `config.yml` | Specifies machine configuration (e.g., GPU type, memory, environment variables). |
13
+
14
+ Other files — e.g., model weights, utility scripts, or dependencies — are optional and can be included as needed for your model.
15
+
16
+ > **Note**: Any required assets must be defined or contained within this repo, which is fully open-source, since all network-related operations (downloading challenge data, weights, etc.) are disabled inside the Chute.
17
+
18
+ ## Overview
19
+
20
+ Below is a high-level diagram showing the interaction between Huggingface, Chutes and Turbovision:
21
+
22
+ ```
23
+ ┌─────────────┐ ┌──────────┐ ┌──────────────┐
24
+ │ HuggingFace │ ───> │ Chutes │ ───> │ Turbovision │
25
+ │ Hub │ │ .ai │ │ Validator │
26
+ └─────────────┘ └──────────┘ └──────────────┘
27
+ ```
28
+
29
+ ## Local Testing
30
+
31
+ After editing the `config.yml` and `miner.py` and saving it into your Huggingface Repo, you will want to test it works locally.
32
+
33
+ 1. **Copy the template file** `scorevision/chute_template/turbovision_chute.py.j2` as a python file called `my_chute.py` and fill in the missing variables:
34
+
35
+ ```python
36
+ HF_REPO_NAME = "{{ huggingface_repository_name }}"
37
+ HF_REPO_REVISION = "{{ huggingface_repository_revision }}"
38
+ CHUTES_USERNAME = "{{ chute_username }}"
39
+ CHUTE_NAME = "{{ chute_name }}"
40
+ ```
41
+
42
+ 2. **Run the following command to build the chute locally** (Caution: there are known issues with the docker location when running this on a mac):
43
+
44
+ ```bash
45
+ chutes build my_chute:chute --local --public
46
+ ```
47
+
48
+ 3. **Run the name of the docker image just built** (i.e. `CHUTE_NAME`) and enter it:
49
+
50
+ ```bash
51
+ docker run -p 8000:8000 -e CHUTES_EXECUTION_CONTEXT=REMOTE -it <image-name> /bin/bash
52
+ ```
53
+
54
+ 4. **Run the file from within the container**:
55
+
56
+ ```bash
57
+ chutes run my_chute:chute --dev --debug
58
+ ```
59
+
60
+ 5. **In another terminal, test the local endpoints** to ensure there are no bugs:
61
+
62
+ ```bash
63
+ # Health check
64
+ curl -X POST http://localhost:8000/health -d '{}'
65
+
66
+ # Prediction test
67
+ curl -X POST http://localhost:8000/predict -d '{"url": "https://scoredata.me/2025_03_14/35ae7a/h1_0f2ca0.mp4","meta": {}}'
68
+ ```
69
+
70
+ ## Live Testing
71
+
72
+ If you have any chute with the same name (i.e. from a previous deployment), ensure you delete that first (or you will get an error when trying to build).
73
+
74
+ 1. **List existing chutes**:
75
+
76
+ ```bash
77
+ chutes chutes list
78
+ ```
79
+
80
+ Take note of the chute id that you wish to delete (if any):
81
+
82
+ ```bash
83
+ chutes chutes delete <chute-id>
84
+ ```
85
+
86
+ 2. **You should also delete its associated image**:
87
+
88
+ ```bash
89
+ chutes images list
90
+ ```
91
+
92
+ Take note of the chute image id:
93
+
94
+ ```bash
95
+ chutes images delete <chute-image-id>
96
+ ```
97
+
98
+ 3. **Use Turbovision's CLI to build, deploy and commit on-chain**:
99
+
100
+ ```bash
101
+ sv -vv push
102
+ ```
103
+
104
+ > **Note**: You can skip the on-chain commit using `--no-commit`. You can also specify a past huggingface revision to point to using `--revision` and/or the local files you want to upload to your huggingface repo using `--model-path`.
105
+
106
+ 4. **When completed, warm up the chute** (if its cold 🧊):
107
+
108
+ You can confirm its status using `chutes chutes list` or `chutes chutes get <chute-id>` if you already know its id.
109
+
110
+ > **Note**: Warming up can sometimes take a while but if the chute runs without errors (should be if you've tested locally first) and there are sufficient nodes (i.e. machines) available matching the `config.yml` you specified, the chute should become hot 🔥!
111
+
112
+ ```bash
113
+ chutes warmup <chute-id>
114
+ ```
115
+
116
+ 5. **Test the chute's endpoints**:
117
+
118
+ ```bash
119
+ # Health check
120
+ curl -X POST https://<YOUR-CHUTE-SLUG>.chutes.ai/health -d '{}' -H "Authorization: Bearer $CHUTES_API_KEY"
121
+
122
+ # Prediction
123
+ curl -X POST https://<YOUR-CHUTE-SLUG>.chutes.ai/predict -d '{"url": "https://scoredata.me/2025_03_14/35ae7a/h1_0f2ca0.mp4","meta": {}}' -H "Authorization: Bearer $CHUTES_API_KEY"
124
+ ```
125
+
126
+ 6. **Test what your chute would get on a validator**:
127
+
128
+ This also applies any validation/integrity checks which may fail if you did not use the Turbovision CLI above to deploy the chute:
129
+
130
+ ```bash
131
+ sv -vv run-once
132
+ ```
SV_kp.engine ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f99452eb79e064189e2758abd20a78845a5b639fc8b9c4bc650519c83e13e8db
3
+ size 368289641
config.yml ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Image:
2
+ from_base: parachutes/python:3.12
3
+ run_command:
4
+ - pip install --upgrade setuptools wheel
5
+ - pip install ultralytics==8.3.222 opencv-python-headless numpy pydantic
6
+ - pip install scikit-learn
7
+ - pip install onnxruntime-gpu
8
+ set_workdir: /app
9
+
10
+ NodeSelector:
11
+ gpu_count: 1
12
+ min_vram_gb_per_gpu: 16
13
+ exclude:
14
+ - "5090"
15
+ - b200
16
+ - h200
17
+ - mi300x
18
+
19
+ Chute:
20
+ timeout_seconds: 900
21
+ concurrency: 4
22
+ max_instances: 5
23
+ scaling_threshold: 0.5
24
+ shutdown_after_seconds: 3600
detection.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7b51470cb703f5a9a789df38674b67d4bbe7f8f31846d69dbc97ce484f790cf9
3
+ size 10245169
detection.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2ad3e89b658d2626c34174f6799d240ffd37cfe45752c0ce6ef73b05935042e0
3
+ size 52014742
football_object_detection.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8bbacfcb38e38b1b8816788e9e6e845160533719a0b87b693d58b932380d0d28
3
+ size 152961687
football_pitch_template.png ADDED
hrnetv2_w48.yaml ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ MODEL:
2
+ IMAGE_SIZE: [960, 540]
3
+ NUM_JOINTS: 58
4
+ PRETRAIN: ''
5
+ EXTRA:
6
+ FINAL_CONV_KERNEL: 1
7
+ STAGE1:
8
+ NUM_MODULES: 1
9
+ NUM_BRANCHES: 1
10
+ BLOCK: BOTTLENECK
11
+ NUM_BLOCKS: [4]
12
+ NUM_CHANNELS: [64]
13
+ FUSE_METHOD: SUM
14
+ STAGE2:
15
+ NUM_MODULES: 1
16
+ NUM_BRANCHES: 2
17
+ BLOCK: BASIC
18
+ NUM_BLOCKS: [4, 4]
19
+ NUM_CHANNELS: [48, 96]
20
+ FUSE_METHOD: SUM
21
+ STAGE3:
22
+ NUM_MODULES: 4
23
+ NUM_BRANCHES: 3
24
+ BLOCK: BASIC
25
+ NUM_BLOCKS: [4, 4, 4]
26
+ NUM_CHANNELS: [48, 96, 192]
27
+ FUSE_METHOD: SUM
28
+ STAGE4:
29
+ NUM_MODULES: 3
30
+ NUM_BRANCHES: 4
31
+ BLOCK: BASIC
32
+ NUM_BLOCKS: [4, 4, 4, 4]
33
+ NUM_CHANNELS: [48, 96, 192, 384]
34
+ FUSE_METHOD: SUM
35
+
keypoint ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7ea78fa76aaf94976a8eca428d6e3c59697a93430cba1a4603e20284b61f5113
3
+ size 264964645
keypoint.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6dd10dba85895c92760cdb5a99c5cfca899c68f361a66c5448f38a187280ee1f
3
+ size 6849672
keypoint_evaluation.py ADDED
@@ -0,0 +1,321 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Keypoint evaluation module for football pitch keypoint detection.
3
+ Evaluates keypoint accuracy by comparing projected template lines with detected lines.
4
+ """
5
+ import logging
6
+ from typing import List, Tuple, Optional
7
+ from pathlib import Path
8
+ import numpy as np
9
+ from numpy import ndarray, array, float32, uint8
10
+
11
+ import cv2
12
+
13
+ # Import cv2 functions
14
+ bitwise_and = cv2.bitwise_and
15
+ findHomography = cv2.findHomography
16
+ warpPerspective = cv2.warpPerspective
17
+ cvtColor = cv2.cvtColor
18
+ COLOR_BGR2GRAY = cv2.COLOR_BGR2GRAY
19
+ threshold = cv2.threshold
20
+ THRESH_BINARY = cv2.THRESH_BINARY
21
+ getStructuringElement = cv2.getStructuringElement
22
+ MORPH_RECT = cv2.MORPH_RECT
23
+ MORPH_TOPHAT = cv2.MORPH_TOPHAT
24
+ GaussianBlur = cv2.GaussianBlur
25
+ morphologyEx = cv2.morphologyEx
26
+ Canny = cv2.Canny
27
+ connectedComponents = cv2.connectedComponents
28
+ perspectiveTransform = cv2.perspectiveTransform
29
+ RETR_EXTERNAL = cv2.RETR_EXTERNAL
30
+ CHAIN_APPROX_SIMPLE = cv2.CHAIN_APPROX_SIMPLE
31
+ findContours = cv2.findContours
32
+ boundingRect = cv2.boundingRect
33
+ dilate = cv2.dilate
34
+
35
+ logger = logging.getLogger(__name__)
36
+
37
+ # Template keypoints constant - define your keypoints here
38
+ # Format: List of (x, y) tuples representing keypoint coordinates on the template image
39
+ TEMPLATE_KEYPOINTS: list[tuple[int, int]] = [
40
+ (5, 5), # 1
41
+ (5, 140), # 2
42
+ (5, 250), # 3
43
+ (5, 430), # 4
44
+ (5, 540), # 5
45
+ (5, 675), # 6
46
+ # -------------
47
+ (55, 250), # 7
48
+ (55, 430), # 8
49
+ # -------------
50
+ (110, 340), # 9
51
+ # -------------
52
+ (165, 140), # 10
53
+ (165, 270), # 11
54
+ (165, 410), # 12
55
+ (165, 540), # 13
56
+ # -------------
57
+ (527, 5), # 14
58
+ (527, 253), # 15
59
+ (527, 433), # 16
60
+ (527, 675), # 17
61
+ # -------------
62
+ (888, 140), # 18
63
+ (888, 270), # 19
64
+ (888, 410), # 20
65
+ (888, 540), # 21
66
+ # -------------
67
+ (940, 340), # 22
68
+ # -------------
69
+ (998, 250), # 23
70
+ (998, 430), # 24
71
+ # -------------
72
+ (1045, 5), # 25
73
+ (1045, 140), # 26
74
+ (1045, 250), # 27
75
+ (1045, 430), # 28
76
+ (1045, 540), # 29
77
+ (1045, 675), # 30
78
+ # -------------
79
+ (435, 340), # 31
80
+ (615, 340), # 32
81
+ ]
82
+
83
+ INDEX_KEYPOINT_CORNER_BOTTOM_LEFT = 5
84
+ INDEX_KEYPOINT_CORNER_BOTTOM_RIGHT = 29
85
+ INDEX_KEYPOINT_CORNER_TOP_LEFT = 0
86
+ INDEX_KEYPOINT_CORNER_TOP_RIGHT = 24
87
+
88
+
89
+ class InvalidMask(Exception):
90
+ """Exception raised when mask validation fails."""
91
+ pass
92
+
93
+
94
+ def has_a_wide_line(mask: ndarray, max_aspect_ratio: float = 1.0) -> bool:
95
+ contours, _ = findContours(mask, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE)
96
+ for cnt in contours:
97
+ x, y, w, h = boundingRect(cnt)
98
+ aspect_ratio = min(w, h) / max(w, h)
99
+ if aspect_ratio >= max_aspect_ratio:
100
+ return True
101
+ return False
102
+
103
+
104
+ def is_bowtie(points: ndarray) -> bool:
105
+ def segments_intersect(p1: int, p2: int, q1: int, q2: int) -> bool:
106
+ def ccw(a: int, b: int, c: int):
107
+ return (c[1] - a[1]) * (b[0] - a[0]) > (b[1] - a[1]) * (c[0] - a[0])
108
+
109
+ return (ccw(p1, q1, q2) != ccw(p2, q1, q2)) and (
110
+ ccw(p1, p2, q1) != ccw(p1, p2, q2)
111
+ )
112
+
113
+ pts = points.reshape(-1, 2)
114
+ edges = [(pts[0], pts[1]), (pts[1], pts[2]), (pts[2], pts[3]), (pts[3], pts[0])]
115
+ return segments_intersect(*edges[0], *edges[2]) or segments_intersect(
116
+ *edges[1], *edges[3]
117
+ )
118
+
119
+ def validate_mask_lines(mask: ndarray) -> None:
120
+ if mask.sum() == 0:
121
+ raise InvalidMask("No projected lines")
122
+ if mask.sum() == mask.size:
123
+ raise InvalidMask("Projected lines cover the entire image surface")
124
+ if has_a_wide_line(mask=mask):
125
+ raise InvalidMask("A projected line is too wide")
126
+
127
+
128
+ def validate_mask_ground(mask: ndarray) -> None:
129
+ num_labels, _ = connectedComponents(mask)
130
+ num_distinct_regions = num_labels - 1
131
+ if num_distinct_regions > 1:
132
+ raise InvalidMask(
133
+ f"Projected ground should be a single object, detected {num_distinct_regions}"
134
+ )
135
+ area_covered = mask.sum() / mask.size
136
+ if area_covered >= 0.9:
137
+ raise InvalidMask(
138
+ f"Projected ground covers more than {area_covered:.2f}% of the image surface which is unrealistic"
139
+ )
140
+
141
+
142
+ def validate_projected_corners(
143
+ source_keypoints: list[tuple[int, int]], homography_matrix: ndarray
144
+ ) -> None:
145
+ src_corners = array(
146
+ [
147
+ source_keypoints[INDEX_KEYPOINT_CORNER_BOTTOM_LEFT],
148
+ source_keypoints[INDEX_KEYPOINT_CORNER_BOTTOM_RIGHT],
149
+ source_keypoints[INDEX_KEYPOINT_CORNER_TOP_RIGHT],
150
+ source_keypoints[INDEX_KEYPOINT_CORNER_TOP_LEFT],
151
+ ],
152
+ dtype="float32",
153
+ )[None, :, :]
154
+
155
+ warped_corners = perspectiveTransform(src_corners, homography_matrix)[0]
156
+
157
+ if is_bowtie(warped_corners):
158
+ raise InvalidMask("Projection twisted!")
159
+
160
+
161
+ def project_image_using_keypoints(
162
+ image: ndarray,
163
+ source_keypoints: List[Tuple[int, int]],
164
+ destination_keypoints: List[Tuple[int, int]],
165
+ destination_width: int,
166
+ destination_height: int,
167
+ inverse: bool = False,
168
+ ) -> ndarray:
169
+ """Project image using homography from source to destination keypoints."""
170
+ filtered_src = []
171
+ filtered_dst = []
172
+
173
+ for src_pt, dst_pt in zip(source_keypoints, destination_keypoints):
174
+ if dst_pt[0] == 0.0 and dst_pt[1] == 0.0: # ignore default / missing points
175
+ continue
176
+ filtered_src.append(src_pt)
177
+ filtered_dst.append(dst_pt)
178
+
179
+ if len(filtered_src) < 4:
180
+ raise ValueError("At least 4 valid keypoints are required for homography.")
181
+
182
+ source_points = array(filtered_src, dtype=float32)
183
+ destination_points = array(filtered_dst, dtype=float32)
184
+
185
+ if inverse:
186
+ H_inv, _ = findHomography(destination_points, source_points)
187
+ return warpPerspective(image, H_inv, (destination_width, destination_height))
188
+
189
+ H, _ = findHomography(source_points, destination_points)
190
+ projected_image = warpPerspective(image, H, (destination_width, destination_height))
191
+
192
+ validate_projected_corners(source_keypoints=source_keypoints, homography_matrix=H)
193
+ return projected_image
194
+
195
+
196
+ def extract_masks_for_ground_and_lines(
197
+ image: ndarray,
198
+ ) -> Tuple[ndarray, ndarray]:
199
+ """Extract masks for ground (gray) and lines (white) from template image."""
200
+ gray = cvtColor(image, COLOR_BGR2GRAY)
201
+ _, mask_ground = threshold(gray, 10, 255, THRESH_BINARY)
202
+ _, mask_lines = threshold(gray, 200, 255, THRESH_BINARY)
203
+ mask_ground_binary = (mask_ground > 0).astype(uint8)
204
+ mask_lines_binary = (mask_lines > 0).astype(uint8)
205
+ validate_mask_ground(mask=mask_ground_binary)
206
+ validate_mask_lines(mask=mask_lines_binary)
207
+ return mask_ground_binary, mask_lines_binary
208
+
209
+
210
+ def extract_mask_of_ground_lines_in_image(
211
+ image: ndarray,
212
+ ground_mask: ndarray,
213
+ blur_ksize: int = 5,
214
+ canny_low: int = 30,
215
+ canny_high: int = 100,
216
+ use_tophat: bool = True,
217
+ dilate_kernel_size: int = 3,
218
+ dilate_iterations: int = 3,
219
+ ) -> ndarray:
220
+ """Extract line mask from image using edge detection on ground region."""
221
+ gray = cvtColor(image, COLOR_BGR2GRAY)
222
+
223
+ if use_tophat:
224
+ kernel = getStructuringElement(MORPH_RECT, (31, 31))
225
+ gray = morphologyEx(gray, MORPH_TOPHAT, kernel)
226
+
227
+ if blur_ksize and blur_ksize % 2 == 1:
228
+ gray = GaussianBlur(gray, (blur_ksize, blur_ksize), 0)
229
+
230
+ image_edges = Canny(gray, canny_low, canny_high)
231
+ image_edges_on_ground = bitwise_and(image_edges, image_edges, mask=ground_mask)
232
+
233
+ if dilate_kernel_size > 1:
234
+ dilate_kernel = getStructuringElement(
235
+ MORPH_RECT, (dilate_kernel_size, dilate_kernel_size)
236
+ )
237
+ image_edges_on_ground = dilate(
238
+ image_edges_on_ground, dilate_kernel, iterations=dilate_iterations
239
+ )
240
+
241
+ return (image_edges_on_ground > 0).astype(uint8)
242
+
243
+
244
+ def evaluate_keypoints_for_frame(
245
+ template_keypoints: List[Tuple[int, int]],
246
+ frame_keypoints: List[Tuple[int, int]],
247
+ frame: ndarray,
248
+ floor_markings_template: ndarray,
249
+ ) -> float:
250
+ """
251
+ Evaluate keypoint accuracy for a single frame.
252
+
253
+ Returns score between 0.0 and 1.0 based on overlap between
254
+ projected template lines and detected lines in frame.
255
+ """
256
+ try:
257
+ warped_template = project_image_using_keypoints(
258
+ image=floor_markings_template,
259
+ source_keypoints=template_keypoints,
260
+ destination_keypoints=frame_keypoints,
261
+ destination_width=frame.shape[1],
262
+ destination_height=frame.shape[0],
263
+ )
264
+
265
+ mask_ground, mask_lines_expected = extract_masks_for_ground_and_lines(
266
+ image=warped_template
267
+ )
268
+
269
+ mask_lines_predicted = extract_mask_of_ground_lines_in_image(
270
+ image=frame, ground_mask=mask_ground
271
+ )
272
+
273
+ pixels_overlapping = bitwise_and(
274
+ mask_lines_expected, mask_lines_predicted
275
+ ).sum()
276
+
277
+ pixels_on_lines = mask_lines_expected.sum()
278
+
279
+ score = pixels_overlapping / (pixels_on_lines + 1e-8)
280
+ return min(1.0, max(0.0, score)) # Clamp to [0, 1]
281
+
282
+ except (InvalidMask, ValueError) as e:
283
+ logger.debug(f"Keypoint evaluation failed: {e}")
284
+ return 0.0
285
+ except Exception as e:
286
+ logger.error(f"Unexpected error in keypoint evaluation: {e}")
287
+ return 0.0
288
+
289
+
290
+ def load_template_from_file(
291
+ template_image_path: str,
292
+ ) -> Tuple[ndarray, List[Tuple[int, int]]]:
293
+ """
294
+ Load template image and use TEMPLATE_KEYPOINTS constant for keypoints.
295
+
296
+ Args:
297
+ template_image_path: Path to template image file
298
+
299
+ Returns:
300
+ template_image: Loaded template image
301
+ template_keypoints: List of (x, y) keypoint coordinates from TEMPLATE_KEYPOINTS constant
302
+ """
303
+ # Load template image
304
+ template_image = cv2.imread(template_image_path)
305
+ if template_image is None:
306
+ raise ValueError(f"Could not load template image from {template_image_path}")
307
+
308
+ # Use TEMPLATE_KEYPOINTS constant
309
+ if len(TEMPLATE_KEYPOINTS) == 0:
310
+ raise ValueError(
311
+ "TEMPLATE_KEYPOINTS constant is empty. Please define keypoints in keypoint_evaluation.py"
312
+ )
313
+
314
+ if len(TEMPLATE_KEYPOINTS) < 4:
315
+ raise ValueError(f"TEMPLATE_KEYPOINTS must have at least 4 keypoints, found {len(TEMPLATE_KEYPOINTS)}")
316
+
317
+ logger.info(f"Loaded template image: {template_image_path}")
318
+ logger.info(f"Using TEMPLATE_KEYPOINTS constant with {len(TEMPLATE_KEYPOINTS)} keypoints")
319
+
320
+ return template_image, TEMPLATE_KEYPOINTS
321
+
miner.py ADDED
@@ -0,0 +1,675 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pathlib import Path
2
+ from typing import List, Tuple, Dict, Optional
3
+ import sys
4
+ import os
5
+
6
+ from numpy import ndarray
7
+ from pydantic import BaseModel
8
+
9
+ from keypoint_evaluation import evaluate_keypoints_for_frame, load_template_from_file
10
+ sys.path.append(os.path.dirname(os.path.abspath(__file__)))
11
+
12
+ from ultralytics import YOLO
13
+ from team_cluster import TeamClassifier
14
+ from utils import (
15
+ BoundingBox,
16
+ Constants,
17
+ classify_teams_batch,
18
+ )
19
+
20
+ import time
21
+ import torch
22
+ import gc
23
+ import cv2
24
+ import numpy as np
25
+ from collections import defaultdict
26
+ from pitch import process_batch_input, get_cls_net
27
+ import yaml
28
+
29
+
30
+ class BoundingBox(BaseModel):
31
+ x1: int
32
+ y1: int
33
+ x2: int
34
+ y2: int
35
+ cls_id: int
36
+ conf: float
37
+
38
+
39
+ class TVFrameResult(BaseModel):
40
+ frame_id: int
41
+ boxes: List[BoundingBox]
42
+ keypoints: List[Tuple[int, int]]
43
+
44
+
45
+ class Miner:
46
+ SMALL_CONTAINED_IOA = Constants.SMALL_CONTAINED_IOA
47
+ SMALL_RATIO_MAX = Constants.SMALL_RATIO_MAX
48
+ SINGLE_PLAYER_HUE_PIVOT = Constants.SINGLE_PLAYER_HUE_PIVOT
49
+ CORNER_INDICES = Constants.CORNER_INDICES
50
+ KEYPOINTS_CONFIDENCE = Constants.KEYPOINTS_CONFIDENCE
51
+ CORNER_CONFIDENCE = Constants.CORNER_CONFIDENCE
52
+ GOALKEEPER_POSITION_MARGIN = Constants.GOALKEEPER_POSITION_MARGIN
53
+ MIN_SAMPLES_FOR_FIT = 16 # Minimum player crops needed before fitting TeamClassifier
54
+ MAX_SAMPLES_FOR_FIT = 600 # Maximum samples to avoid overfitting
55
+
56
+ def __init__(self, path_hf_repo: Path) -> None:
57
+ try:
58
+ device = "cuda" if torch.cuda.is_available() else "cpu"
59
+ model_path = path_hf_repo / "detection.onnx"
60
+ self.bbox_model = YOLO(model_path)
61
+
62
+ print(f"BBox Model Loaded: class name {self.bbox_model.names}")
63
+
64
+ team_model_path = path_hf_repo / "osnet_model.pth.tar-100"
65
+ self.team_classifier = TeamClassifier(
66
+ device=device,
67
+ batch_size=32,
68
+ model_name=str(team_model_path)
69
+ )
70
+ print("Team Classifier Loaded")
71
+
72
+ # Team classification state
73
+ self.team_classifier_fitted = False
74
+ self.player_crops_for_fit = []
75
+
76
+ self.keypoints_model_yolo = YOLO(path_hf_repo / "keypoint.pt")
77
+
78
+ model_kp_path = path_hf_repo / 'keypoint'
79
+ config_kp_path = path_hf_repo / 'hrnetv2_w48.yaml'
80
+ cfg_kp = yaml.safe_load(open(config_kp_path, 'r'))
81
+
82
+ loaded_state_kp = torch.load(model_kp_path, map_location=device)
83
+ model = get_cls_net(cfg_kp)
84
+ model.load_state_dict(loaded_state_kp)
85
+ model.to(device)
86
+ model.eval()
87
+
88
+ self.keypoints_model = model
89
+ print("Keypoints Model (keypoint.pt) Loaded")
90
+
91
+ template_image_path = path_hf_repo / "football_pitch_template.png"
92
+ self.template_image, self.template_keypoints = load_template_from_file(str(template_image_path))
93
+
94
+ self.kp_threshold = 0.1
95
+ self.pitch_batch_size = 4
96
+ self.health = "healthy"
97
+ print("✅ Keypoints Model Loaded")
98
+ except Exception as e:
99
+ self.health = "❌ Miner initialization failed: " + str(e)
100
+ print(self.health)
101
+
102
+ def __repr__(self) -> str:
103
+ if self.health == 'healthy':
104
+ return (
105
+ f"health: {self.health}\n"
106
+ f"BBox Model: {type(self.bbox_model).__name__}\n"
107
+ f"Keypoints Model: {type(self.keypoints_model).__name__}"
108
+ )
109
+ else:
110
+ return self.health
111
+
112
+ def _calculate_iou(self, box1: Tuple[float, float, float, float],
113
+ box2: Tuple[float, float, float, float]) -> float:
114
+ """
115
+ Calculate Intersection over Union (IoU) between two bounding boxes.
116
+ Args:
117
+ box1: (x1, y1, x2, y2)
118
+ box2: (x1, y1, x2, y2)
119
+ Returns:
120
+ IoU score (0-1)
121
+ """
122
+ x1_1, y1_1, x2_1, y2_1 = box1
123
+ x1_2, y1_2, x2_2, y2_2 = box2
124
+
125
+ # Calculate intersection area
126
+ x_left = max(x1_1, x1_2)
127
+ y_top = max(y1_1, y1_2)
128
+ x_right = min(x2_1, x2_2)
129
+ y_bottom = min(y2_1, y2_2)
130
+
131
+ if x_right < x_left or y_bottom < y_top:
132
+ return 0.0
133
+
134
+ intersection_area = (x_right - x_left) * (y_bottom - y_top)
135
+
136
+ # Calculate union area
137
+ box1_area = (x2_1 - x1_1) * (y2_1 - y1_1)
138
+ box2_area = (x2_2 - x1_2) * (y2_2 - y1_2)
139
+ union_area = box1_area + box2_area - intersection_area
140
+
141
+ if union_area == 0:
142
+ return 0.0
143
+
144
+ return intersection_area / union_area
145
+
146
+ def _extract_jersey_region(self, crop: ndarray) -> ndarray:
147
+ """
148
+ Extract jersey region (upper body) from player crop.
149
+ For close-ups, focuses on upper 60%, for distant shots uses full crop.
150
+ """
151
+ if crop is None or crop.size == 0:
152
+ return crop
153
+
154
+ h, w = crop.shape[:2]
155
+ if h < 10 or w < 10:
156
+ return crop
157
+
158
+ # For close-up shots, extract upper body (jersey region)
159
+ is_closeup = h > 100 or (h * w) > 12000
160
+ if is_closeup:
161
+ # Upper 60% of the crop (jersey area, avoiding shorts)
162
+ jersey_top = 0
163
+ jersey_bottom = int(h * 0.60)
164
+ jersey_left = max(0, int(w * 0.05))
165
+ jersey_right = min(w, int(w * 0.95))
166
+ return crop[jersey_top:jersey_bottom, jersey_left:jersey_right]
167
+ return crop
168
+
169
+ def _extract_color_signature(self, crop: ndarray) -> Optional[np.ndarray]:
170
+ """
171
+ Extract color signature from jersey region using HSV and LAB color spaces.
172
+ Returns a feature vector with dominant colors and color statistics.
173
+ """
174
+ if crop is None or crop.size == 0:
175
+ return None
176
+
177
+ jersey_region = self._extract_jersey_region(crop)
178
+ if jersey_region.size == 0:
179
+ return None
180
+
181
+ try:
182
+ # Convert to HSV and LAB color spaces
183
+ hsv = cv2.cvtColor(jersey_region, cv2.COLOR_BGR2HSV)
184
+ lab = cv2.cvtColor(jersey_region, cv2.COLOR_BGR2LAB)
185
+
186
+ # Reshape for processing
187
+ hsv_flat = hsv.reshape(-1, 3).astype(np.float32)
188
+ lab_flat = lab.reshape(-1, 3).astype(np.float32)
189
+
190
+ # Compute statistics for HSV
191
+ hsv_mean = np.mean(hsv_flat, axis=0) / 255.0
192
+ hsv_std = np.std(hsv_flat, axis=0) / 255.0
193
+
194
+ # Compute statistics for LAB
195
+ lab_mean = np.mean(lab_flat, axis=0) / 255.0
196
+ lab_std = np.std(lab_flat, axis=0) / 255.0
197
+
198
+ # Dominant color (most frequent hue)
199
+ hue_hist, _ = np.histogram(hsv_flat[:, 0], bins=36, range=(0, 180))
200
+ dominant_hue = np.argmax(hue_hist) * 5 # Convert to hue value
201
+
202
+ # Combine features
203
+ color_features = np.concatenate([
204
+ hsv_mean,
205
+ hsv_std,
206
+ lab_mean[:2], # L and A channels (B is less informative)
207
+ lab_std[:2],
208
+ [dominant_hue / 180.0] # Normalized dominant hue
209
+ ])
210
+
211
+ return color_features
212
+ except Exception as e:
213
+ print(f"Error extracting color signature: {e}")
214
+ return None
215
+
216
+ def _get_spatial_position(self, bbox: Tuple[float, float, float, float],
217
+ frame_width: int, frame_height: int) -> Tuple[float, float]:
218
+ """
219
+ Get normalized spatial position of player on the pitch.
220
+ Returns (x_normalized, y_normalized) where 0,0 is top-left.
221
+ """
222
+ x1, y1, x2, y2 = bbox
223
+ center_x = (x1 + x2) / 2.0
224
+ center_y = (y1 + y2) / 2.0
225
+
226
+ # Normalize to [0, 1]
227
+ x_norm = center_x / frame_width if frame_width > 0 else 0.5
228
+ y_norm = center_y / frame_height if frame_height > 0 else 0.5
229
+
230
+ return (x_norm, y_norm)
231
+
232
+ def _find_best_match(self, target_box: Tuple[float, float, float, float],
233
+ predicted_frame_data: Dict[int, Tuple[Tuple, str]],
234
+ iou_threshold: float) -> Tuple[Optional[str], float]:
235
+ """
236
+ Find best matching box in predicted frame data using IoU.
237
+ """
238
+ best_iou = 0.0
239
+ best_team_id = None
240
+
241
+ for idx, (bbox, team_cls_id) in predicted_frame_data.items():
242
+ iou = self._calculate_iou(target_box, bbox)
243
+ if iou > best_iou and iou >= iou_threshold:
244
+ best_iou = iou
245
+ best_team_id = team_cls_id
246
+
247
+ return (best_team_id, best_iou)
248
+
249
+ def _detect_objects_batch(self, decoded_images: List[ndarray]) -> Dict[int, List[BoundingBox]]:
250
+ batch_size = 16
251
+ detection_results = []
252
+ n_frames = len(decoded_images)
253
+ for frame_number in range(0, n_frames, batch_size):
254
+ batch_images = decoded_images[frame_number: frame_number + batch_size]
255
+ detections = self.bbox_model(batch_images, verbose=False, save=False)
256
+ detection_results.extend(detections)
257
+
258
+ return detection_results
259
+
260
+ def _team_classify(self, detection_results, decoded_images, offset):
261
+ """
262
+ Hybrid team classification combining:
263
+ 1. Appearance features (OSNet)
264
+ 2. Color signatures (HSV/LAB)
265
+ 3. Spatial priors (left/right side of pitch)
266
+ 4. Temporal tracking (same player = same team)
267
+ """
268
+ start = time.time()
269
+
270
+ # Phase 1: Collect samples and fit appearance-based classifier
271
+ fit_sample_size = min(self.MAX_SAMPLES_FOR_FIT, len(detection_results) * 10)
272
+ player_crops_for_fit = []
273
+
274
+ for frame_id in range(len(detection_results)):
275
+ detection_box = detection_results[frame_id].boxes.data
276
+ if len(detection_box) < 4:
277
+ continue
278
+
279
+ if len(player_crops_for_fit) < fit_sample_size:
280
+ frame_image = decoded_images[frame_id]
281
+ for box in detection_box:
282
+ x1, y1, x2, y2, conf, cls_id = box.tolist()
283
+ if conf < 0.5 or cls_id != 2:
284
+ continue
285
+ crop = frame_image[int(y1):int(y2), int(x1):int(x2)]
286
+ if crop.size > 0:
287
+ player_crops_for_fit.append(crop)
288
+
289
+ if self.team_classifier and not self.team_classifier_fitted and len(player_crops_for_fit) >= fit_sample_size:
290
+ print(f"Fitting TeamClassifier (OSNet) with {len(player_crops_for_fit)} player crops")
291
+ self.team_classifier.fit(player_crops_for_fit)
292
+ self.team_classifier_fitted = True
293
+ break
294
+
295
+ if not self.team_classifier_fitted and len(player_crops_for_fit) >= self.MIN_SAMPLES_FOR_FIT:
296
+ print(f"Fallback: Fitting TeamClassifier with {len(player_crops_for_fit)} player crops")
297
+ self.team_classifier.fit(player_crops_for_fit)
298
+ self.team_classifier_fitted = True
299
+
300
+ print(f"Fitting time: {time.time() - start:.2f}s")
301
+
302
+ # Phase 2: Hybrid classification for all frames
303
+ start = time.time()
304
+ bboxes: dict[int, list[BoundingBox]] = {}
305
+
306
+ # Temporal tracking: {track_id: (team_id, confidence, last_frame)}
307
+ player_tracks: Dict[Tuple, Tuple[int, float, int]] = {}
308
+
309
+ # Spatial team assignment: track which team is on which side
310
+ left_side_team = None
311
+ right_side_team = None
312
+
313
+ for frame_id in range(len(detection_results)):
314
+ detection_box = detection_results[frame_id].boxes.data
315
+ frame_image = decoded_images[frame_id]
316
+ frame_h, frame_w = frame_image.shape[:2]
317
+ boxes = []
318
+
319
+ # Collect all players in this frame
320
+ player_data = [] # (idx, crop, bbox, spatial_pos, color_sig)
321
+
322
+ for idx, box in enumerate(detection_box):
323
+ x1, y1, x2, y2, conf, cls_id = box.tolist()
324
+ if cls_id != 2 or conf < 0.6:
325
+ continue
326
+
327
+ crop = frame_image[int(y1):int(y2), int(x1):int(x2)]
328
+ if crop.size == 0:
329
+ continue
330
+
331
+ bbox = (x1, y1, x2, y2)
332
+ spatial_pos = self._get_spatial_position(bbox, frame_w, frame_h)
333
+ color_sig = self._extract_color_signature(crop)
334
+
335
+ player_data.append((idx, crop, bbox, spatial_pos, color_sig))
336
+
337
+ if len(player_data) == 0:
338
+ bboxes[offset + frame_id] = []
339
+ continue
340
+
341
+ # Step 1: Get appearance-based predictions (OSNet)
342
+ appearance_predictions = {}
343
+ if self.team_classifier and self.team_classifier_fitted:
344
+ crops = [data[1] for data in player_data]
345
+ appearance_team_ids = self.team_classifier.predict(crops)
346
+ for (idx, _, _, _, _), team_id in zip(player_data, appearance_team_ids):
347
+ appearance_predictions[idx] = team_id
348
+
349
+ # Step 2: Extract color signatures and cluster
350
+ color_signatures = []
351
+ color_indices = []
352
+ for idx, _, _, _, color_sig in player_data:
353
+ if color_sig is not None:
354
+ color_signatures.append(color_sig)
355
+ color_indices.append(idx)
356
+
357
+ color_predictions = {}
358
+ if len(color_signatures) >= 4:
359
+ try:
360
+ from sklearn.cluster import KMeans
361
+ color_kmeans = KMeans(n_clusters=2, random_state=42, n_init=10)
362
+ color_clusters = color_kmeans.fit_predict(color_signatures)
363
+ for idx, cluster_id in zip(color_indices, color_clusters):
364
+ color_predictions[idx] = cluster_id
365
+ except Exception as e:
366
+ print(f"Color clustering failed: {e}")
367
+
368
+ # Step 3: Apply spatial priors
369
+ # Determine which team is on which side based on majority
370
+ if left_side_team is None or right_side_team is None:
371
+ left_side_players = [p for p in player_data if p[3][0] < 0.5] # x < 0.5
372
+ right_side_players = [p for p in player_data if p[3][0] >= 0.5] # x >= 0.5
373
+
374
+ if len(left_side_players) >= 2 and len(right_side_players) >= 2:
375
+ # Use appearance predictions to determine sides
376
+ left_teams = [appearance_predictions.get(p[0]) for p in left_side_players
377
+ if p[0] in appearance_predictions]
378
+ right_teams = [appearance_predictions.get(p[0]) for p in right_side_players
379
+ if p[0] in appearance_predictions]
380
+
381
+ if left_teams and right_teams:
382
+ left_team_mode = max(set(left_teams), key=left_teams.count)
383
+ right_team_mode = max(set(right_teams), key=right_teams.count)
384
+
385
+ if left_team_mode != right_team_mode:
386
+ left_side_team = left_team_mode
387
+ right_side_team = right_team_mode
388
+
389
+ # Step 4: Combine predictions with voting
390
+ final_predictions = {}
391
+ for idx, _, bbox, spatial_pos, _ in player_data:
392
+ votes = []
393
+ weights = []
394
+
395
+ # Appearance vote (weight: 0.4)
396
+ if idx in appearance_predictions:
397
+ votes.append(appearance_predictions[idx])
398
+ weights.append(0.4)
399
+
400
+ # Color vote (weight: 0.3)
401
+ if idx in color_predictions:
402
+ votes.append(color_predictions[idx])
403
+ weights.append(0.3)
404
+
405
+ # Spatial vote (weight: 0.3)
406
+ if left_side_team is not None and right_side_team is not None:
407
+ x_pos, _ = spatial_pos
408
+ if x_pos < 0.5:
409
+ spatial_team = left_side_team
410
+ else:
411
+ spatial_team = right_side_team
412
+ votes.append(spatial_team)
413
+ weights.append(0.3)
414
+
415
+ # Temporal vote (weight: 0.2) - match with previous frames
416
+ if len(votes) > 0:
417
+ # Simple temporal matching: find similar bbox in previous frames
418
+ best_track_match = None
419
+ best_track_iou = 0.0
420
+ for track_key, (track_team, track_conf, track_frame) in player_tracks.items():
421
+ if abs(track_frame - frame_id) <= 5: # Within 5 frames
422
+ track_bbox = track_key
423
+ iou = self._calculate_iou(bbox, track_bbox)
424
+ if iou > best_track_iou and iou > 0.3:
425
+ best_track_iou = iou
426
+ best_track_match = track_team
427
+
428
+ if best_track_match is not None:
429
+ votes.append(best_track_match)
430
+ weights.append(0.2)
431
+
432
+ # Weighted voting
433
+ if len(votes) > 0:
434
+ team_0_score = sum(w for v, w in zip(votes, weights) if v == 0)
435
+ team_1_score = sum(w for v, w in zip(votes, weights) if v == 1)
436
+
437
+ if team_0_score > team_1_score:
438
+ final_team = 0
439
+ elif team_1_score > team_0_score:
440
+ final_team = 1
441
+ else:
442
+ # Tie: use appearance prediction or first vote
443
+ final_team = votes[0] if votes else 0
444
+
445
+ final_predictions[idx] = final_team
446
+
447
+ # Update tracking
448
+ track_key = bbox
449
+ player_tracks[track_key] = (final_team, max(team_0_score, team_1_score), frame_id)
450
+
451
+ # Step 5: Generate output boxes
452
+ for idx, box in enumerate(detection_box):
453
+ x1, y1, x2, y2, conf, cls_id = box.tolist()
454
+ if cls_id == 2 and conf < 0.6:
455
+ continue
456
+
457
+ # Check overlap with staff
458
+ overlap_staff = False
459
+ for idy, boxy in enumerate(detection_box):
460
+ s_x1, s_y1, s_x2, s_y2, s_conf, s_cls_id = boxy.tolist()
461
+ if cls_id == 2 and s_cls_id == 4:
462
+ staff_iou = self._calculate_iou(box[:4], boxy[:4])
463
+ if staff_iou >= 0.8:
464
+ overlap_staff = True
465
+ break
466
+ if overlap_staff:
467
+ continue
468
+
469
+ mapped_cls_id = str(int(cls_id))
470
+
471
+ # Override with team prediction
472
+ if idx in final_predictions:
473
+ mapped_cls_id = str(6 + int(final_predictions[idx]))
474
+
475
+ if mapped_cls_id != '4':
476
+ if int(mapped_cls_id) == 3 and conf < 0.5:
477
+ continue
478
+ boxes.append(
479
+ BoundingBox(
480
+ x1=int(x1),
481
+ y1=int(y1),
482
+ x2=int(x2),
483
+ y2=int(y2),
484
+ cls_id=int(mapped_cls_id),
485
+ conf=float(conf),
486
+ )
487
+ )
488
+
489
+ # Handle footballs - keep only the best one
490
+ footballs = [bb for bb in boxes if int(bb.cls_id) == 0]
491
+ if len(footballs) > 1:
492
+ best_ball = max(footballs, key=lambda b: b.conf)
493
+ boxes = [bb for bb in boxes if int(bb.cls_id) != 0]
494
+ boxes.append(best_ball)
495
+
496
+ bboxes[offset + frame_id] = boxes
497
+
498
+ print(f"Hybrid team classification time: {time.time() - start:.2f}s")
499
+ return bboxes
500
+
501
+
502
+ def predict_batch(self, batch_images: List[ndarray], offset: int, n_keypoints: int) -> List[TVFrameResult]:
503
+ start = time.time()
504
+ detection_results = self._detect_objects_batch(batch_images)
505
+ end = time.time()
506
+ print(f"Detection time: {end - start}")
507
+
508
+ # Use hybrid team classification
509
+ start = time.time()
510
+ bboxes = self._team_classify(detection_results, batch_images, offset)
511
+ end = time.time()
512
+ print(f"Team classify time: {end - start}")
513
+
514
+ # Phase 3: Keypoint Detection
515
+ keypoints_yolo: Dict[int, List[Tuple[int, int]]] = {}
516
+
517
+ keypoints_yolo = self._detect_keypoints_batch(batch_images, offset, n_keypoints)
518
+
519
+
520
+ pitch_batch_size = min(self.pitch_batch_size, len(batch_images))
521
+ keypoints: Dict[int, List[Tuple[int, int]]] = {}
522
+
523
+ start = time.time()
524
+ while True:
525
+ gc.collect()
526
+ if torch.cuda.is_available():
527
+ torch.cuda.empty_cache()
528
+ torch.cuda.synchronize()
529
+ device_str = "cuda"
530
+ keypoints_result = process_batch_input(
531
+ batch_images,
532
+ self.keypoints_model,
533
+ self.kp_threshold,
534
+ device_str,
535
+ batch_size=pitch_batch_size,
536
+ )
537
+ if keypoints_result is not None and len(keypoints_result) > 0:
538
+ for frame_number_in_batch, kp_dict in enumerate(keypoints_result):
539
+ if frame_number_in_batch >= len(batch_images):
540
+ break
541
+ frame_keypoints: List[Tuple[int, int]] = []
542
+ try:
543
+ height, width = batch_images[frame_number_in_batch].shape[:2]
544
+ if kp_dict is not None and isinstance(kp_dict, dict):
545
+ for idx in range(32):
546
+ x, y = 0, 0
547
+ kp_idx = idx + 1
548
+ if kp_idx in kp_dict:
549
+ try:
550
+ kp_data = kp_dict[kp_idx]
551
+ if isinstance(kp_data, dict) and "x" in kp_data and "y" in kp_data:
552
+ x = int(kp_data["x"] * width)
553
+ y = int(kp_data["y"] * height)
554
+ except (KeyError, TypeError, ValueError):
555
+ pass
556
+ frame_keypoints.append((x, y))
557
+ except (IndexError, ValueError, AttributeError):
558
+ frame_keypoints = [(0, 0)] * 32
559
+ if len(frame_keypoints) < n_keypoints:
560
+ frame_keypoints.extend([(0, 0)] * (n_keypoints - len(frame_keypoints)))
561
+ else:
562
+ frame_keypoints = frame_keypoints[:n_keypoints]
563
+
564
+ valid_keypoints = [kp for kp in frame_keypoints if kp[0] != 0.0 or kp[1] != 0.0]
565
+
566
+ frame_keypoints_yolo = keypoints_yolo.get(offset + frame_number_in_batch, frame_keypoints)
567
+ valid_keypoints_yolo = [kp for kp in frame_keypoints_yolo if kp[0] != 0.0 or kp[1] != 0.0]
568
+
569
+ if len(valid_keypoints_yolo) > 3 and len(valid_keypoints) > 3:
570
+ try:
571
+ score = evaluate_keypoints_for_frame(
572
+ template_keypoints=self.template_keypoints,
573
+ frame_keypoints=frame_keypoints,
574
+ frame=batch_images[frame_number_in_batch],
575
+ floor_markings_template=self.template_image.copy(),
576
+ )
577
+
578
+ score_yolo = evaluate_keypoints_for_frame(
579
+ template_keypoints=self.template_keypoints,
580
+ frame_keypoints=frame_keypoints_yolo,
581
+ frame=batch_images[frame_number_in_batch],
582
+ floor_markings_template=self.template_image.copy(),
583
+ )
584
+
585
+ if score_yolo > score:
586
+ frame_keypoints = frame_keypoints_yolo
587
+ except Exception as e:
588
+ pass
589
+ elif len(valid_keypoints_yolo) > 3:
590
+ frame_keypoints = frame_keypoints_yolo
591
+
592
+
593
+ keypoints[offset + frame_number_in_batch] = frame_keypoints
594
+ break
595
+ end = time.time()
596
+ print(f"Keypoint time: {end - start}")
597
+
598
+ results: List[TVFrameResult] = []
599
+ for frame_number in range(offset, offset + len(batch_images)):
600
+ frame_boxes = bboxes.get(frame_number, [])
601
+ result = TVFrameResult(
602
+ frame_id=frame_number,
603
+ boxes=frame_boxes,
604
+ keypoints=keypoints.get(
605
+ frame_number,
606
+ [(0, 0) for _ in range(n_keypoints)],
607
+ ),
608
+ )
609
+ results.append(result)
610
+
611
+ gc.collect()
612
+ if torch.cuda.is_available():
613
+ torch.cuda.empty_cache()
614
+ torch.cuda.synchronize()
615
+
616
+ return results
617
+
618
+ def _detect_keypoints_batch(self, batch_images: List[ndarray],
619
+ offset: int, n_keypoints: int) -> Dict[int, List[Tuple[int, int]]]:
620
+ """
621
+ Phase 3: Keypoint detection for all frames in batch.
622
+
623
+ Args:
624
+ batch_images: List of images to process
625
+ offset: Frame offset for numbering
626
+ n_keypoints: Number of keypoints expected
627
+
628
+ Returns:
629
+ Dictionary mapping frame_id to list of keypoint coordinates
630
+ """
631
+ keypoints: Dict[int, List[Tuple[int, int]]] = {}
632
+ keypoints_model_results = self.keypoints_model_yolo.predict(batch_images)
633
+
634
+ if keypoints_model_results is None:
635
+ return keypoints
636
+
637
+ for frame_idx_in_batch, detection in enumerate(keypoints_model_results):
638
+ if not hasattr(detection, "keypoints") or detection.keypoints is None:
639
+ continue
640
+
641
+ # Extract keypoints with confidence
642
+ frame_keypoints_with_conf: List[Tuple[int, int, float]] = []
643
+ for i, part_points in enumerate(detection.keypoints.data):
644
+ for k_id, (x, y, _) in enumerate(part_points):
645
+ confidence = float(detection.keypoints.conf[i][k_id])
646
+ frame_keypoints_with_conf.append((int(x), int(y), confidence))
647
+
648
+ # Pad or truncate to expected number of keypoints
649
+ if len(frame_keypoints_with_conf) < n_keypoints:
650
+ frame_keypoints_with_conf.extend(
651
+ [(0, 0, 0.0)] * (n_keypoints - len(frame_keypoints_with_conf))
652
+ )
653
+ else:
654
+ frame_keypoints_with_conf = frame_keypoints_with_conf[:n_keypoints]
655
+
656
+ # Filter keypoints based on confidence thresholds
657
+ filtered_keypoints: List[Tuple[int, int]] = []
658
+ for idx, (x, y, confidence) in enumerate(frame_keypoints_with_conf):
659
+ if idx in self.CORNER_INDICES:
660
+ # Corner keypoints have lower confidence threshold
661
+ if confidence < 0.3:
662
+ filtered_keypoints.append((0, 0))
663
+ else:
664
+ filtered_keypoints.append((int(x), int(y)))
665
+ else:
666
+ # Regular keypoints
667
+ if confidence < 0.5:
668
+ filtered_keypoints.append((0, 0))
669
+ else:
670
+ filtered_keypoints.append((int(x), int(y)))
671
+
672
+ frame_id = offset + frame_idx_in_batch
673
+ keypoints[frame_id] = filtered_keypoints
674
+
675
+ return keypoints
miner1.py ADDED
@@ -0,0 +1,509 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pathlib import Path
2
+ from typing import List, Tuple, Dict, Optional
3
+ import sys, os
4
+ sys.path.append(os.path.dirname(os.path.abspath(__file__)))
5
+ import onnxruntime as ort
6
+ import numpy as np
7
+ import cv2
8
+ from torchvision.ops import batched_nms
9
+ import torch
10
+ from ultralytics import YOLO
11
+ from numpy import ndarray
12
+ from pydantic import BaseModel
13
+ from team_cluster import TeamClassifier
14
+ from utils import (
15
+ BoundingBox,
16
+ Constants,
17
+ suppress_small_contained_boxes,
18
+ classify_teams_batch,
19
+ )
20
+
21
+
22
+ class TVFrameResult(BaseModel):
23
+ frame_id: int
24
+ boxes: List[BoundingBox]
25
+ keypoints: List[Tuple[int, int]]
26
+
27
+
28
+ class Miner:
29
+ """
30
+ Football video analysis system for object detection and team classification.
31
+ """
32
+ # Use constants from utils
33
+ SMALL_CONTAINED_IOA = Constants.SMALL_CONTAINED_IOA
34
+ SMALL_RATIO_MAX = Constants.SMALL_RATIO_MAX
35
+ SINGLE_PLAYER_HUE_PIVOT = Constants.SINGLE_PLAYER_HUE_PIVOT
36
+ CORNER_INDICES = Constants.CORNER_INDICES
37
+ KEYPOINTS_CONFIDENCE = Constants.KEYPOINTS_CONFIDENCE
38
+ CORNER_CONFIDENCE = Constants.CORNER_CONFIDENCE
39
+ GOALKEEPER_POSITION_MARGIN = Constants.GOALKEEPER_POSITION_MARGIN
40
+ MIN_SAMPLES_FOR_FIT = 16 # Minimum player crops needed before fitting TeamClassifier
41
+ MAX_SAMPLES_FOR_FIT = 500 # Maximum samples to avoid overfitting
42
+
43
+ def __init__(self, path_hf_repo: Path) -> None:
44
+ providers = [
45
+ 'CUDAExecutionProvider',
46
+ 'CPUExecutionProvider'
47
+ ]
48
+ model_path = path_hf_repo / "detection.onnx"
49
+ session = ort.InferenceSession(model_path, providers=providers)
50
+
51
+ input_name = session.get_inputs()[0].name
52
+ height = width = 640
53
+ dummy = np.zeros((1, 3, height, width), dtype=np.float32)
54
+ session.run(None, {input_name: dummy})
55
+ model = session
56
+ self.bbox_model = model
57
+
58
+ print("BBox Model Loaded")
59
+ self.keypoints_model = YOLO(path_hf_repo / "keypoint.pt")
60
+ print("Keypoints Model (keypoint.pt) Loaded")
61
+ # Initialize team classifier with OSNet model
62
+ team_model_path = path_hf_repo / "osnet_model.pth.tar-100"
63
+ device = 'cuda'
64
+ self.team_classifier = TeamClassifier(
65
+ device=device,
66
+ batch_size=32,
67
+ model_name=str(team_model_path)
68
+ )
69
+ print("Team Classifier Loaded")
70
+
71
+ # Team classification state
72
+ self.team_classifier_fitted = False
73
+ self.player_crops_for_fit = [] # Collect samples across frames
74
+
75
+ def __repr__(self) -> str:
76
+ return (
77
+ f"BBox Model: {type(self.bbox_model).__name__}\n"
78
+ f"Keypoints Model: {type(self.keypoints_model).__name__}"
79
+ )
80
+
81
+
82
+
83
+ def _handle_multiple_goalkeepers(self, boxes: List[BoundingBox]) -> List[BoundingBox]:
84
+ """
85
+ Handle goalkeeper detection issues:
86
+ 1. Fix misplaced goalkeepers (standing in middle of field)
87
+ 2. Limit to maximum 2 goalkeepers (one from each team)
88
+
89
+ Returns:
90
+ Filtered list of boxes with corrected goalkeepers
91
+ """
92
+ # Step 1: Fix misplaced goalkeepers first
93
+ # Convert goalkeepers in middle of field to regular players
94
+ boxes = self._fix_misplaced_goalkeepers(boxes)
95
+
96
+ # Step 2: Handle multiple goalkeepers (after fixing misplaced ones)
97
+ gk_idxs = [i for i, bb in enumerate(boxes) if int(bb.cls_id) == 1]
98
+ if len(gk_idxs) <= 2:
99
+ return boxes
100
+
101
+ # Sort goalkeepers by confidence (highest first)
102
+ gk_idxs_sorted = sorted(gk_idxs, key=lambda i: boxes[i].conf, reverse=True)
103
+ keep_gk_idxs = set(gk_idxs_sorted[:2]) # Keep top 2 goalkeepers
104
+
105
+ # Create new list keeping only top 2 goalkeepers
106
+ filtered_boxes = []
107
+ for i, box in enumerate(boxes):
108
+ if int(box.cls_id) == 1:
109
+ # Only keep the top 2 goalkeepers by confidence
110
+ if i in keep_gk_idxs:
111
+ filtered_boxes.append(box)
112
+ # Skip extra goalkeepers
113
+ else:
114
+ # Keep all non-goalkeeper boxes
115
+ filtered_boxes.append(box)
116
+
117
+ return filtered_boxes
118
+
119
+ def _fix_misplaced_goalkeepers(self, boxes: List[BoundingBox]) -> List[BoundingBox]:
120
+ """
121
+ """
122
+ gk_idxs = [i for i, bb in enumerate(boxes) if int(bb.cls_id) == 1]
123
+ player_idxs = [i for i, bb in enumerate(boxes) if int(bb.cls_id) == 2]
124
+
125
+ if len(gk_idxs) == 0 or len(player_idxs) < 2:
126
+ return boxes
127
+
128
+ updated_boxes = boxes.copy()
129
+
130
+ for gk_idx in gk_idxs:
131
+ if boxes[gk_idx].conf < 0.3:
132
+ updated_boxes[gk_idx].cls_id = 2
133
+
134
+ return updated_boxes
135
+
136
+
137
+ def _pre_process_img(self, frames: List[np.ndarray], scale: float = 640.0) -> np.ndarray:
138
+ """
139
+ Preprocess images for ONNX inference.
140
+
141
+ Args:
142
+ frames: List of BGR frames
143
+ scale: Target scale for resizing
144
+
145
+ Returns:
146
+ Preprocessed numpy array ready for ONNX inference
147
+ """
148
+ imgs = np.stack([cv2.resize(frame, (int(scale), int(scale))) for frame in frames])
149
+ imgs = imgs.transpose(0, 3, 1, 2) # BHWC to BCHW
150
+ imgs = imgs.astype(np.float32) / 255.0 # Normalize to [0, 1]
151
+ return imgs
152
+
153
+ def _post_process_output(self, outputs: np.ndarray, x_scale: float, y_scale: float,
154
+ conf_thresh: float = 0.6, nms_thresh: float = 0.55) -> List[List[Tuple]]:
155
+ """
156
+ Post-process ONNX model outputs to get detections.
157
+
158
+ Args:
159
+ outputs: Raw ONNX model outputs
160
+ x_scale: X-axis scaling factor
161
+ y_scale: Y-axis scaling factor
162
+ conf_thresh: Confidence threshold
163
+ nms_thresh: NMS threshold
164
+
165
+ Returns:
166
+ List of detections for each frame: [(box, conf, class_id), ...]
167
+ """
168
+ B, C, N = outputs.shape
169
+ outputs = torch.from_numpy(outputs)
170
+ outputs = outputs.permute(0, 2, 1) # B,C,N -> B,N,C
171
+
172
+ boxes = outputs[..., :4]
173
+ class_scores = 1 / (1 + torch.exp(-outputs[..., 4:])) # Sigmoid activation
174
+ conf, class_id = class_scores.max(dim=2)
175
+
176
+ mask = conf > conf_thresh
177
+
178
+ # Special handling for balls - keep best one even with lower confidence
179
+ for i in range(class_id.shape[0]): # loop over batch
180
+ # Find detections that are balls
181
+ ball_mask = class_id[i] == 0
182
+ ball_idx = ball_mask.nonzero(as_tuple=True)[0]
183
+ if ball_idx.numel() > 0:
184
+ # Pick the one with the highest confidence
185
+ best_ball_idx = ball_idx[conf[i, ball_idx].argmax()]
186
+ if conf[i, best_ball_idx] >= 0.55: # apply confidence threshold
187
+ mask[i, best_ball_idx] = True
188
+
189
+ batch_idx, pred_idx = mask.nonzero(as_tuple=True)
190
+
191
+ if len(batch_idx) == 0:
192
+ return [[] for _ in range(B)]
193
+
194
+ boxes = boxes[batch_idx, pred_idx]
195
+ conf = conf[batch_idx, pred_idx]
196
+ class_id = class_id[batch_idx, pred_idx]
197
+
198
+ # Convert from center format to xyxy format
199
+ x, y, w, h = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
200
+ x1 = (x - w / 2) * x_scale
201
+ y1 = (y - h / 2) * y_scale
202
+ x2 = (x + w / 2) * x_scale
203
+ y2 = (y + h / 2) * y_scale
204
+ boxes_xyxy = torch.stack([x1, y1, x2, y2], dim=1)
205
+
206
+ # Apply batched NMS
207
+ max_coord = 1e4
208
+ offset = batch_idx.to(boxes_xyxy) * max_coord
209
+ boxes_for_nms = boxes_xyxy + offset[:, None]
210
+
211
+ keep = batched_nms(boxes_for_nms, conf, batch_idx, nms_thresh)
212
+
213
+ boxes_final = boxes_xyxy[keep]
214
+ conf_final = conf[keep]
215
+ class_final = class_id[keep]
216
+ batch_final = batch_idx[keep]
217
+
218
+ # Group results by batch
219
+ results = [[] for _ in range(B)]
220
+ for b in range(B):
221
+ mask_b = batch_final == b
222
+ if mask_b.sum() == 0:
223
+ continue
224
+ results[b] = list(zip(boxes_final[mask_b].numpy(),
225
+ conf_final[mask_b].numpy(),
226
+ class_final[mask_b].numpy()))
227
+ return results
228
+
229
+ def _ioa(self, a: BoundingBox, b: BoundingBox) -> float:
230
+ inter = self._intersect_area(a, b)
231
+ aa = self._area(a)
232
+ if aa <= 0:
233
+ return 0.0
234
+ return inter / aa
235
+
236
+ def suppress_small_contained(self, boxes: List[BoundingBox]) -> List[BoundingBox]:
237
+ if len(boxes) <= 1:
238
+ return boxes
239
+ keep = [True] * len(boxes)
240
+ areas = [self._area(bb) for bb in boxes]
241
+ for i in range(len(boxes)):
242
+ if not keep[i]:
243
+ continue
244
+ for j in range(len(boxes)):
245
+ if i == j or not keep[j]:
246
+ continue
247
+ ai, aj = areas[i], areas[j]
248
+ if ai == 0 or aj == 0:
249
+ continue
250
+ if ai <= aj:
251
+ ratio = ai / aj
252
+ if ratio <= self.SMALL_RATIO_MAX:
253
+ ioa_i_in_j = self._ioa(boxes[i], boxes[j])
254
+ if ioa_i_in_j >= self.SMALL_CONTAINED_IOA:
255
+ keep[i] = False
256
+ break
257
+ else:
258
+ ratio = aj / ai
259
+ if ratio <= self.SMALL_RATIO_MAX:
260
+ ioa_j_in_i = self._ioa(boxes[j], boxes[i])
261
+ if ioa_j_in_i >= self.SMALL_CONTAINED_IOA:
262
+ keep[j] = False
263
+ return [bb for bb, k in zip(boxes, keep) if k]
264
+
265
+ def _detect_objects_batch(self, batch_images: List[ndarray], offset: int) -> Dict[int, List[BoundingBox]]:
266
+ """
267
+ Phase 1: Object detection for all frames in batch.
268
+ Returns detected objects with players still having class_id=2 (before team classification).
269
+
270
+ Args:
271
+ batch_images: List of images to process
272
+ offset: Frame offset for numbering
273
+
274
+ Returns:
275
+ Dictionary mapping frame_id to list of detected boxes
276
+ """
277
+ bboxes: Dict[int, List[BoundingBox]] = {}
278
+
279
+ if len(batch_images) == 0:
280
+ return bboxes
281
+
282
+ print(f"Processing batch of {len(batch_images)} images")
283
+
284
+ # Get original image dimensions for scaling
285
+ height, width = batch_images[0].shape[:2]
286
+ scale = 640.0
287
+ x_scale = width / scale
288
+ y_scale = height / scale
289
+
290
+ # Memory optimization: Process smaller batches if needed
291
+ max_batch_size = 32 # Reduce batch size further to prevent memory issues
292
+ if len(batch_images) > max_batch_size:
293
+ print(f"Large batch detected ({len(batch_images)} images), splitting into smaller batches of {max_batch_size}")
294
+ # Process in smaller chunks
295
+ all_bboxes = {}
296
+ for chunk_start in range(0, len(batch_images), max_batch_size):
297
+ chunk_end = min(chunk_start + max_batch_size, len(batch_images))
298
+ chunk_images = batch_images[chunk_start:chunk_end]
299
+ chunk_offset = offset + chunk_start
300
+ print(f"Processing chunk {chunk_start//max_batch_size + 1}: images {chunk_start}-{chunk_end-1}")
301
+ chunk_bboxes = self._detect_objects_batch(chunk_images, chunk_offset)
302
+ all_bboxes.update(chunk_bboxes)
303
+ return all_bboxes
304
+
305
+ # Preprocess images for ONNX inference
306
+ imgs = self._pre_process_img(batch_images, scale)
307
+ actual_batch_size = len(batch_images)
308
+
309
+ # Handle batch size mismatch - pad if needed
310
+ model_batch_size = self.bbox_model.get_inputs()[0].shape[0]
311
+ print(f"Model input shape: {self.bbox_model.get_inputs()[0].shape}, batch_size: {model_batch_size}")
312
+
313
+ if model_batch_size is not None:
314
+ try:
315
+ # Handle dynamic batch size (None, -1, 'None')
316
+ if str(model_batch_size) in ['None', '-1'] or model_batch_size == -1:
317
+ model_batch_size = None
318
+ else:
319
+ model_batch_size = int(model_batch_size)
320
+ except (ValueError, TypeError):
321
+ model_batch_size = None
322
+
323
+ print(f"Processed model_batch_size: {model_batch_size}, actual_batch_size: {actual_batch_size}")
324
+
325
+ if model_batch_size and actual_batch_size < model_batch_size:
326
+ padding_size = model_batch_size - actual_batch_size
327
+ dummy_img = np.zeros((1, 3, int(scale), int(scale)), dtype=np.float32)
328
+ padding = np.repeat(dummy_img, padding_size, axis=0)
329
+ imgs = np.vstack([imgs, padding])
330
+
331
+ # ONNX inference with error handling
332
+ try:
333
+ input_name = self.bbox_model.get_inputs()[0].name
334
+ import time
335
+ start_time = time.time()
336
+ outputs = self.bbox_model.run(None, {input_name: imgs})[0]
337
+ inference_time = time.time() - start_time
338
+ print(f"Inference time: {inference_time:.3f}s for {actual_batch_size} images")
339
+
340
+ # Remove padded results if we added padding
341
+ if model_batch_size and isinstance(model_batch_size, int) and actual_batch_size < model_batch_size:
342
+ outputs = outputs[:actual_batch_size]
343
+
344
+ # Post-process outputs to get detections
345
+ raw_results = self._post_process_output(np.array(outputs), x_scale, y_scale)
346
+
347
+ except Exception as e:
348
+ print(f"Error during ONNX inference: {e}")
349
+ return bboxes
350
+
351
+ if not raw_results:
352
+ return bboxes
353
+
354
+ # Convert raw results to BoundingBox objects and apply processing
355
+ for frame_idx_in_batch, frame_detections in enumerate(raw_results):
356
+ if not frame_detections:
357
+ continue
358
+
359
+ # Convert to BoundingBox objects
360
+ boxes: List[BoundingBox] = []
361
+ for box, conf, cls_id in frame_detections:
362
+ x1, y1, x2, y2 = box
363
+ if int(cls_id) < 4:
364
+ boxes.append(
365
+ BoundingBox(
366
+ x1=int(x1),
367
+ y1=int(y1),
368
+ x2=int(x2),
369
+ y2=int(y2),
370
+ cls_id=int(cls_id),
371
+ conf=float(conf),
372
+ )
373
+ )
374
+
375
+ # Handle footballs - keep only the best one
376
+ footballs = [bb for bb in boxes if int(bb.cls_id) == 0]
377
+ if len(footballs) > 1:
378
+ best_ball = max(footballs, key=lambda b: b.conf)
379
+ boxes = [bb for bb in boxes if int(bb.cls_id) != 0]
380
+ boxes.append(best_ball)
381
+
382
+ # Remove overlapping small boxes
383
+ boxes = suppress_small_contained_boxes(boxes, self.SMALL_CONTAINED_IOA, self.SMALL_RATIO_MAX)
384
+
385
+ # Handle goalkeeper detection issues:
386
+ # 1. Fix misplaced goalkeepers (convert to players if standing in middle)
387
+ # 2. Allow up to 2 goalkeepers maximum (one from each team)
388
+ # Goalkeepers remain class_id = 1 (no team assignment)
389
+ boxes = self._handle_multiple_goalkeepers(boxes)
390
+
391
+ # Store results (players still have class_id=2, will be classified in phase 2)
392
+ frame_id = offset + frame_idx_in_batch
393
+ bboxes[frame_id] = boxes
394
+
395
+ return bboxes
396
+
397
+
398
+ def predict_batch(
399
+ self,
400
+ batch_images: List[ndarray],
401
+ offset: int,
402
+ n_keypoints: int,
403
+ task_type: Optional[str] = None,
404
+ ) -> List[TVFrameResult]:
405
+ process_objects = task_type is None or task_type == "object"
406
+ process_keypoints = task_type is None or task_type == "keypoint"
407
+
408
+ # Phase 1: Object Detection for all frames
409
+ bboxes: Dict[int, List[BoundingBox]] = {}
410
+ if process_objects:
411
+ bboxes = self._detect_objects_batch(batch_images, offset)
412
+
413
+ import time
414
+ time_start = time.time()
415
+ # Phase 2: Team Classification for all detected players
416
+ if process_objects and bboxes:
417
+ bboxes, self.team_classifier_fitted, self.player_crops_for_fit = classify_teams_batch(
418
+ self.team_classifier,
419
+ self.team_classifier_fitted,
420
+ self.player_crops_for_fit,
421
+ batch_images,
422
+ bboxes,
423
+ offset,
424
+ self.MIN_SAMPLES_FOR_FIT,
425
+ self.MAX_SAMPLES_FOR_FIT,
426
+ self.SINGLE_PLAYER_HUE_PIVOT
427
+ )
428
+ self.team_classifier_fitted = False
429
+ self.player_crops_for_fit = []
430
+ print(f"Time Team Classification: {time.time() - time_start} s")
431
+
432
+ # Phase 3: Keypoint Detection
433
+ keypoints: Dict[int, List[Tuple[int, int]]] = {}
434
+ if process_keypoints:
435
+ keypoints = self._detect_keypoints_batch(batch_images, offset, n_keypoints)
436
+
437
+ # Phase 4: Combine results
438
+ results: List[TVFrameResult] = []
439
+ for frame_number in range(offset, offset + len(batch_images)):
440
+ results.append(
441
+ TVFrameResult(
442
+ frame_id=frame_number,
443
+ boxes=bboxes.get(frame_number, []),
444
+ keypoints=keypoints.get(
445
+ frame_number,
446
+ [(0, 0) for _ in range(n_keypoints)],
447
+ ),
448
+ )
449
+ )
450
+ return results
451
+
452
+ def _detect_keypoints_batch(self, batch_images: List[ndarray],
453
+ offset: int, n_keypoints: int) -> Dict[int, List[Tuple[int, int]]]:
454
+ """
455
+ Phase 3: Keypoint detection for all frames in batch.
456
+
457
+ Args:
458
+ batch_images: List of images to process
459
+ offset: Frame offset for numbering
460
+ n_keypoints: Number of keypoints expected
461
+
462
+ Returns:
463
+ Dictionary mapping frame_id to list of keypoint coordinates
464
+ """
465
+ keypoints: Dict[int, List[Tuple[int, int]]] = {}
466
+ keypoints_model_results = self.keypoints_model.predict(batch_images)
467
+
468
+ if keypoints_model_results is None:
469
+ return keypoints
470
+
471
+ for frame_idx_in_batch, detection in enumerate(keypoints_model_results):
472
+ if not hasattr(detection, "keypoints") or detection.keypoints is None:
473
+ continue
474
+
475
+ # Extract keypoints with confidence
476
+ frame_keypoints_with_conf: List[Tuple[int, int, float]] = []
477
+ for i, part_points in enumerate(detection.keypoints.data):
478
+ for k_id, (x, y, _) in enumerate(part_points):
479
+ confidence = float(detection.keypoints.conf[i][k_id])
480
+ frame_keypoints_with_conf.append((int(x), int(y), confidence))
481
+
482
+ # Pad or truncate to expected number of keypoints
483
+ if len(frame_keypoints_with_conf) < n_keypoints:
484
+ frame_keypoints_with_conf.extend(
485
+ [(0, 0, 0.0)] * (n_keypoints - len(frame_keypoints_with_conf))
486
+ )
487
+ else:
488
+ frame_keypoints_with_conf = frame_keypoints_with_conf[:n_keypoints]
489
+
490
+ # Filter keypoints based on confidence thresholds
491
+ filtered_keypoints: List[Tuple[int, int]] = []
492
+ for idx, (x, y, confidence) in enumerate(frame_keypoints_with_conf):
493
+ if idx in self.CORNER_INDICES:
494
+ # Corner keypoints have lower confidence threshold
495
+ if confidence < 0.3:
496
+ filtered_keypoints.append((0, 0))
497
+ else:
498
+ filtered_keypoints.append((int(x), int(y)))
499
+ else:
500
+ # Regular keypoints
501
+ if confidence < 0.5:
502
+ filtered_keypoints.append((0, 0))
503
+ else:
504
+ filtered_keypoints.append((int(x), int(y)))
505
+
506
+ frame_id = offset + frame_idx_in_batch
507
+ keypoints[frame_id] = filtered_keypoints
508
+
509
+ return keypoints
object-detection.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:05112479be8cb59494e9ae23a57af43becd5aa1f448b0e5ed33fcb6b4c2bbbc3
3
+ size 273322667
osnet_ain.pyc ADDED
Binary file (24.2 kB). View file
 
osnet_model.pth.tar-100 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:64873ef0e8abf28df31facd113f27634e2d085a2dcf8d19123409b1d0e2566c8
3
+ size 36189526
pitch.py ADDED
@@ -0,0 +1,687 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import absolute_import
2
+ from __future__ import division
3
+ from __future__ import print_function
4
+
5
+ import os
6
+ import sys
7
+ import time
8
+ from typing import List, Optional, Tuple
9
+
10
+ import cv2
11
+ import numpy as np
12
+ import torch
13
+ import torch.nn as nn
14
+ import torch.nn.functional as F
15
+ import torchvision.transforms as T
16
+ import torchvision.transforms.functional as f
17
+ from pydantic import BaseModel
18
+
19
+ import logging
20
+ logger = logging.getLogger(__name__)
21
+
22
+
23
+ class BoundingBox(BaseModel):
24
+ x1: int
25
+ y1: int
26
+ x2: int
27
+ y2: int
28
+ cls_id: int
29
+ conf: float
30
+
31
+
32
+ class TVFrameResult(BaseModel):
33
+ frame_id: int
34
+ boxes: list[BoundingBox]
35
+ keypoints: list[tuple[int, int]]
36
+
37
+ BatchNorm2d = nn.BatchNorm2d
38
+ BN_MOMENTUM = 0.1
39
+
40
+ def conv3x3(in_planes, out_planes, stride=1):
41
+ """3x3 convolution with padding"""
42
+ return nn.Conv2d(in_planes, out_planes, kernel_size=3,
43
+ stride=stride, padding=1, bias=False)
44
+
45
+
46
+ class BasicBlock(nn.Module):
47
+ expansion = 1
48
+
49
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
50
+ super(BasicBlock, self).__init__()
51
+ self.conv1 = conv3x3(inplanes, planes, stride)
52
+ self.bn1 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
53
+ self.relu = nn.ReLU(inplace=True)
54
+ self.conv2 = conv3x3(planes, planes)
55
+ self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
56
+ self.downsample = downsample
57
+ self.stride = stride
58
+
59
+ def forward(self, x):
60
+ residual = x
61
+
62
+ out = self.conv1(x)
63
+ out = self.bn1(out)
64
+ out = self.relu(out)
65
+
66
+ out = self.conv2(out)
67
+ out = self.bn2(out)
68
+
69
+ if self.downsample is not None:
70
+ residual = self.downsample(x)
71
+
72
+ out += residual
73
+ out = self.relu(out)
74
+
75
+ return out
76
+
77
+
78
+ class Bottleneck(nn.Module):
79
+ expansion = 4
80
+
81
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
82
+ super(Bottleneck, self).__init__()
83
+ self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
84
+ self.bn1 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
85
+ self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
86
+ padding=1, bias=False)
87
+ self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
88
+ self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
89
+ bias=False)
90
+ self.bn3 = BatchNorm2d(planes * self.expansion,
91
+ momentum=BN_MOMENTUM)
92
+ self.relu = nn.ReLU(inplace=True)
93
+ self.downsample = downsample
94
+ self.stride = stride
95
+
96
+ def forward(self, x):
97
+ residual = x
98
+
99
+ out = self.conv1(x)
100
+ out = self.bn1(out)
101
+ out = self.relu(out)
102
+
103
+ out = self.conv2(out)
104
+ out = self.bn2(out)
105
+ out = self.relu(out)
106
+
107
+ out = self.conv3(out)
108
+ out = self.bn3(out)
109
+
110
+ if self.downsample is not None:
111
+ residual = self.downsample(x)
112
+
113
+ out += residual
114
+ out = self.relu(out)
115
+
116
+ return out
117
+
118
+
119
+ class HighResolutionModule(nn.Module):
120
+ def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
121
+ num_channels, fuse_method, multi_scale_output=True):
122
+ super(HighResolutionModule, self).__init__()
123
+ self._check_branches(
124
+ num_branches, blocks, num_blocks, num_inchannels, num_channels)
125
+
126
+ self.num_inchannels = num_inchannels
127
+ self.fuse_method = fuse_method
128
+ self.num_branches = num_branches
129
+
130
+ self.multi_scale_output = multi_scale_output
131
+
132
+ self.branches = self._make_branches(
133
+ num_branches, blocks, num_blocks, num_channels)
134
+ self.fuse_layers = self._make_fuse_layers()
135
+ self.relu = nn.ReLU(inplace=True)
136
+
137
+ def _check_branches(self, num_branches, blocks, num_blocks,
138
+ num_inchannels, num_channels):
139
+ if num_branches != len(num_blocks):
140
+ error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(
141
+ num_branches, len(num_blocks))
142
+ logger.error(error_msg)
143
+ raise ValueError(error_msg)
144
+
145
+ if num_branches != len(num_channels):
146
+ error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(
147
+ num_branches, len(num_channels))
148
+ logger.error(error_msg)
149
+ raise ValueError(error_msg)
150
+
151
+ if num_branches != len(num_inchannels):
152
+ error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(
153
+ num_branches, len(num_inchannels))
154
+ logger.error(error_msg)
155
+ raise ValueError(error_msg)
156
+
157
+ def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
158
+ stride=1):
159
+ downsample = None
160
+ if stride != 1 or \
161
+ self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
162
+ downsample = nn.Sequential(
163
+ nn.Conv2d(self.num_inchannels[branch_index],
164
+ num_channels[branch_index] * block.expansion,
165
+ kernel_size=1, stride=stride, bias=False),
166
+ BatchNorm2d(num_channels[branch_index] * block.expansion,
167
+ momentum=BN_MOMENTUM),
168
+ )
169
+
170
+ layers = []
171
+ layers.append(block(self.num_inchannels[branch_index],
172
+ num_channels[branch_index], stride, downsample))
173
+ self.num_inchannels[branch_index] = \
174
+ num_channels[branch_index] * block.expansion
175
+ for i in range(1, num_blocks[branch_index]):
176
+ layers.append(block(self.num_inchannels[branch_index],
177
+ num_channels[branch_index]))
178
+
179
+ return nn.Sequential(*layers)
180
+
181
+ def _make_branches(self, num_branches, block, num_blocks, num_channels):
182
+ branches = []
183
+
184
+ for i in range(num_branches):
185
+ branches.append(
186
+ self._make_one_branch(i, block, num_blocks, num_channels))
187
+
188
+ return nn.ModuleList(branches)
189
+
190
+ def _make_fuse_layers(self):
191
+ if self.num_branches == 1:
192
+ return None
193
+
194
+ num_branches = self.num_branches
195
+ num_inchannels = self.num_inchannels
196
+ fuse_layers = []
197
+ for i in range(num_branches if self.multi_scale_output else 1):
198
+ fuse_layer = []
199
+ for j in range(num_branches):
200
+ if j > i:
201
+ fuse_layer.append(nn.Sequential(
202
+ nn.Conv2d(num_inchannels[j],
203
+ num_inchannels[i],
204
+ 1,
205
+ 1,
206
+ 0,
207
+ bias=False),
208
+ BatchNorm2d(num_inchannels[i], momentum=BN_MOMENTUM)))
209
+ # nn.Upsample(scale_factor=2**(j-i), mode='nearest')))
210
+ elif j == i:
211
+ fuse_layer.append(None)
212
+ else:
213
+ conv3x3s = []
214
+ for k in range(i - j):
215
+ if k == i - j - 1:
216
+ num_outchannels_conv3x3 = num_inchannels[i]
217
+ conv3x3s.append(nn.Sequential(
218
+ nn.Conv2d(num_inchannels[j],
219
+ num_outchannels_conv3x3,
220
+ 3, 2, 1, bias=False),
221
+ BatchNorm2d(num_outchannels_conv3x3, momentum=BN_MOMENTUM)))
222
+ else:
223
+ num_outchannels_conv3x3 = num_inchannels[j]
224
+ conv3x3s.append(nn.Sequential(
225
+ nn.Conv2d(num_inchannels[j],
226
+ num_outchannels_conv3x3,
227
+ 3, 2, 1, bias=False),
228
+ BatchNorm2d(num_outchannels_conv3x3,
229
+ momentum=BN_MOMENTUM),
230
+ nn.ReLU(inplace=True)))
231
+ fuse_layer.append(nn.Sequential(*conv3x3s))
232
+ fuse_layers.append(nn.ModuleList(fuse_layer))
233
+
234
+ return nn.ModuleList(fuse_layers)
235
+
236
+ def get_num_inchannels(self):
237
+ return self.num_inchannels
238
+
239
+ def forward(self, x):
240
+ if self.num_branches == 1:
241
+ return [self.branches[0](x[0])]
242
+
243
+ for i in range(self.num_branches):
244
+ x[i] = self.branches[i](x[i])
245
+
246
+ x_fuse = []
247
+ for i in range(len(self.fuse_layers)):
248
+ y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
249
+ for j in range(1, self.num_branches):
250
+ if i == j:
251
+ y = y + x[j]
252
+ elif j > i:
253
+ y = y + F.interpolate(
254
+ self.fuse_layers[i][j](x[j]),
255
+ size=[x[i].shape[2], x[i].shape[3]],
256
+ mode='bilinear')
257
+ else:
258
+ y = y + self.fuse_layers[i][j](x[j])
259
+ x_fuse.append(self.relu(y))
260
+
261
+ return x_fuse
262
+
263
+
264
+ blocks_dict = {
265
+ 'BASIC': BasicBlock,
266
+ 'BOTTLENECK': Bottleneck
267
+ }
268
+
269
+
270
+ class HighResolutionNet(nn.Module):
271
+
272
+ def __init__(self, config, **kwargs):
273
+ self.inplanes = 64
274
+ extra = config['MODEL']['EXTRA']
275
+ super(HighResolutionNet, self).__init__()
276
+
277
+ # stem net
278
+ self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=2, padding=1,
279
+ bias=False)
280
+ self.bn1 = BatchNorm2d(self.inplanes, momentum=BN_MOMENTUM)
281
+ self.conv2 = nn.Conv2d(self.inplanes, self.inplanes, kernel_size=3, stride=2, padding=1,
282
+ bias=False)
283
+ self.bn2 = BatchNorm2d(self.inplanes, momentum=BN_MOMENTUM)
284
+ self.relu = nn.ReLU(inplace=True)
285
+ self.sf = nn.Softmax(dim=1)
286
+ self.layer1 = self._make_layer(Bottleneck, 64, 64, 4)
287
+
288
+ self.stage2_cfg = extra['STAGE2']
289
+ num_channels = self.stage2_cfg['NUM_CHANNELS']
290
+ block = blocks_dict[self.stage2_cfg['BLOCK']]
291
+ num_channels = [
292
+ num_channels[i] * block.expansion for i in range(len(num_channels))]
293
+ self.transition1 = self._make_transition_layer(
294
+ [256], num_channels)
295
+ self.stage2, pre_stage_channels = self._make_stage(
296
+ self.stage2_cfg, num_channels)
297
+
298
+ self.stage3_cfg = extra['STAGE3']
299
+ num_channels = self.stage3_cfg['NUM_CHANNELS']
300
+ block = blocks_dict[self.stage3_cfg['BLOCK']]
301
+ num_channels = [
302
+ num_channels[i] * block.expansion for i in range(len(num_channels))]
303
+ self.transition2 = self._make_transition_layer(
304
+ pre_stage_channels, num_channels)
305
+ self.stage3, pre_stage_channels = self._make_stage(
306
+ self.stage3_cfg, num_channels)
307
+
308
+ self.stage4_cfg = extra['STAGE4']
309
+ num_channels = self.stage4_cfg['NUM_CHANNELS']
310
+ block = blocks_dict[self.stage4_cfg['BLOCK']]
311
+ num_channels = [
312
+ num_channels[i] * block.expansion for i in range(len(num_channels))]
313
+ self.transition3 = self._make_transition_layer(
314
+ pre_stage_channels, num_channels)
315
+ self.stage4, pre_stage_channels = self._make_stage(
316
+ self.stage4_cfg, num_channels, multi_scale_output=True)
317
+
318
+ self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
319
+ final_inp_channels = sum(pre_stage_channels) + self.inplanes
320
+
321
+ self.head = nn.Sequential(nn.Sequential(
322
+ nn.Conv2d(
323
+ in_channels=final_inp_channels,
324
+ out_channels=final_inp_channels,
325
+ kernel_size=1),
326
+ BatchNorm2d(final_inp_channels, momentum=BN_MOMENTUM),
327
+ nn.ReLU(inplace=True),
328
+ nn.Conv2d(
329
+ in_channels=final_inp_channels,
330
+ out_channels=config['MODEL']['NUM_JOINTS'],
331
+ kernel_size=extra['FINAL_CONV_KERNEL']),
332
+ nn.Softmax(dim=1)))
333
+
334
+
335
+
336
+ def _make_head(self, x, x_skip):
337
+ x = self.upsample(x)
338
+ x = torch.cat([x, x_skip], dim=1)
339
+ x = self.head(x)
340
+
341
+ return x
342
+
343
+ def _make_transition_layer(
344
+ self, num_channels_pre_layer, num_channels_cur_layer):
345
+ num_branches_cur = len(num_channels_cur_layer)
346
+ num_branches_pre = len(num_channels_pre_layer)
347
+
348
+ transition_layers = []
349
+ for i in range(num_branches_cur):
350
+ if i < num_branches_pre:
351
+ if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
352
+ transition_layers.append(nn.Sequential(
353
+ nn.Conv2d(num_channels_pre_layer[i],
354
+ num_channels_cur_layer[i],
355
+ 3,
356
+ 1,
357
+ 1,
358
+ bias=False),
359
+ BatchNorm2d(
360
+ num_channels_cur_layer[i], momentum=BN_MOMENTUM),
361
+ nn.ReLU(inplace=True)))
362
+ else:
363
+ transition_layers.append(None)
364
+ else:
365
+ conv3x3s = []
366
+ for j in range(i + 1 - num_branches_pre):
367
+ inchannels = num_channels_pre_layer[-1]
368
+ outchannels = num_channels_cur_layer[i] \
369
+ if j == i - num_branches_pre else inchannels
370
+ conv3x3s.append(nn.Sequential(
371
+ nn.Conv2d(
372
+ inchannels, outchannels, 3, 2, 1, bias=False),
373
+ BatchNorm2d(outchannels, momentum=BN_MOMENTUM),
374
+ nn.ReLU(inplace=True)))
375
+ transition_layers.append(nn.Sequential(*conv3x3s))
376
+
377
+ return nn.ModuleList(transition_layers)
378
+
379
+ def _make_layer(self, block, inplanes, planes, blocks, stride=1):
380
+ downsample = None
381
+ if stride != 1 or inplanes != planes * block.expansion:
382
+ downsample = nn.Sequential(
383
+ nn.Conv2d(inplanes, planes * block.expansion,
384
+ kernel_size=1, stride=stride, bias=False),
385
+ BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
386
+ )
387
+
388
+ layers = []
389
+ layers.append(block(inplanes, planes, stride, downsample))
390
+ inplanes = planes * block.expansion
391
+ for i in range(1, blocks):
392
+ layers.append(block(inplanes, planes))
393
+
394
+ return nn.Sequential(*layers)
395
+
396
+ def _make_stage(self, layer_config, num_inchannels,
397
+ multi_scale_output=True):
398
+ num_modules = layer_config['NUM_MODULES']
399
+ num_branches = layer_config['NUM_BRANCHES']
400
+ num_blocks = layer_config['NUM_BLOCKS']
401
+ num_channels = layer_config['NUM_CHANNELS']
402
+ block = blocks_dict[layer_config['BLOCK']]
403
+ fuse_method = layer_config['FUSE_METHOD']
404
+
405
+ modules = []
406
+ for i in range(num_modules):
407
+ # multi_scale_output is only used last module
408
+ if not multi_scale_output and i == num_modules - 1:
409
+ reset_multi_scale_output = False
410
+ else:
411
+ reset_multi_scale_output = True
412
+ modules.append(
413
+ HighResolutionModule(num_branches,
414
+ block,
415
+ num_blocks,
416
+ num_inchannels,
417
+ num_channels,
418
+ fuse_method,
419
+ reset_multi_scale_output)
420
+ )
421
+ num_inchannels = modules[-1].get_num_inchannels()
422
+
423
+ return nn.Sequential(*modules), num_inchannels
424
+
425
+ def forward(self, x):
426
+ # h, w = x.size(2), x.size(3)
427
+ x = self.conv1(x)
428
+ x_skip = x.clone()
429
+ x = self.bn1(x)
430
+ x = self.relu(x)
431
+ x = self.conv2(x)
432
+ x = self.bn2(x)
433
+ x = self.relu(x)
434
+ x = self.layer1(x)
435
+
436
+ x_list = []
437
+ for i in range(self.stage2_cfg['NUM_BRANCHES']):
438
+ if self.transition1[i] is not None:
439
+ x_list.append(self.transition1[i](x))
440
+ else:
441
+ x_list.append(x)
442
+ y_list = self.stage2(x_list)
443
+
444
+ x_list = []
445
+ for i in range(self.stage3_cfg['NUM_BRANCHES']):
446
+ if self.transition2[i] is not None:
447
+ x_list.append(self.transition2[i](y_list[-1]))
448
+ else:
449
+ x_list.append(y_list[i])
450
+ y_list = self.stage3(x_list)
451
+
452
+ x_list = []
453
+ for i in range(self.stage4_cfg['NUM_BRANCHES']):
454
+ if self.transition3[i] is not None:
455
+ x_list.append(self.transition3[i](y_list[-1]))
456
+ else:
457
+ x_list.append(y_list[i])
458
+ x = self.stage4(x_list)
459
+
460
+ # Head Part
461
+ height, width = x[0].size(2), x[0].size(3)
462
+ x1 = F.interpolate(x[1], size=(height, width), mode='bilinear', align_corners=False)
463
+ x2 = F.interpolate(x[2], size=(height, width), mode='bilinear', align_corners=False)
464
+ x3 = F.interpolate(x[3], size=(height, width), mode='bilinear', align_corners=False)
465
+ x = torch.cat([x[0], x1, x2, x3], 1)
466
+ x = self._make_head(x, x_skip)
467
+
468
+ return x
469
+
470
+ def init_weights(self, pretrained=''):
471
+ for m in self.modules():
472
+ if isinstance(m, nn.Conv2d):
473
+ nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
474
+ #nn.init.normal_(m.weight, std=0.001)
475
+ #nn.init.constant_(m.bias, 0)
476
+ elif isinstance(m, nn.BatchNorm2d):
477
+ nn.init.constant_(m.weight, 1)
478
+ nn.init.constant_(m.bias, 0)
479
+ if pretrained != '':
480
+ if os.path.isfile(pretrained):
481
+ pretrained_dict = torch.load(pretrained)
482
+ model_dict = self.state_dict()
483
+ pretrained_dict = {k: v for k, v in pretrained_dict.items()
484
+ if k in model_dict.keys()}
485
+ model_dict.update(pretrained_dict)
486
+ self.load_state_dict(model_dict)
487
+ else:
488
+ sys.exit(f'Weights {pretrained} not found.')
489
+
490
+
491
+ def get_cls_net(config, pretrained='', **kwargs):
492
+ """Create keypoint detection model with softmax activation"""
493
+ model = HighResolutionNet(config, **kwargs)
494
+ model.init_weights(pretrained)
495
+ return model
496
+
497
+
498
+ def get_cls_net_l(config, pretrained='', **kwargs):
499
+ """Create line detection model with sigmoid activation"""
500
+ model = HighResolutionNet(config, **kwargs)
501
+ model.init_weights(pretrained)
502
+
503
+ # After loading weights, replace just the activation function
504
+ # The saved model expects the nested Sequential structure
505
+ inner_seq = model.head[0]
506
+ # Replace softmax (index 4) with sigmoid
507
+ model.head[0][4] = nn.Sigmoid()
508
+
509
+ return model
510
+
511
+ # Simplified utility functions - removed complex Gaussian generation functions
512
+ # These were mainly used for training data generation, not inference
513
+
514
+
515
+
516
+ # generate_gaussian_array_vectorized_dist_l function removed - not used in current implementation
517
+ @torch.inference_mode()
518
+ def run_inference(model, input_tensor: torch.Tensor, device):
519
+ input_tensor = input_tensor.to(device).to(memory_format=torch.channels_last)
520
+ output = model.module().forward(input_tensor)
521
+ return output
522
+
523
+ def preprocess_batch_fast(frames):
524
+ """Ultra-fast batch preprocessing using optimized tensor operations"""
525
+ target_size = (540, 960) # H, W format for model input
526
+ batch = []
527
+ for i, frame in enumerate(frames):
528
+ frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
529
+ img = cv2.resize(frame_rgb, (target_size[1], target_size[0]))
530
+ img = img.astype(np.float32) / 255.0
531
+ img = np.transpose(img, (2, 0, 1)) # HWC -> CHW
532
+ batch.append(img)
533
+ batch = torch.from_numpy(np.stack(batch)).float()
534
+
535
+ return batch
536
+
537
+ def extract_keypoints_from_heatmap(heatmap: torch.Tensor, scale: int = 2, max_keypoints: int = 1):
538
+ """Optimized keypoint extraction from heatmaps"""
539
+ batch_size, n_channels, height, width = heatmap.shape
540
+
541
+ # Find local maxima using max pooling (keep on GPU)
542
+ kernel = 3
543
+ pad = 1
544
+ max_pooled = F.max_pool2d(heatmap, kernel, stride=1, padding=pad)
545
+ local_maxima = (max_pooled == heatmap)
546
+ heatmap = heatmap * local_maxima
547
+
548
+ # Get top keypoints (keep on GPU longer)
549
+ scores, indices = torch.topk(heatmap.view(batch_size, n_channels, -1), max_keypoints, sorted=False)
550
+ y_coords = torch.div(indices, width, rounding_mode="floor")
551
+ x_coords = indices % width
552
+
553
+ # Optimized tensor operations
554
+ x_coords = x_coords * scale
555
+ y_coords = y_coords * scale
556
+
557
+ # Create result tensor directly on GPU
558
+ results = torch.stack([x_coords.float(), y_coords.float(), scores], dim=-1)
559
+
560
+ return results
561
+
562
+
563
+ def extract_keypoints_from_heatmap_fast(heatmap: torch.Tensor, scale: int = 2, max_keypoints: int = 1):
564
+ """Ultra-fast keypoint extraction optimized for speed"""
565
+ batch_size, n_channels, height, width = heatmap.shape
566
+
567
+ # Simplified local maxima detection (faster but slightly less accurate)
568
+ max_pooled = F.max_pool2d(heatmap, 3, stride=1, padding=1)
569
+ local_maxima = (max_pooled == heatmap)
570
+
571
+ # Apply mask and get top keypoints in one go
572
+ masked_heatmap = heatmap * local_maxima
573
+ flat_heatmap = masked_heatmap.view(batch_size, n_channels, -1)
574
+ scores, indices = torch.topk(flat_heatmap, max_keypoints, dim=-1, sorted=False)
575
+
576
+ # Vectorized coordinate calculation
577
+ y_coords = torch.div(indices, width, rounding_mode="floor") * scale
578
+ x_coords = (indices % width) * scale
579
+
580
+ # Stack results efficiently
581
+ results = torch.stack([x_coords.float(), y_coords.float(), scores], dim=-1)
582
+ return results
583
+
584
+
585
+ def process_keypoints_vectorized(kp_coords, kp_threshold, w, h, batch_size):
586
+ """Ultra-fast vectorized keypoint processing"""
587
+ batch_results = []
588
+
589
+ # Convert to numpy once for faster CPU operations
590
+ kp_np = kp_coords.cpu().numpy()
591
+
592
+ for batch_idx in range(batch_size):
593
+ kp_dict = {}
594
+ # Vectorized threshold check
595
+ valid_kps = kp_np[batch_idx, :, 0, 2] > kp_threshold
596
+ valid_indices = np.where(valid_kps)[0]
597
+
598
+ for ch_idx in valid_indices:
599
+ x = float(kp_np[batch_idx, ch_idx, 0, 0]) / w
600
+ y = float(kp_np[batch_idx, ch_idx, 0, 1]) / h
601
+ p = float(kp_np[batch_idx, ch_idx, 0, 2])
602
+ kp_dict[ch_idx + 1] = {'x': x, 'y': y, 'p': p}
603
+
604
+ batch_results.append(kp_dict)
605
+
606
+ return batch_results
607
+
608
+ def inference_batch(frames, model, kp_threshold, device, batch_size=8):
609
+ """Optimized batch inference for multiple frames"""
610
+ results = []
611
+ num_frames = len(frames)
612
+
613
+ # Get the device from the model itself
614
+ model_device = next(model.parameters()).device
615
+
616
+ # Process all frames in optimally-sized batches
617
+ for i in range(0, num_frames, batch_size):
618
+ current_batch_size = min(batch_size, num_frames - i)
619
+ batch_frames = frames[i:i + current_batch_size]
620
+
621
+ # Fast preprocessing - create on CPU first
622
+ batch = preprocess_batch_fast(batch_frames)
623
+ b, c, h, w = batch.size()
624
+
625
+ # Move batch to model device
626
+ batch = batch.to(model_device)
627
+
628
+ with torch.no_grad():
629
+ heatmaps = model(batch)
630
+
631
+ # Ultra-fast keypoint extraction
632
+ kp_coords = extract_keypoints_from_heatmap_fast(heatmaps[:,:-1,:,:], scale=2, max_keypoints=1)
633
+
634
+ # Vectorized batch processing - no loops
635
+ batch_results = process_keypoints_vectorized(kp_coords, kp_threshold, 960, 540, current_batch_size)
636
+ results.extend(batch_results)
637
+
638
+ # Minimal cleanup
639
+ del heatmaps, kp_coords, batch
640
+
641
+ return results
642
+
643
+ # Keypoint mapping from detection indices to standard football pitch keypoint IDs
644
+ map_keypoints = {
645
+ 1: 1, 2: 14, 3: 25, 4: 2, 5: 10, 6: 18, 7: 26, 8: 3, 9: 7, 10: 23,
646
+ 11: 27, 20: 4, 21: 8, 22: 24, 23: 28, 24: 5, 25: 13, 26: 21, 27: 29,
647
+ 28: 6, 29: 17, 30: 30, 31: 11, 32: 15, 33: 19, 34: 12, 35: 16, 36: 20,
648
+ 45: 9, 50: 31, 52: 32, 57: 22
649
+ }
650
+
651
+ def get_mapped_keypoints(kp_points):
652
+ """Apply keypoint mapping to detection results"""
653
+ mapped_points = {}
654
+ for key, value in kp_points.items():
655
+ if key in map_keypoints:
656
+ mapped_key = map_keypoints[key]
657
+ mapped_points[mapped_key] = value
658
+ # else:
659
+ # Keep unmapped keypoints with original key
660
+ # mapped_points[key] = value
661
+ return mapped_points
662
+
663
+ def process_batch_input(frames, model, kp_threshold, device, batch_size=8):
664
+ """Process multiple input images in batch"""
665
+ # Batch inference
666
+ kp_results = inference_batch(frames, model, kp_threshold, device, batch_size)
667
+ kp_results = [get_mapped_keypoints(kp) for kp in kp_results]
668
+ # Draw results and save
669
+ # for i, (frame, kp_points, input_path) in enumerate(zip(frames, kp_results, valid_paths)):
670
+ # height, width = frame.shape[:2]
671
+
672
+ # # Apply mapping to get standard keypoint IDs
673
+ # mapped_kp_points = get_mapped_keypoints(kp_points)
674
+
675
+ # for key, value in mapped_kp_points.items():
676
+ # x = int(value['x'] * width)
677
+ # y = int(value['y'] * height)
678
+ # cv2.circle(frame, (x, y), 5, (0, 255, 0), -1) # Green circles
679
+ # cv2.putText(frame, str(key), (x+10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
680
+
681
+ # # Save result
682
+ # output_path = input_path.replace('.png', '_result.png').replace('.jpg', '_result.jpg')
683
+ # cv2.imwrite(output_path, frame)
684
+
685
+ # print(f"Batch processing complete. Processed {len(frames)} images.")
686
+
687
+ return kp_results
player.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ce9fc31f61e6f156f786077abb8eef36b0836bda1ef07d1d0ba82d43ae0ecd0b
3
+ size 22540152
player.py ADDED
@@ -0,0 +1,389 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import numpy as np
3
+ from sklearn.cluster import KMeans
4
+ import warnings
5
+ import time
6
+
7
+ import torch
8
+ from torchvision.ops import batched_nms
9
+ from numpy import ndarray
10
+ # Suppress ALL runtime and sklearn warnings
11
+ warnings.filterwarnings('ignore', category=RuntimeWarning)
12
+ warnings.filterwarnings('ignore', category=FutureWarning)
13
+ warnings.filterwarnings('ignore', category=UserWarning)
14
+
15
+ # Suppress sklearn warnings specifically
16
+ import logging
17
+ logging.getLogger('sklearn').setLevel(logging.ERROR)
18
+
19
+ def get_grass_color(img):
20
+ # Convert image to HSV color space
21
+ hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
22
+
23
+ # Define range of green color in HSV
24
+ lower_green = np.array([30, 40, 40])
25
+ upper_green = np.array([80, 255, 255])
26
+
27
+ # Threshold the HSV image to get only green colors
28
+ mask = cv2.inRange(hsv, lower_green, upper_green)
29
+
30
+ # Calculate the mean value of the pixels that are not masked
31
+ masked_img = cv2.bitwise_and(img, img, mask=mask)
32
+ grass_color = cv2.mean(img, mask=mask)
33
+ return grass_color[:3]
34
+
35
+ def get_players_boxes(frame, result):
36
+ players_imgs = []
37
+ players_boxes = []
38
+ for (box, score, cls) in result:
39
+ label = int(cls)
40
+ if label == 0:
41
+ x1, y1, x2, y2 = box.astype(int)
42
+ player_img = frame[y1: y2, x1: x2]
43
+ players_imgs.append(player_img)
44
+ players_boxes.append([box, score, cls])
45
+ return players_imgs, players_boxes
46
+
47
+ def get_kits_colors(players, grass_hsv=None, frame=None):
48
+ kits_colors = []
49
+ if grass_hsv is None:
50
+ grass_color = get_grass_color(frame)
51
+ grass_hsv = cv2.cvtColor(np.uint8([[list(grass_color)]]), cv2.COLOR_BGR2HSV)
52
+
53
+ for player_img in players:
54
+ # Skip empty or invalid images
55
+ if player_img is None or player_img.size == 0 or len(player_img.shape) != 3:
56
+ continue
57
+
58
+ # Convert image to HSV color space
59
+ hsv = cv2.cvtColor(player_img, cv2.COLOR_BGR2HSV)
60
+
61
+ # Define range of green color in HSV
62
+ lower_green = np.array([grass_hsv[0, 0, 0] - 10, 40, 40])
63
+ upper_green = np.array([grass_hsv[0, 0, 0] + 10, 255, 255])
64
+
65
+ # Threshold the HSV image to get only green colors
66
+ mask = cv2.inRange(hsv, lower_green, upper_green)
67
+
68
+ # Bitwise-AND mask and original image
69
+ mask = cv2.bitwise_not(mask)
70
+ upper_mask = np.zeros(player_img.shape[:2], np.uint8)
71
+ upper_mask[0:player_img.shape[0]//2, 0:player_img.shape[1]] = 255
72
+ mask = cv2.bitwise_and(mask, upper_mask)
73
+
74
+ kit_color = np.array(cv2.mean(player_img, mask=mask)[:3])
75
+
76
+ kits_colors.append(kit_color)
77
+ return kits_colors
78
+
79
+ def get_kits_classifier(kits_colors):
80
+ if len(kits_colors) == 0:
81
+ return None
82
+ if len(kits_colors) == 1:
83
+ # Only one kit color, create a dummy classifier
84
+ return None
85
+ kits_kmeans = KMeans(n_clusters=2)
86
+ kits_kmeans.fit(kits_colors)
87
+ return kits_kmeans
88
+
89
+ def classify_kits(kits_classifer, kits_colors):
90
+ if kits_classifer is None or len(kits_colors) == 0:
91
+ return np.array([0]) # Default to team 0
92
+ team = kits_classifer.predict(kits_colors)
93
+ return team
94
+
95
+ def get_left_team_label(players_boxes, kits_colors, kits_clf):
96
+ left_team_label = 0
97
+ team_0 = []
98
+ team_1 = []
99
+
100
+ for i in range(len(players_boxes)):
101
+ x1, y1, x2, y2 = players_boxes[i][0].astype(int)
102
+ team = classify_kits(kits_clf, [kits_colors[i]]).item()
103
+ if team == 0:
104
+ team_0.append(np.array([x1]))
105
+ else:
106
+ team_1.append(np.array([x1]))
107
+
108
+ team_0 = np.array(team_0)
109
+ team_1 = np.array(team_1)
110
+
111
+ # Safely calculate averages with fallback for empty arrays
112
+ avg_team_0 = np.average(team_0) if len(team_0) > 0 else 0
113
+ avg_team_1 = np.average(team_1) if len(team_1) > 0 else 0
114
+
115
+ if avg_team_0 - avg_team_1 > 0:
116
+ left_team_label = 1
117
+
118
+ return left_team_label
119
+
120
+ def check_box_boundaries(boxes, img_height, img_width):
121
+ """
122
+ Check if bounding boxes are within image boundaries and clip them if necessary.
123
+
124
+ Args:
125
+ boxes: numpy array of shape (N, 4) with [x1, y1, x2, y2] format
126
+ img_height: height of the image
127
+ img_width: width of the image
128
+
129
+ Returns:
130
+ valid_boxes: numpy array of valid boxes within boundaries
131
+ valid_indices: indices of valid boxes
132
+ """
133
+ x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
134
+
135
+ # Check if boxes are within boundaries
136
+ valid_mask = (x1 >= 0) & (y1 >= 0) & (x2 < img_width) & (y2 < img_height) & (x1 < x2) & (y1 < y2)
137
+
138
+ if not np.any(valid_mask):
139
+ return np.array([]), np.array([])
140
+
141
+ valid_boxes = boxes[valid_mask]
142
+ valid_indices = np.where(valid_mask)[0]
143
+
144
+ # Clip boxes to image boundaries
145
+ valid_boxes[:, 0] = np.clip(valid_boxes[:, 0], 0, img_width - 1) # x1
146
+ valid_boxes[:, 1] = np.clip(valid_boxes[:, 1], 0, img_height - 1) # y1
147
+ valid_boxes[:, 2] = np.clip(valid_boxes[:, 2], 0, img_width - 1) # x2
148
+ valid_boxes[:, 3] = np.clip(valid_boxes[:, 3], 0, img_height - 1) # y2
149
+
150
+ return valid_boxes, valid_indices
151
+
152
+ def process_team_identification_batch(frames, results, kits_clf, left_team_label, grass_hsv):
153
+ """
154
+ Process team identification and label formatting for batch results.
155
+
156
+ Args:
157
+ frames: list of frames
158
+ results: list of detection results for each frame
159
+ kits_clf: trained kit classifier
160
+ left_team_label: label for left team
161
+ grass_hsv: grass color in HSV format
162
+
163
+ Returns:
164
+ processed_results: list of processed results with team identification
165
+ """
166
+ processed_results = []
167
+
168
+ for frame_idx, frame in enumerate(frames):
169
+ frame_results = []
170
+ frame_detections = results[frame_idx]
171
+
172
+ if not frame_detections:
173
+ processed_results.append([])
174
+ continue
175
+
176
+ # Extract player boxes and images
177
+ players_imgs = []
178
+ players_boxes = []
179
+ player_indices = []
180
+
181
+ for idx, (box, score, cls) in enumerate(frame_detections):
182
+ label = int(cls)
183
+ if label == 0: # Player detection
184
+ x1, y1, x2, y2 = box.astype(int)
185
+
186
+ # Check boundaries
187
+ if (x1 >= 0 and y1 >= 0 and x2 < frame.shape[1] and y2 < frame.shape[0] and x1 < x2 and y1 < y2):
188
+ player_img = frame[y1:y2, x1:x2]
189
+ if player_img.size > 0: # Ensure valid image
190
+ players_imgs.append(player_img)
191
+ players_boxes.append([box, score, cls])
192
+ player_indices.append(idx)
193
+
194
+ # Initialize player team mapping
195
+ player_team_map = {}
196
+
197
+ # Process team identification if we have players
198
+ if players_imgs and kits_clf is not None:
199
+ kits_colors = get_kits_colors(players_imgs, grass_hsv)
200
+ teams = classify_kits(kits_clf, kits_colors)
201
+
202
+ # Create mapping from player index to team
203
+ for i, team in enumerate(teams):
204
+ player_team_map[player_indices[i]] = team.item()
205
+
206
+ id = 0
207
+ # Process all detections with team identification
208
+ for idx, (box, score, cls) in enumerate(frame_detections):
209
+ label = int(cls)
210
+ x1, y1, x2, y2 = box.astype(int)
211
+
212
+ # Check boundaries
213
+ valid_boxes, valid_indices = check_box_boundaries(
214
+ np.array([[x1, y1, x2, y2]]), frame.shape[0], frame.shape[1]
215
+ )
216
+
217
+ if len(valid_boxes) == 0:
218
+ continue
219
+
220
+ x1, y1, x2, y2 = valid_boxes[0].astype(int)
221
+
222
+ # Apply team identification logic
223
+ if label == 0: # Player
224
+ if players_imgs and kits_clf is not None and idx in player_team_map:
225
+ team = player_team_map[idx]
226
+ if team == left_team_label:
227
+ final_label = 6 # Player-L (Left team)
228
+ else:
229
+ final_label = 7 # Player-R (Right team)
230
+ else:
231
+ final_label = 6 # Default player label
232
+
233
+ elif label == 1: # Goalkeeper
234
+ final_label = 1 # GK
235
+
236
+ elif label == 2: # Ball
237
+ final_label = 0 # Ball
238
+
239
+ elif label == 3 or label == 4: # Referee or other
240
+ final_label = 3 # Referee
241
+
242
+ else:
243
+ continue
244
+ # final_label = int(label) # Keep original label, ensure it's int
245
+
246
+ frame_results.append({
247
+ "id": int(id),
248
+ "bbox": [int(x1), int(y1), int(x2), int(y2)],
249
+ "class_id": int(final_label),
250
+ "conf": float(score)
251
+ })
252
+ id = id + 1
253
+
254
+ processed_results.append(frame_results)
255
+
256
+ return processed_results
257
+
258
+ def convert_numpy_types(obj):
259
+ """Convert numpy types to native Python types for JSON serialization."""
260
+ if isinstance(obj, np.integer):
261
+ return int(obj)
262
+ elif isinstance(obj, np.floating):
263
+ return float(obj)
264
+ elif isinstance(obj, np.ndarray):
265
+ return obj.tolist()
266
+ elif isinstance(obj, dict):
267
+ return {key: convert_numpy_types(value) for key, value in obj.items()}
268
+ elif isinstance(obj, list):
269
+ return [convert_numpy_types(item) for item in obj]
270
+ else:
271
+ return obj
272
+
273
+ def pre_process_img(frames, scale):
274
+ imgs = np.stack([cv2.resize(frame, (int(scale), int(scale))) for frame in frames])
275
+ imgs = imgs.transpose(0, 3, 1, 2)
276
+ imgs = imgs.astype(np.float32) / 255.0 # Normalize
277
+ return imgs
278
+
279
+ def post_process_output(outputs, x_scale, y_scale, conf_thresh=0.6, nms_thresh=0.75):
280
+ B, C, N = outputs.shape
281
+ outputs = torch.from_numpy(outputs)
282
+ outputs = outputs.permute(0, 2, 1)
283
+ boxes = outputs[..., :4]
284
+ class_scores = 1 / (1 + torch.exp(-outputs[..., 4:]))
285
+ conf, class_id = class_scores.max(dim=2)
286
+
287
+ mask = conf > conf_thresh
288
+
289
+ for i in range(class_id.shape[0]): # loop over batch
290
+ # Find detections that are balls
291
+ ball_idx = np.where(class_id[i] == 2)[0]
292
+ if ball_idx.size > 0:
293
+ # Pick the one with the highest confidence
294
+ top = ball_idx[np.argmax(conf[i, ball_idx])]
295
+ if conf[i, top] > 0.55: # apply confidence threshold
296
+ mask[i, top] = True
297
+
298
+ # ball_mask = (class_id == 2) & (conf > 0.51)
299
+ # mask = mask | ball_mask
300
+
301
+ batch_idx, pred_idx = mask.nonzero(as_tuple=True)
302
+
303
+ if len(batch_idx) == 0:
304
+ return [[] for _ in range(B)]
305
+
306
+ boxes = boxes[batch_idx, pred_idx]
307
+ conf = conf[batch_idx, pred_idx]
308
+ class_id = class_id[batch_idx, pred_idx]
309
+
310
+ x, y, w, h = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
311
+ x1 = (x - w / 2) * x_scale
312
+ y1 = (y - h / 2) * y_scale
313
+ x2 = (x + w / 2) * x_scale
314
+ y2 = (y + h / 2) * y_scale
315
+ boxes_xyxy = torch.stack([x1, y1, x2, y2], dim=1)
316
+
317
+ max_coord = 1e4
318
+ offset = batch_idx.to(boxes_xyxy) * max_coord
319
+ boxes_for_nms = boxes_xyxy + offset[:, None]
320
+
321
+ keep = batched_nms(boxes_for_nms, conf, batch_idx, nms_thresh)
322
+
323
+ boxes_final = boxes_xyxy[keep]
324
+ conf_final = conf[keep]
325
+ class_final = class_id[keep]
326
+ batch_final = batch_idx[keep]
327
+
328
+ results = [[] for _ in range(B)]
329
+ for b in range(B):
330
+ mask_b = batch_final == b
331
+ if mask_b.sum() == 0:
332
+ continue
333
+ results[b] = list(zip(boxes_final[mask_b].numpy(),
334
+ conf_final[mask_b].numpy(),
335
+ class_final[mask_b].numpy()))
336
+ return results
337
+
338
+ def player_detection_result(frames: list[ndarray], batch_size, model, kits_clf=None, left_team_label=None, grass_hsv=None):
339
+ start_time = time.time()
340
+ # input_layer = model.input(0)
341
+ # output_layer = model.output(0)
342
+ height, width = frames[0].shape[:2]
343
+ scale = 640.0
344
+ x_scale = width / scale
345
+ y_scale = height / scale
346
+
347
+ # infer_queue = AsyncInferQueue(model, len(frames))
348
+
349
+ infer_time = time.time()
350
+ kits_clf = kits_clf
351
+ left_team_label = left_team_label
352
+ grass_hsv = grass_hsv
353
+ results = []
354
+ for i in range(0, len(frames), batch_size):
355
+ if i + batch_size > len(frames):
356
+ batch_size = len(frames) - i
357
+ batch_frames = frames[i:i + batch_size]
358
+ imgs = pre_process_img(batch_frames, scale)
359
+
360
+ input_name = model.get_inputs()[0].name
361
+ outputs = model.run(None, {input_name: imgs})[0]
362
+ raw_results = post_process_output(np.array(outputs), x_scale, y_scale)
363
+
364
+ if kits_clf is None or left_team_label is None or grass_hsv is None:
365
+ # Use first frame to initialize team classification
366
+ first_frame = batch_frames[0]
367
+ first_frame_results = raw_results[0] if raw_results else []
368
+
369
+ if first_frame_results:
370
+ players_imgs, players_boxes = get_players_boxes(first_frame, first_frame_results)
371
+ if players_imgs:
372
+ grass_color = get_grass_color(first_frame)
373
+ grass_hsv = cv2.cvtColor(np.uint8([[list(grass_color)]]), cv2.COLOR_BGR2HSV)
374
+ kits_colors = get_kits_colors(players_imgs, grass_hsv)
375
+ if kits_colors: # Only proceed if we have valid kit colors
376
+ kits_clf = get_kits_classifier(kits_colors)
377
+ if kits_clf is not None:
378
+ left_team_label = int(get_left_team_label(players_boxes, kits_colors, kits_clf))
379
+
380
+ # Process team identification and boundary checking
381
+ processed_results = process_team_identification_batch(
382
+ batch_frames, raw_results, kits_clf, left_team_label, grass_hsv
383
+ )
384
+
385
+ processed_results = convert_numpy_types(processed_results)
386
+ results.extend(processed_results)
387
+
388
+ # Return the same format as before for compatibility
389
+ return results, kits_clf, left_team_label, grass_hsv
team_cluster.pyc ADDED
Binary file (7.62 kB). View file
 
utils.pyc ADDED
Binary file (20.6 kB). View file