hursor_data_collector.py

"""
Hursor Data Collector
======================
Collect training data for gesture classification.

Gesture Labels:
- 0 = CLICK (fist gesture)
- 1 = SCROLL mode
- 2 = ZOOM mode
- 3 = CURSOR mode
- 4 = IDLE (gateway)

Controls:
- Press 0-4 to select gesture label
- Press ENTER to auto-collect 30 samples
- Press Q to quit

Author: Hursor - Hand Gesture Control System
"""

import cv2
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
from mediapipe.framework.formats import landmark_pb2
import csv
import numpy as np
import os
import time
import atexit
import urllib.request


class HursorDataCollector:
    """High-performance gesture data collector for Hursor."""
    
    # Gesture label mapping
    GESTURE_NAMES = {
        0: "CLICK",
        1: "SCROLL",
        2: "ZOOM",
        3: "CURSOR",
        4: "IDLE"
    }
    
    def __init__(self, csv_file='hursor_dataset.csv'):
        self.csv_file = csv_file
        self.model_file = 'hand_landmarker.task'
        
        # Buffer system for speed
        self.data_buffer = []
        self.buffer_limit = 50  # Write to disk every 50 samples
        
        # Download model if not exists
        self._ensure_model()
        
        # Initialize MediaPipe
        self._init_mediapipe()
        
        # Drawing utils (Faster than manual loops)
        self.mp_drawing = mp.solutions.drawing_utils
        self.mp_drawing_styles = mp.solutions.drawing_styles

        # CSV Header (X,Y,Z for 21 landmarks + label)
        self.header = ([f'x{i}' for i in range(21)] + 
                       [f'y{i}' for i in range(21)] + 
                       [f'z{i}' for i in range(21)] + ['label'])

        self._check_dataset()
        
        # Register cleanup to save remaining data on exit
        atexit.register(self._flush_buffer)

        # State
        self.current_label = None
        self.is_auto_collecting = False
        self.auto_target = 30
        self.auto_count = 0
        self.sample_total = 0
        
        # Per-class sample counts
        self.class_counts = {i: 0 for i in range(5)}
        self._count_existing_samples()

    def _ensure_model(self):
        """Download hand_landmarker.task if not exists."""
        if not os.path.exists(self.model_file):
            print("[INFO] Downloading hand_landmarker.task model...")
            url = "https://storage.googleapis.com/mediapipe-models/hand_landmarker/hand_landmarker/float16/1/hand_landmarker.task"
            try:
                urllib.request.urlretrieve(url, self.model_file)
                print("[OK] Model downloaded successfully!")
            except Exception as e:
                print(f"[ERROR] Failed to download model: {e}")
                print("[INFO] Please download manually from:")
                print(f"       {url}")
                raise

    def _init_mediapipe(self):
        """Initialize MediaPipe Hand Landmarker."""
        base_options = python.BaseOptions(model_asset_path=self.model_file)
        options = vision.HandLandmarkerOptions(
            base_options=base_options,
            num_hands=1,
            min_hand_detection_confidence=0.6,
            min_hand_presence_confidence=0.6,
            min_tracking_confidence=0.6
        )
        self.hand_landmarker = vision.HandLandmarker.create_from_options(options)

    def _check_dataset(self):
        """Creates dataset if not exists."""
        if not os.path.exists(self.csv_file):
            with open(self.csv_file, 'w', newline='') as f:
                csv.writer(f).writerow(self.header)
            self.sample_total = 0
            print(f"[OK] Created new dataset: {self.csv_file}")
        else:
            # Count lines quickly
            with open(self.csv_file, 'r') as f:
                self.sample_total = sum(1 for _ in f) - 1
            print(f"[OK] Loaded existing dataset with {self.sample_total} samples")

    def _count_existing_samples(self):
        """Count samples per class in existing dataset."""
        if not os.path.exists(self.csv_file):
            return
        
        try:
            with open(self.csv_file, 'r') as f:
                reader = csv.reader(f)
                next(reader)  # Skip header
                for row in reader:
                    if row:
                        label = int(row[-1])
                        if label in self.class_counts:
                            self.class_counts[label] += 1
        except Exception as e:
            print(f"[WARN] Could not count existing samples: {e}")

    def _flush_buffer(self):
        """Writes memory buffer to disk (Fast I/O)."""
        if not self.data_buffer:
            return
        
        with open(self.csv_file, 'a', newline='') as f:
            writer = csv.writer(f)
            writer.writerows(self.data_buffer)
        
        print(f"[SAVE] Flushed {len(self.data_buffer)} samples to disk.")
        self.sample_total += len(self.data_buffer)
        self.data_buffer = []

    def process_landmarks(self, hand_landmarks):
        """
        Advanced Normalization:
        1. Relative to Wrist (Position Invariant)
        2. Scaled by Hand Size (Scale Invariant)
        
        Returns:
            numpy array of 63 features (21 landmarks * 3 coords)
        """
        landmarks = np.array([[lm.x, lm.y, lm.z] for lm in hand_landmarks])
        
        # 1. Center around wrist
        wrist = landmarks[0]
        centered = landmarks - wrist
        
        # 2. Scale Normalization
        # Calculate size: Distance from Wrist(0) to Middle Finger Base(9)
        # Using index 9 is more stable than tip(12) which moves a lot
        palm_size = np.linalg.norm(centered[9]) 
        
        # Avoid division by zero
        if palm_size < 1e-6:
            palm_size = 1.0
            
        normalized = centered / palm_size

        # Flatten: [x0,x1,...,x20, y0,y1,...,y20, z0,z1,...,z20]
        return np.concatenate([normalized[:,0], normalized[:,1], normalized[:,2]])

    def draw_fast(self, frame, hand_landmarks):
        """Uses MediaPipe built-in drawing (Optimized C++)."""
        proto_list = landmark_pb2.NormalizedLandmarkList()
        proto_list.landmark.extend([
            landmark_pb2.NormalizedLandmark(x=lm.x, y=lm.y, z=lm.z) 
            for lm in hand_landmarks
        ])
        
        self.mp_drawing.draw_landmarks(
            frame,
            proto_list,
            mp.solutions.hands.HAND_CONNECTIONS,
            self.mp_drawing_styles.get_default_hand_landmarks_style(),
            self.mp_drawing_styles.get_default_hand_connections_style()
        )

    def draw_ui(self, frame, status_color, fps):
        """Draw comprehensive UI overlay."""
        h, w = frame.shape[:2]
        
        # Header background
        cv2.rectangle(frame, (0, 0), (w, 100), (20, 20, 20), -1)
        
        # Title
        cv2.putText(frame, "HURSOR Data Collector", (10, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
        
        # FPS
        cv2.putText(frame, f"FPS: {int(fps)}", (w - 100, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
        
        # Total samples
        total = self.sample_total + len(self.data_buffer)
        cv2.putText(frame, f"Total: {total}", (10, 55), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1)
        
        # Per-class counts
        y_offset = 55
        for label, name in self.GESTURE_NAMES.items():
            count = self.class_counts[label]
            color = (0, 255, 0) if count >= 100 else (0, 200, 255) if count >= 50 else (100, 100, 100)
            cv2.putText(frame, f"{label}:{name}={count}", (150 + label * 110, y_offset), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1)
        
        # Instructions
        cv2.putText(frame, "Press 0-4 to select gesture, ENTER to collect 30 samples, Q to quit", 
                    (10, 85), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (150, 150, 150), 1)
        
        # Current label indicator
        if self.current_label is not None:
            label_name = self.GESTURE_NAMES.get(self.current_label, "?")
            cv2.putText(frame, f"Selected: {self.current_label} ({label_name})", 
                        (w - 250, 85), cv2.FONT_HERSHEY_SIMPLEX, 0.5, status_color, 2)
        
        # Auto-collection progress bar
        if self.is_auto_collecting:
            progress = self.auto_count / self.auto_target
            bar_width = int(w * progress)
            cv2.rectangle(frame, (0, 100), (bar_width, 115), (0, 255, 0), -1)
            cv2.rectangle(frame, (0, 100), (w, 115), (100, 100, 100), 2)
            cv2.putText(frame, f"Collecting: {self.auto_count}/{self.auto_target}", 
                        (w//2 - 60, 112), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
        
        # Recording indicator
        if self.current_label is not None:
            label_name = self.GESTURE_NAMES.get(self.current_label, "?")
            
            # Large center text
            text = f"{self.current_label}: {label_name}"
            text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 1.5, 3)[0]
            text_x = (w - text_size[0]) // 2
            
            # Background
            cv2.rectangle(frame, (text_x - 20, h - 80), (text_x + text_size[0] + 20, h - 30), 
                          (30, 30, 30), -1)
            cv2.putText(frame, text, (text_x, h - 45), 
                        cv2.FONT_HERSHEY_SIMPLEX, 1.5, status_color, 3)
            
            if not self.is_auto_collecting:
                cv2.putText(frame, "Press ENTER to collect", (text_x, h - 15), 
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (180, 180, 180), 1)

    def run(self):
        """Main data collection loop."""
        cap = cv2.VideoCapture(0)
        
        # Set MJPG codec for faster camera IO (if supported)
        cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG'))
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        cap.set(cv2.CAP_PROP_FPS, 30)

        print("\n" + "=" * 60)
        print("    HURSOR Data Collector")
        print("=" * 60)
        print("\nGesture Labels:")
        for label, name in self.GESTURE_NAMES.items():
            print(f"  {label} = {name}")
        print("\nControls:")
        print("  0-4   = Select gesture to record")
        print("  ENTER = Auto-collect 30 samples")
        print("  Q     = Quit and save")
        print("=" * 60 + "\n")
        
        last_time = time.time()
        
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
            
            # FPS Calculation
            current_time = time.time()
            fps = 1.0 / max(current_time - last_time, 0.001)
            last_time = current_time
            
            # Flip for mirror effect
            frame = cv2.flip(frame, 1)
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            mp_img = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)
            
            result = self.hand_landmarker.detect(mp_img)
            
            status_color = (0, 0, 255)  # Red - no hand
            
            if result.hand_landmarks:
                landmarks_obj = result.hand_landmarks[0]
                
                # Draw hand
                self.draw_fast(frame, landmarks_obj)
                
                # Process landmarks
                features = self.process_landmarks(landmarks_obj)
                
                # Collection Logic
                if self.is_auto_collecting:
                    status_color = (0, 255, 255)  # Yellow - collecting
                    if self.auto_count < self.auto_target:
                        self.data_buffer.append(features.tolist() + [self.current_label])
                        self.class_counts[self.current_label] += 1
                        self.auto_count += 1
                    else:
                        self.is_auto_collecting = False
                        self._flush_buffer()  # Save immediately after batch
                        print(f"[OK] Collected {self.auto_target} samples for {self.GESTURE_NAMES[self.current_label]}")
                        self.current_label = None

                elif self.current_label is not None:
                    status_color = (0, 255, 0)  # Green - ready to collect
            else:
                if self.is_auto_collecting:
                    status_color = (0, 100, 255)  # Orange - waiting for hand
            
            # Flush buffer if full
            if len(self.data_buffer) >= self.buffer_limit:
                self._flush_buffer()

            # Draw UI
            self.draw_ui(frame, status_color, fps)

            cv2.imshow("Hursor Data Collector", frame)
            
            # Handle keyboard input
            k = cv2.waitKey(1) & 0xFF
            
            if k == ord('q') or k == 27:  # Q or ESC
                break
            elif k == 13:  # ENTER - start auto-collection
                if self.current_label is not None and not self.is_auto_collecting:
                    self.is_auto_collecting = True
                    self.auto_count = 0
                    print(f"[REC] Starting collection for: {self.current_label} ({self.GESTURE_NAMES[self.current_label]})")
            elif k in [ord('0'), ord('1'), ord('2'), ord('3'), ord('4')]:
                # Select label
                self.current_label = k - ord('0')
                self.is_auto_collecting = False
                print(f"[SELECT] Label: {self.current_label} ({self.GESTURE_NAMES[self.current_label]})")

        # Final save
        self._flush_buffer()
        
        print("\n" + "=" * 60)
        print("[DONE] Data Collection Complete!")
        print("-" * 40)
        print("Samples per class:")
        for label, name in self.GESTURE_NAMES.items():
            count = self.class_counts[label]
            status = "OK" if count >= 100 else "LOW" if count >= 50 else "NEED MORE"
            print(f"  {label} ({name}): {count} [{status}]")
        print(f"\nTotal: {self.sample_total} samples in {self.csv_file}")
        print("=" * 60)
        
        cap.release()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    HursorDataCollector().run()