| from pathlib import Path
|
| from typing import List, Tuple, Dict
|
| import sys
|
| import os
|
|
|
| from numpy import ndarray
|
| from pydantic import BaseModel
|
| sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
|
|
| from ultralytics import YOLO
|
| from team_cluster import TeamClassifier
|
| from utils import (
|
| BoundingBox,
|
| Constants,
|
| )
|
| from inference import predict_batch
|
| import torch
|
| from pitch import get_cls_net
|
| import yaml
|
| os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
|
|
|
|
| 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:
|
| SMALL_CONTAINED_IOA = Constants.SMALL_CONTAINED_IOA
|
| SMALL_RATIO_MAX = Constants.SMALL_RATIO_MAX
|
| SINGLE_PLAYER_HUE_PIVOT = Constants.SINGLE_PLAYER_HUE_PIVOT
|
| CORNER_INDICES = Constants.CORNER_INDICES
|
| KEYPOINTS_CONFIDENCE = Constants.KEYPOINTS_CONFIDENCE
|
| CORNER_CONFIDENCE = Constants.CORNER_CONFIDENCE
|
| GOALKEEPER_POSITION_MARGIN = Constants.GOALKEEPER_POSITION_MARGIN
|
| MIN_SAMPLES_FOR_FIT = 16
|
| MAX_SAMPLES_FOR_FIT = 1000
|
|
|
| def __init__(self, path_hf_repo: Path) -> None:
|
| try:
|
| device = "cuda" if torch.cuda.is_available() else "cpu"
|
| model_path = path_hf_repo / "football_object_detection.onnx"
|
| self.bbox_model = YOLO(model_path)
|
|
|
| print("BBox Model Loaded")
|
|
|
| team_model_path = path_hf_repo / "osnet_model.pth.tar-100"
|
| self.team_classifier = TeamClassifier(
|
| device=device,
|
| batch_size=32,
|
| model_name=str(team_model_path)
|
| )
|
| print("Team Classifier Loaded")
|
|
|
|
|
| self.team_classifier_fitted = False
|
| self.player_crops_for_fit = []
|
|
|
| model_kp_path = path_hf_repo / 'keypoint'
|
| config_kp_path = path_hf_repo / 'hrnetv2_w48.yaml'
|
| cfg_kp = yaml.safe_load(open(config_kp_path, 'r'))
|
|
|
| loaded_state_kp = torch.load(model_kp_path, map_location=device)
|
| model = get_cls_net(cfg_kp)
|
| model.load_state_dict(loaded_state_kp)
|
| model.to(device)
|
| model.eval()
|
|
|
| self.keypoints_model = model
|
| self.kp_threshold = 0.1
|
| self.pitch_batch_size = 4
|
| self.health = "healthy"
|
| self.path_hf_repo = path_hf_repo
|
| print("✅ Keypoints Model Loaded")
|
| except Exception as e:
|
| self.health = "❌ Miner initialization failed: " + str(e)
|
| print(self.health)
|
|
|
| def __repr__(self) -> str:
|
| if self.health == 'healthy':
|
| return (
|
| f"health: {self.health}\n"
|
| f"BBox Model: {type(self.bbox_model).__name__}\n"
|
| f"Keypoints Model: {type(self.keypoints_model).__name__}"
|
| )
|
| else:
|
| return self.health
|
|
|
| def predict_batch(self, batch_images: List[ndarray], offset: int, n_keypoints: int) -> List[TVFrameResult]:
|
| results = predict_batch(
|
| self.bbox_model,
|
| self.team_classifier,
|
| self.keypoints_model,
|
| batch_images,
|
| offset,
|
| n_keypoints,
|
| self.pitch_batch_size,
|
| self.kp_threshold,
|
| self.path_hf_repo
|
| )
|
| return results |
