first commit

miner.py

@@ -0,0 +1,520 @@
+from pathlib import Path
+from ultralytics import YOLO
+from numpy import ndarray
+from pydantic import BaseModel
+from typing import List, Tuple, Optional
+import numpy as np
+import cv2
+from sklearn.cluster import KMeans
+import base64
+import boto3
+import json
+import uuid
+import torch
+from torchvision.models import resnet50, ResNet50_Weights
+import torchvision.transforms as transforms
+
+
+########################################
+# Helper utilities for R2 storage
+########################################
+
+def init_r2_client():
+    """
+    Khởi tạo S3 client cho Cloudflare R2.
+    Returns:
+        tuple: (s3_client, bucket_name, can_upload)
+    """
+    try:
+        r2_account_id = "f5ac691bc782b80f90edb38eba5534ad"
+        r2_access_key_id = "54f3343f68621c563d7ca29d3b356122"
+        r2_secret_access_key = "41484baa8a10838e197f528b7eefbb824e1f38ffe13abc4e6b5fa7b68ad6d82d"
+        bucket_name = "my-miner-sn44"
+
+        can_upload = all([r2_account_id, r2_access_key_id, r2_secret_access_key, bucket_name])
+
+        if can_upload:
+            s3_client = boto3.client(
+                's3',
+                endpoint_url=f"https://{r2_account_id}.r2.cloudflarestorage.com",
+                aws_access_key_id=r2_access_key_id,
+                aws_secret_access_key=r2_secret_access_key,
+                region_name='auto'
+            )
+            print(f"✅ R2 client initialized for bucket: {bucket_name}")
+            return s3_client, bucket_name, True
+        else:
+            print("⚠️  Thiếu một hoặc nhiều secret của R2, sẽ không lưu frames.")
+            return None, None, False
+
+    except Exception as e:
+        print(f"⚠️  Không thể khởi tạo S3 client: {e}")
+        return None, None, False
+
+
+def image_to_base64(image: np.ndarray, quality: int = 85) -> str:
+    """
+    Convert numpy image array to base64 string.
+    Args:
+        image: numpy array (BGR format from OpenCV)
+        quality: JPEG quality (1-100, default 85)
+    Returns:
+        str: base64 encoded string
+    """
+    # Encode image as JPEG
+    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
+    _, buffer = cv2.imencode('.jpg', image, encode_param)
+    # Convert to base64
+    base64_str = base64.b64encode(buffer).decode('utf-8')
+    return base64_str
+
+
+def upload_frames_to_r2(
+    s3_client, 
+    bucket_name: str, 
+    frames_base64: List[dict], 
+    challenge_id: str
+) -> bool:
+    """
+    Upload danh sách frames (base64) lên Cloudflare R2 dưới dạng JSON.
+    Args:
+        s3_client: boto3 S3 client
+        bucket_name: Tên bucket R2
+        frames_base64: List of dicts with frame_id and base64 data
+        challenge_id: ID của challenge (dùng làm tên file)
+    Returns:
+        bool: True nếu upload thành công
+    """
+    try:
+        json_filename = f"{challenge_id}_frames.json"
+        json_data = json.dumps(frames_base64)
+        
+        s3_client.put_object(
+            Bucket=bucket_name,
+            Key=json_filename,
+            Body=json_data.encode('utf-8'),
+            ContentType='application/json'
+        )
+        print(f"✅ {len(frames_base64)} frames đã được lưu vào R2: {json_filename}")
+        return True
+    except Exception as e:
+        print(f"⚠️  Lỗi khi tải frames lên R2: {e}")
+        return False
+
+
+########################################
+# Helper utilities for grass & color clustering
+########################################
+
+def get_grass_color(img: np.ndarray) -> Tuple[int, int, int]:
+    """Estimate dominant green (grass) color from the image in BGR."""
+    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
+    lower_green = np.array([30, 40, 40])
+    upper_green = np.array([80, 255, 255])
+    mask = cv2.inRange(hsv, lower_green, upper_green)
+    grass_color = cv2.mean(img, mask=mask)
+    return grass_color[:3]
+
+
+def get_players_boxes(result):
+    """Extract player crops and boxes from YOLO result.
+    
+    Model class mapping:
+    0: 'Player', 1: 'GoalKeeper', 2: 'Ball', 3: 'Main Referee', 
+    4: 'Side Referee', 5: 'Staff Member', 6: 'left team', 7: 'right team'
+    """
+    players_imgs, players_boxes = [], []
+    for box in result.boxes:
+        label = int(box.cls.cpu().numpy()[0])
+        if label == 0:  # Player class (cls_id=0 is Player)
+            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
+            crop = result.orig_img[y1:y2, x1:x2]
+            if crop.size > 0:
+                players_imgs.append(crop)
+                players_boxes.append((x1, y1, x2, y2))
+    return players_imgs, players_boxes
+
+
+def get_kits_colors(players, grass_hsv=None, frame=None):
+    """Extract average kit colors from player crops."""
+    kits_colors = []
+    if grass_hsv is None:
+        grass_color = get_grass_color(frame)
+        grass_hsv = cv2.cvtColor(np.uint8([[list(grass_color)]]), cv2.COLOR_BGR2HSV)
+    for player_img in players:
+        hsv = cv2.cvtColor(player_img, cv2.COLOR_BGR2HSV)
+        lower_green = np.array([grass_hsv[0, 0, 0] - 10, 40, 40])
+        upper_green = np.array([grass_hsv[0, 0, 0] + 10, 255, 255])
+        mask = cv2.inRange(hsv, lower_green, upper_green)
+        mask = cv2.bitwise_not(mask)
+        upper_mask = np.zeros(player_img.shape[:2], np.uint8)
+        upper_mask[0:player_img.shape[0] // 2, :] = 255
+        mask = cv2.bitwise_and(mask, upper_mask)
+        kit_color = np.array(cv2.mean(player_img, mask=mask)[:3])
+        kits_colors.append(kit_color)
+    return kits_colors
+
+
+# ============================================================================
+# Team Classification using ResNet50 Features
+# ============================================================================
+class TeamClassifierResNet:
+    def __init__(self, device="cuda"):
+        self.device = device
+        self.model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1).to(device).eval()
+        self.preprocess = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225],
+            ),
+        ])
+        self.kmeans = None
+        self.left_team = None
+
+    def get_feature(self, img):
+        t = self.preprocess(img).unsqueeze(0).to(self.device)
+        with torch.no_grad():
+            f = self.model(t)
+        # ✅ convert to float64 here to be safe for sklearn
+        return f.squeeze(0).cpu().numpy().astype(np.float64)
+
+    def fit(self, player_crops, player_centers):
+        feats = []
+        for crop in player_crops:
+            feats.append(self.get_feature(crop))
+        feats = np.array(feats, dtype=np.float64)
+
+        # KMeans feature clustering
+        self.kmeans = KMeans(n_clusters=2, random_state=0)
+        labels = self.kmeans.fit_predict(feats)
+
+        # Determine which team is on the left side
+        mean_x = {0: [], 1: []}
+        for lab, (x, y) in zip(labels, player_centers):
+            mean_x[lab].append(x)
+
+        left = 0 if np.mean(mean_x[0]) < np.mean(mean_x[1]) else 1
+        self.left_team = left
+        return labels
+
+########################################
+# Data models
+########################################
+
+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]]
+
+
+########################################
+# Main Miner class
+########################################
+
+class Miner:
+    """
+    Main class for sn44-compatible inference pipeline.
+    Integrates YOLO + team color classification (HSV-based).
+    """
+    CORNER_INDICES = {0, 5, 24, 29}
+
+    def __init__(
+        self, 
+        path_hf_repo: Path,
+    ) -> None:
+        """Load models from the repository.
+        
+        Model class mapping:
+        0: 'Player', 1: 'GoalKeeper', 2: 'Ball', 3: 'Main Referee', 
+        4: 'Side Referee', 5: 'Staff Member', 6: 'left team', 7: 'right team'
+        
+        Args:
+            path_hf_repo: Path to HuggingFace repo with models
+            enable_frame_storage: If True, collect frames as base64 for R2 upload
+            storage_quality: JPEG quality for stored frames (1-100)
+            challenge_id: Challenge ID for R2 upload (required if enable_frame_storage=True)
+        """
+        enable_frame_storage = True
+        storage_quality = 85
+
+        challenge_id = f"challenge_{uuid.uuid4().hex[:12]}"
+        
+        # Option 2: Timestamp-based (unique theo thời gian)
+        # challenge_id = f"challenge_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}"
+        
+        print(f"✅ Auto-generated challenge_id: {challenge_id}")
+
+
+        self.bbox_model = YOLO(path_hf_repo / "251110-football-detection.pt")
+        print("✅ BBox Model Loaded")
+        self.keypoints_model = YOLO(path_hf_repo / "17112025_keypoint.pt")
+        print("✅ Keypoints Model (Pose) Loaded")
+
+        self.team_kmeans = None
+        self.left_team_label = 0
+        self.grass_hsv = None
+        self.team_classifier_fitted = False
+        
+        # Frame storage setup
+        self.enable_frame_storage = enable_frame_storage
+        self.storage_quality = storage_quality
+        self.stored_frames: List[dict] = []  # Store frames as base64
+        self.challenge_id = challenge_id
+        
+        # R2 client setup
+        if enable_frame_storage:
+            self.s3_client, self.r2_bucket, self.can_upload = init_r2_client()
+            if not challenge_id:
+                print("⚠️  WARNING: enable_frame_storage=True nhưng chưa set challenge_id")
+        else:
+            self.s3_client = None
+            self.r2_bucket = None
+            self.can_upload = False
+
+    def __repr__(self) -> str:
+        return (
+            f"BBox Model: {type(self.bbox_model).__name__}\n"
+            f"Keypoints Model: {type(self.keypoints_model).__name__}\n"
+            f"Team Clustering: HSV + KMeans"
+        )
+
+    def fit_team_classifier(self, frame):
+        print("[INFO] Extracting players from first frame for team classifier...")
+
+        result = self.bbox_model(frame, conf=0.2, verbose=False)[0]
+
+        players_imgs = []
+        player_centers = []
+
+        if result and result.boxes is not None:
+            for box in result.boxes:
+                cls_id = int(box.cls.cpu().numpy()[0])
+                if cls_id == 0:  # player
+                    x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
+                    crop = frame[y1:y2, x1:x2]
+                    players_imgs.append(crop)
+                    player_centers.append(((x1 + x2) / 2, (y1 + y2) / 2))
+
+        if len(players_imgs) < 2:
+            print("[WARN] Not enough players to fit KMeans. Skip.")
+            self.team_classifier_fitted = True
+            return None
+
+        # Init classifier
+        self.team_classifier = TeamClassifierResNet()
+
+        # Extract features
+        feats = []
+        for crop in players_imgs:
+            try:
+                f = self.team_classifier.get_feature(crop)
+                feats.append(f)
+            except:
+                feats.append(np.zeros(512, dtype=np.float64))
+
+        feats = np.array(feats, dtype=np.float64)  # ✅ convert to float64
+
+        # Fit KMeans
+        print("[INFO] Fitting KMeans on ResNet player features...")
+        self.team_kmeans = KMeans(n_clusters=2, random_state=0)
+        teams = self.team_kmeans.fit_predict(feats)
+
+        # Determine left team
+        left_cluster = np.argmin([
+            np.mean([c for c, t in zip([x for x, y in player_centers], teams) if t == cluster])
+            for cluster in [0, 1]
+        ])
+        self.left_team_label = left_cluster
+        self.team_classifier_fitted = True
+        print("[INFO] Team classifier fitted using ResNet50.")
+
+    
+    def _auto_upload_frames(self) -> None:
+        """Internal method to auto-upload frames after last batch."""
+        if not self.challenge_id:
+            print("❌ Không thể upload: challenge_id chưa được set!")
+            return
+        
+        total_frames = len(self.stored_frames)
+        size_mb = self.get_stored_frames_size_mb()
+        
+        print(f"📊 Tổng frames đã lưu: {total_frames}")
+        print(f"💾 Size trong memory: {size_mb:.2f} MB")
+        print(f"📤 Đang upload lên R2...")
+        
+        success = upload_frames_to_r2(
+            self.s3_client,
+            self.r2_bucket,
+            self.stored_frames,
+            self.challenge_id
+        )
+        
+        if success:
+            print(f"✅ Upload thành công {total_frames} frames!")
+            print(f"📁 File trên R2: {self.challenge_id}_frames.json")
+            # Clear frames after successful upload
+            self.clear_stored_frames()
+        else:
+            print(f"❌ Upload thất bại!")
+            print(f"💡 Frames vẫn còn trong memory. Có thể retry bằng: miner.upload_stored_frames('{self.challenge_id}')")
+    
+    def upload_stored_frames(self, challenge_id: str) -> bool:
+        """
+        Upload all stored frames to R2.
+        Args:
+            challenge_id: ID của challenge để đặt tên file
+        Returns:
+            bool: True nếu upload thành công
+        """
+        if not self.can_upload:
+            print("⚠️  R2 client chưa được khởi tạo, không thể upload frames.")
+            return False
+        
+        if len(self.stored_frames) == 0:
+            print("⚠️  Không có frames nào để upload.")
+            return False
+        
+        print(f"📤 Đang upload {len(self.stored_frames)} frames lên R2...")
+        success = upload_frames_to_r2(
+            self.s3_client,
+            self.r2_bucket,
+            self.stored_frames,
+            challenge_id
+        )
+        
+        if success:
+            print(f"✅ Đã upload thành công {len(self.stored_frames)} frames")
+            return True
+        else:
+            print("Chưa upload được.")
+            return False
+    
+    def clear_stored_frames(self) -> None:
+        """Clear all stored frames from memory."""
+        self.stored_frames = []
+        print("🗑️  Đã xóa stored frames khỏi memory")
+    
+    def get_stored_frames_count(self) -> int:
+        """Get number of stored frames."""
+        return len(self.stored_frames)
+    
+    def get_stored_frames_size_mb(self) -> float:
+        """Get approximate size of stored frames in MB."""
+        if len(self.stored_frames) == 0:
+            return 0.0
+        total_size = sum(len(frame["data"]) for frame in self.stored_frames)
+        # Base64 encoding adds ~33% overhead, but we calculate as-is
+        return total_size / (1024 * 1024)
+
+    def predict_batch(self, batch_images: list[ndarray], offset: int, n_keypoints: int) -> list[TVFrameResult]:
+        results: list[TVFrameResult] = []
+
+        for i, frame in enumerate(batch_images):
+            frame_id = offset + i
+
+            if not self.team_classifier_fitted:
+                self.fit_team_classifier(frame)
+
+            bbox_result = self.bbox_model(frame, conf=0.2, verbose=False)[0]
+            boxes = []
+
+            if bbox_result and bbox_result.boxes is not None:
+                players_imgs, players_boxes = get_players_boxes(bbox_result)
+
+                # Extract features
+                player_features = []
+                for crop in players_imgs:
+                    try:
+                        feat = self.team_classifier.get_feature(crop)
+                        player_features.append(feat)
+                    except:
+                        player_features.append(np.zeros(512, dtype=np.float64))
+
+                # Predict teams
+                teams = []
+                if len(player_features) > 0 and self.team_kmeans is not None:
+                    player_features = np.array(player_features, dtype=np.float64)  # ✅ convert to float64
+                    teams = self.team_kmeans.predict(player_features)
+
+                # Map teams to boxes
+                player_indices = [idx for idx, box in enumerate(bbox_result.boxes) if int(box.cls.cpu().numpy()[0]) == 0]
+                team_predictions = {}
+                if len(player_indices) > 0 and len(teams) > 0:
+                    for player_idx, team_id in zip(player_indices, teams):
+                        team_predictions[player_idx] = 6 if team_id == self.left_team_label else 7
+
+                # Create BoundingBox list
+                for idx, box in enumerate(bbox_result.boxes):
+                    x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
+                    conf = float(box.conf.cpu().numpy()[0])
+                    cls_id = int(box.cls.cpu().numpy()[0])
+
+                    if idx in team_predictions:
+                        cls_id = team_predictions[idx]
+                    elif cls_id == 0:
+                        cls_id = 2
+                    elif cls_id == 1:
+                        cls_id = 1
+                    elif cls_id == 2:
+                        cls_id = 0
+                    elif cls_id in [3, 4]:
+                        cls_id = 3
+                    else:
+                        continue
+
+                    boxes.append(BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2, cls_id=cls_id, conf=conf))
+
+            # -----------------------------------------
+            # Keypoint detection using YOLO pose model
+            # -----------------------------------------
+            keypoints_result = self.keypoints_model(frame, verbose=False)[0]
+            frame_keypoints: List[Tuple[int, int]] = [(0, 0)] * n_keypoints
+            if keypoints_result and hasattr(keypoints_result, "keypoints") and keypoints_result.keypoints is not None:
+                frame_keypoints_with_conf = []
+                for i, part_points in enumerate(keypoints_result.keypoints.data):
+                    for k_id, (x, y, _) in enumerate(part_points):
+                        confidence = float(keypoints_result.keypoints.conf[i][k_id])
+                        frame_keypoints_with_conf.append((int(x), int(y), confidence))
+
+                if len(frame_keypoints_with_conf) < n_keypoints:
+                    frame_keypoints_with_conf.extend([(0, 0, 0.0)] * (n_keypoints - len(frame_keypoints_with_conf)))
+                else:
+                    frame_keypoints_with_conf = frame_keypoints_with_conf[:n_keypoints]
+
+                filtered_keypoints = []
+                for idx, (x, y, confidence) in enumerate(frame_keypoints_with_conf):
+                    if idx in self.CORNER_INDICES:
+                        filtered_keypoints.append((int(x), int(y)) if confidence >= 0.3 else (0, 0))
+                    else:
+                        filtered_keypoints.append((int(x), int(y)) if confidence >= 0.5 else (0, 0))
+                frame_keypoints = filtered_keypoints
+
+            results.append(TVFrameResult(frame_id=frame_id, boxes=boxes, keypoints=frame_keypoints))
+
+            # Auto-upload when reaching frame 750
+            if frame_id == 749 and self.enable_frame_storage and self.can_upload:
+                try:
+                    if len(self.stored_frames) > 0:
+                        print(f"\n{'='*60}")
+                        print(f"🏁 FRAME 750 REACHED - Tự động upload {len(self.stored_frames)} frames lên R2")
+                        print(f"{'='*60}")
+                        self._auto_upload_frames()
+                    else:
+                        print("⚠️  Frame 750 reached nhưng không có frames nào để upload.")
+                except Exception as e:
+                    print(f"⚠️  Lỗi khi upload R2: {e}")
+                    print(f"💡 Tiếp tục trả về results. Frames vẫn còn trong memory.")
+
+        return results