Update cursor_tracker.py

cursor_tracker.py

@@ -1,90 +1,21 @@
-"""
-Cursor detection and action logging utilities.
-"""
-import cv2
-import numpy as np
-import json
-from pathlib import Path
-import os
-import smtplib
-from email.message import EmailMessage
+from typing import List, Tuple
+from PIL import Image
+from io import BytesIO
+import requests
 
-def to_rgb(img):
-    if img is None:
-        return None
-    if len(img.shape) == 2:
-        return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
-    if img.shape[2] == 4:
-        return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
-    return img
+def track_cursor_from_images(
+    images: List[Tuple[str, np.ndarray]],
+    cursor_templates_dir: str,
+    output_json_path: str,
+    threshold=0.8,
+    email_results=False
+):
+    """Detect cursor in a list of in-memory images (filename, image)."""
 
-def get_mask_from_alpha(template_img):
-    if template_img is not None and len(template_img.shape) == 3 and template_img.shape[2] == 4:
-        # Use alpha channel as mask (nonzero alpha = 255)
-        return (template_img[:, :, 3] > 0).astype(np.uint8) * 255
-    return None
-
-def detect_cursor_in_frame_multi(frame, cursor_templates, threshold=0.8):
-    """Detect cursor position in a frame using multiple templates. Returns best match above threshold."""
-    best_pos = None
-    best_conf = -1
-    best_template_name = None
-    frame_rgb = to_rgb(frame)
-    for template_name, cursor_template in cursor_templates.items():
-        template_rgb = to_rgb(cursor_template)
-        mask = get_mask_from_alpha(cursor_template)
-        if template_rgb is None or frame_rgb is None or template_rgb.shape[2] != frame_rgb.shape[2]:
-            print(f"[WARN] Skipping template {template_name} due to channel mismatch or load error.")
-            continue
-        try:
-            result = cv2.matchTemplate(frame_rgb, template_rgb, cv2.TM_CCOEFF_NORMED, mask=mask)
-        except Exception as e:
-            print(f"[WARN] matchTemplate failed for {template_name}: {e}")
-            continue
-        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
-        if max_val > best_conf:
-            best_conf = max_val
-            if max_val >= threshold:
-                cursor_w, cursor_h = template_rgb.shape[1], template_rgb.shape[0]
-                cursor_x = max_loc[0] + cursor_w // 2
-                cursor_y = max_loc[1] + cursor_h // 2
-                best_pos = (cursor_x, cursor_y)
-                best_template_name = template_name
-    if best_conf >= threshold:
-        return best_pos, best_conf, best_template_name
-    return None, best_conf, None
-
-def ensure_dir(path):
-    os.makedirs(path, exist_ok=True)
-
-def send_email_with_attachment(subject, body, to_email, from_email, app_password, attachment_path):
-    msg = EmailMessage()
-    msg['Subject'] = subject
-    msg['From'] = from_email
-    msg['To'] = to_email
-    msg.set_content(body)
-    with open(attachment_path, 'rb') as f:
-        file_data = f.read()
-        file_name = Path(attachment_path).name
-    msg.add_attachment(file_data, maintype='application', subtype='octet-stream', filename=file_name)
-    try:
-        with smtplib.SMTP_SSL('smtp.gmail.com', 465) as smtp:
-            smtp.login(from_email, app_password)
-            smtp.send_message(msg)
-        print(f"[SUCCESS] Email sent to {to_email}")
-    except Exception as e:
-        print(f"[ERROR] Failed to send email: {e}")
-
-def track_cursor(frames_dir, cursor_templates_dir, output_json_path, threshold=0.8, start_frame=1, email_results=False):
-    """Detect cursor in each frame using multiple templates, print status, and write positions to a JSON file."""
-    frames_dir = Path(frames_dir).resolve()
-    output_json_path = Path(output_json_path).resolve()
     cursor_templates_dir = Path(cursor_templates_dir).resolve()
-    print(f"[DEBUG] frames_dir: {frames_dir}")
-    print(f"[DEBUG] cursor_templates_dir: {cursor_templates_dir}")
-    print(f"[DEBUG] output_json_path: {output_json_path}")
-    ensure_dir(frames_dir)
+    output_json_path = Path(output_json_path).resolve()
     ensure_dir(output_json_path.parent)
+
     # Load all PNG templates from the cursor_templates_dir
     cursor_templates = {}
     for template_file in cursor_templates_dir.glob('*.png'):
@@ -95,20 +26,19 @@ def track_cursor(frames_dir, cursor_templates_dir, output_json_path, threshold=0
             print(f"[WARN] Could not load template: {template_file}")
     if not cursor_templates:
         raise FileNotFoundError(f"No cursor templates found in: {cursor_templates_dir}")
+
     results = []
-    for frame_file in sorted(frames_dir.glob('*.png')):
-        frame_num = int(frame_file.stem)
-        if frame_num < start_frame:
-            continue
-        frame = cv2.imread(str(frame_file), cv2.IMREAD_UNCHANGED)
+
+    for frame_filename, frame in images:
         if frame is None:
-            print(f"[WARN] Could not load frame: {frame_file}")
+            print(f"[WARN] Frame {frame_filename} is empty")
             continue
+
         pos, conf, template_name = detect_cursor_in_frame_multi(frame, cursor_templates, threshold)
         if pos is not None:
-            print(f"{frame_file.name}: Cursor at {pos} (template: {template_name})")
+            print(f"{frame_filename}: Cursor at {pos} (template: {template_name})")
             results.append({
-                'frame': frame_file.name,
+                'frame': frame_filename,
                 'cursor_active': True,
                 'x': pos[0],
                 'y': pos[1],
@@ -116,15 +46,16 @@ def track_cursor(frames_dir, cursor_templates_dir, output_json_path, threshold=0
                 'template': template_name
             })
         else:
-            print(f"{frame_file.name}: Cursor disabled")
+            print(f"{frame_filename}: Cursor disabled")
             results.append({
-                'frame': frame_file.name,
+                'frame': frame_filename,
                 'cursor_active': False,
                 'x': None,
                 'y': None,
                 'confidence': conf,
                 'template': None
             })
+
     try:
         with open(output_json_path, 'w') as f:
             json.dump(results, f, indent=2)
@@ -148,12 +79,3 @@ def track_cursor(frames_dir, cursor_templates_dir, output_json_path, threshold=0
     except Exception as e:
         print(f"[ERROR] Failed to write output JSON: {e}")
         raise
-
-if __name__ == "__main__":
-    # Pre-set all paths and parameters for direct execution
-    FRAMES_DIR = 'frames/unzipped/blender_day1'
-    CURSOR_TEMPLATES_DIR = 'cursors'  # Directory containing all cursor PNG templates
-    OUTPUT_JSON = 'annotations/blender1.json'
-    THRESHOLD = 0.8
-    print(f"[INFO] Running with preset values:\n  frames_dir={FRAMES_DIR}\n  cursor_templates_dir={CURSOR_TEMPLATES_DIR}\n  output_json={OUTPUT_JSON}\n  threshold={THRESHOLD}\n  start_frame=0\n")
-    track_cursor(FRAMES_DIR, CURSOR_TEMPLATES_DIR, OUTPUT_JSON, THRESHOLD, start_frame=0, email_results=True)