Upload 17 files

cursor_tracker.py

@@ -0,0 +1,118 @@
+"""
+Cursor detection and action logging utilities for notebooks or scripts.
+"""
+import cv2
+import numpy as np
+import json
+from pathlib import Path
+import os
+
+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 get_mask_from_alpha(template_img):
+    if template_img is not None and len(template_img.shape) == 3 and template_img.shape[2] == 4:
+        return (template_img[:, :, 3] > 0).astype(np.uint8) * 255
+    return None
+
+def detect_cursor_in_frame_multi(frame, cursor_templates, threshold=0.8):
+    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
+        
+        _, max_val, _, 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):
+    if not os.path.exists(path):
+        os.makedirs(path, exist_ok=True)
+
+def track_cursor(frames_dir, cursor_templates_dir, output_json_path, threshold=0.8, start_frame=1):
+    frames_dir = Path(frames_dir)
+    ensure_dir(Path(output_json_path).parent)
+
+    # Load cursor templates
+    cursor_templates = {}
+    for template_file in Path(cursor_templates_dir).glob('*.png'):
+        template_img = cv2.imread(str(template_file), cv2.IMREAD_UNCHANGED)
+        if template_img is not None:
+            cursor_templates[template_file.name] = template_img
+        else:
+            print(f"[WARN] Could not load template: {template_file}")
+    if not cursor_templates:
+        raise FileNotFoundError(f"No cursor templates found in: {cursor_templates_dir}")
+
+    # Process frames
+    results = []
+    for frame_file in sorted(frames_dir.glob('*.png')):
+        try:
+            frame_num = int(frame_file.stem)
+        except:
+            continue
+
+        if frame_num < start_frame:
+            continue
+
+        frame = cv2.imread(str(frame_file), cv2.IMREAD_UNCHANGED)
+        if frame is None:
+            print(f"[WARN] Could not load frame: {frame_file}")
+            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})")
+            results.append({
+                'frame': frame_file.name,
+                'cursor_active': True,
+                'x': pos[0],
+                'y': pos[1],
+                'confidence': conf,
+                'template': template_name
+            })
+        else:
+            print(f"{frame_file.name}: Cursor disabled")
+            results.append({
+                'frame': frame_file.name,
+                'cursor_active': False,
+                'x': None,
+                'y': None,
+                'confidence': conf,
+                'template': None
+            })
+
+    with open(output_json_path, 'w') as f:
+        json.dump(results, f, indent=2)
+
+    print(f"[DONE] Cursor tracking results saved to: {output_json_path}")