Update vision_analyzer.py

vision_analyzer.py

@@ -8,46 +8,25 @@ import re
 import threading
 from typing import Dict, List, Set, Optional
 from huggingface_hub import HfApi, list_repo_files
-from fastapi import FastAPI, File, UploadFile, Form
-from fastapi.responses import JSONResponse
+
+import cv2
+import numpy as np
 from pathlib import Path
 import smtplib
 from email.message import EmailMessage
-import tempfile
-import rarfile
-import zipfile
-import cv2
-import numpy as np
-from PIL import Image
-import torch
-from transformers import AutoProcessor, AutoModelForCausalLM
-from fastapi.staticfiles import StaticFiles
-
-from openai import OpenAI
-
-# Initialize FastAPI
-app = FastAPI()
 
 # ==== CONFIGURATION ====
 HF_TOKEN = os.getenv("HF_TOKEN", "")
 SOURCE_REPO_ID = os.getenv("SOURCE_REPO", "Fred808/BG1")
 
-BASE_URL = "https://openrouter.ai/api/v1"
-SITE_URL = os.getenv("SITE_URL", "https://fred808-vis1.hf.space")
-SITE_TITLE = os.getenv("SITE_TITLE", "FrameAnalyzer")
-LOG_PATH = os.getenv("LOG_PATH", "openrouter_responses.json")
-
-
 # Path Configuration
 DOWNLOAD_FOLDER = "downloads"
 EXTRACT_FOLDER = "extracted"
 FRAMES_OUTPUT_FOLDER = "extracted_frames"
-ANALYSIS_OUTPUT_FOLDER = "analysis_results"
 
 os.makedirs(DOWNLOAD_FOLDER, exist_ok=True)
 os.makedirs(EXTRACT_FOLDER, exist_ok=True)
 os.makedirs(FRAMES_OUTPUT_FOLDER, exist_ok=True)
-os.makedirs(ANALYSIS_OUTPUT_FOLDER, exist_ok=True)
 
 # State Files
 DOWNLOAD_STATE_FILE = "download_progress.json"
@@ -61,7 +40,10 @@ MAX_RETRIES = 3
 MIN_FREE_SPACE_GB = 2  # Minimum free space in GB before processing
 
 # Frame Extraction Parameters
-DEFAULT_FPS = 0.1  # Default frames per second for extraction
+DEFAULT_FPS = 3 # Default frames per second for extraction
+
+# Cursor Tracking Parameters
+CURSOR_THRESHOLD = 0.8
 
 # Initialize HF API
 hf_api = HfApi(token=HF_TOKEN)
@@ -76,14 +58,11 @@ processing_status = {
     "extracted_courses": 0,
     "extracted_videos": 0,
     "extracted_frames_count": 0,
-    "analyzed_frames_count": 0,
+    "tracked_cursors_count": 0,
     "last_update": None,
     "logs": []
 }
 
-
-
-
 def log_message(message: str):
     """Log messages with timestamp"""
     timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
@@ -97,7 +76,7 @@ def log_message(message: str):
 def log_failed_file(filename: str, error: str):
     """Log failed files to persistent file"""
     with open(FAILED_FILES_LOG, "a") as f:
-        f.write(f'{time.strftime("%Y-%m-%d %H:%M:%S")} - {filename}: {error}\n')
+        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - {filename}: {error}\n")
 
 def get_disk_usage(path: str) -> Dict[str, float]:
     """Get disk usage statistics in GB"""
@@ -108,7 +87,7 @@ def get_disk_usage(path: str) -> Dict[str, float]:
     return {"total": total, "free": free, "used": used}
 
 def check_disk_space(path: str = ".") -> bool:
-    """Check if there\'s enough disk space"""
+    """Check if there's enough disk space"""
     disk_info = get_disk_usage(path)
     if disk_info["free"] < MIN_FREE_SPACE_GB:
         log_message(f'⚠️ Low disk space: {disk_info["free"]:.2f}GB free, {disk_info["used"]:.2f}GB used')
@@ -245,7 +224,8 @@ def extract_with_retry(rar_path: str, output_dir: str, max_retries: int = 2) ->
             time.sleep(1)
     
     return False
-
+        
+# --- Frame Extraction Utilities ---
 def ensure_dir(path):
     os.makedirs(path, exist_ok=True)
 
@@ -258,10 +238,12 @@ def extract_frames(video_path, output_dir, fps=DEFAULT_FPS):
         log_message(f"[ERROR] Failed to open video file: {video_path}")
         return 0
     video_fps = cap.get(cv2.CAP_PROP_FPS)
+    # log_message(f"[DEBUG] Video FPS: {video_fps}")
     if not video_fps or video_fps <= 0:
         video_fps = 30  # fallback if FPS is not available
         log_message(f"[WARN] Using fallback FPS: {video_fps}")
     frame_interval = int(round(video_fps / fps))
+    # log_message(f"[DEBUG] Frame interval: {frame_interval}")
     frame_idx = 0
     saved_idx = 1
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
@@ -269,21 +251,38 @@ def extract_frames(video_path, output_dir, fps=DEFAULT_FPS):
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
+            # log_message(f"[DEBUG] No more frames to read at frame_idx {frame_idx}.")
             break
         if frame_idx % frame_interval == 0:
-            if saved_idx <= 10: # Limit to 10 frames for testing
-                frame_name = f"{saved_idx:04d}.png"
-                cv2.imwrite(str(Path(output_dir) / frame_name), frame)
-                saved_idx += 1
-            else:
-                break # Stop extracting after 10 frames
+            frame_name = f"{saved_idx:04d}.png"
+            cv2.imwrite(str(Path(output_dir) / frame_name), frame)
+            # log_message(f"[DEBUG] Saved frame {frame_idx} as {frame_name}")
+            saved_idx += 1
         frame_idx += 1
     cap.release()
     log_message(f"Extracted {saved_idx-1} frames from {video_path} to {output_dir}")
     return saved_idx - 1
 
+# --- Cursor Tracking Utilities ---
+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:
+        # Use alpha channel as mask (nonzero alpha = 255)
+        return (template_img[:, :, 3] > 0).astype(np.uint8) * 255
+    return None
+
+
+
 def process_rar_file(rar_path: str) -> bool:
-    """Process a single RAR file with new frame processing"""
+    """Process a single RAR file - extract, then process videos for frames and cursor tracking"""
     filename = os.path.basename(rar_path)
     processing_status["current_file"] = filename
     
@@ -307,39 +306,39 @@ def process_rar_file(rar_path: str) -> bool:
         if not extract_with_retry(rar_path, extract_dir):
             raise Exception("RAR extraction failed")
         
-        # Process video files
-        video_files = []
-        for root, _, files in os.walk(extract_dir):
+        # Count extracted files
+        file_count = 0
+        video_files_found = []
+        for root, dirs, files in os.walk(extract_dir):
             for file in files:
+                file_count += 1
                 if file.lower().endswith((".mp4", ".avi", ".mov", ".mkv")):
-                    video_files.append(os.path.join(root, file))
+                    video_files_found.append(os.path.join(root, file))
         
         processing_status["extracted_courses"] += 1
-        log_message(f"✅ Extracted {len(video_files)} videos from \'{course_name}\'")
+        log_message(f"✅ Successfully extracted \'{course_name}\' ({file_count} files, {len(video_files_found)} videos)")
 
-        # Process each video
-        for video_path in video_files:
+        # Process video files for frame extraction and cursor tracking
+        for video_path in video_files_found:
             video_filename = Path(video_path).name
-            video_filename_clean = video_filename.replace(".", "_")
-            frames_dir = os.path.join(FRAMES_OUTPUT_FOLDER, f"{course_name}_{video_filename_clean}_frames")
-            ensure_dir(frames_dir)
-            
-            # Extract frames
-            extracted_count = extract_frames(video_path, frames_dir, DEFAULT_FPS)
-            if extracted_count == 0:
-                raise Exception(f"No frames extracted from {video_filename}")
+            # Create a unique output directory for frames for each video
+            frames_output_dir = os.path.join(FRAMES_OUTPUT_FOLDER, f"{course_name}_{video_filename.replace('.', '_')}_frames")
+            ensure_dir(frames_output_dir)
             
-            processing_status["extracted_frames_count"] += extracted_count
-            
-            # Analyze frames
-            video_filename_clean = video_filename.replace(".", "_")
-            analysis_output = os.path.join(ANALYSIS_OUTPUT_FOLDER, f"{course_name}_{video_filename_clean}_analysis.json")
-            if process_video_frames(frames_dir, video_filename, analysis_output):
-                processing_status["analyzed_frames_count"] += extracted_count
+            extracted_frames_count = extract_frames(video_path, frames_output_dir, fps=DEFAULT_FPS)
+            processing_status["extracted_frames_count"] += extracted_frames_count
+            if extracted_frames_count > 0:
                 processing_status["extracted_videos"] += 1
+                log_message(f"[INFO] Extracted {extracted_frames_count} frames from {video_filename}")
+
+                # Perform cursor tracking on the extracted frames
+                cursor_output_json = os.path.join(CURSOR_TRACKING_OUTPUT_FOLDER, f"{course_name}_{video_filename.replace('.', '_')}_cursor_data.json")
+                tracked_cursors = track_cursor(frames_output_dir, CURSOR_TEMPLATES_DIR, cursor_output_json, threshold=CURSOR_THRESHOLD)
+                processing_status["tracked_cursors_count"] += tracked_cursors
+                log_message(f"[INFO] Tracked {tracked_cursors} cursors in frames from {video_filename}")
             else:
-                raise Exception(f"Frame analysis failed for {video_filename}")
-            
+                log_message(f"[WARN] No frames extracted from {video_filename}")
+
         return True
     
     except Exception as e:
@@ -351,18 +350,14 @@ def process_rar_file(rar_path: str) -> bool:
     finally:
         processing_status["current_file"] = None
 
-
-
-        
-
 def main_processing_loop(start_index: int = 0):
-    """Main processing workflow - extraction, frame extraction, and vision analysis"""
+    """Main processing workflow - extraction, frame extraction, and cursor tracking"""
     processing_status["is_running"] = True
     
     try:
         # Load state
         processed_rars = load_json_state(PROCESS_STATE_FILE, {"processed_rars": []})["processed_rars"]
-        download_state = load_json_state(DOWNLOAD_STATE_FILE, {"next_download_index": 0})
+        download_state = load_json_state(DOWNLOAD_STATE_FILE, {"next_download_index": 5})
         
         # Use start_index if provided, otherwise use the saved state
         next_index = start_index if start_index > 0 else download_state["next_download_index"]
@@ -432,11 +427,9 @@ def main_processing_loop(start_index: int = 0):
             
             # Status update
             log_message(f"📊 Progress: {next_index}/{len(rar_files)} files processed")
-            log_message(f"📊 Extracted: {processing_status['extracted_courses']} courses")
-            log_message(f"📊 Videos Processed: {processing_status['extracted_videos']} videos")
-            log_message(f"📊 Frames Extracted: {processing_status['extracted_frames_count']} frames")
-            log_message(f"📊 Frames Analyzed: {processing_status['analyzed_frames_count']} frames")
-            log_message(f"📊 Failed: {processing_status['failed_files']} files")
+            log_message(f'📊 Extracted: {processing_status["extracted_courses"]} courses')
+            log_message(f'📊 Videos Processed: {processing_status["extracted_videos"]}')
+            log_message(f'📊 Frames Extracted: {processing_status["extracted_frames_count"]}')
             
             if next_index < len(rar_files):
                 log_message(f"🔄 Run the script again to process the next file: {os.path.basename(rar_files[next_index])}")
@@ -445,8 +438,8 @@ def main_processing_loop(start_index: int = 0):
         else:
             log_message("✅ All files have been processed!")
         
-        log_message(f"🎉 Processing complete!")
-        log_message(f"📊 Final stats: {processing_status['extracted_courses']} courses extracted, {processing_status['extracted_videos']} videos processed, {processing_status['extracted_frames_count']} frames extracted, {processing_status['analyzed_frames_count']} frames analyzed")
+        log_message("🎉 Processing complete!")
+        log_message(f'📊 Final stats: {processing_status["extracted_courses"]} courses extracted, {processing_status["extracted_videos"]} videos processed, {processing_status["extracted_frames_count"]} frames extracted, {processing_status["tracked_cursors_count"]} cursors tracked')
         
     except KeyboardInterrupt:
         log_message("⏹️ Processing interrupted by user")
@@ -456,31 +449,16 @@ def main_processing_loop(start_index: int = 0):
         processing_status["is_running"] = False
         cleanup_temp_files()
 
-
-
-@app.get("/health")
-async def health_check():
-    """Health check endpoint."""
-    return JSONResponse(content={
-        "status": "healthy",
-        "model": "GIT",
-        "note": "Now using GIT model."
-    })
-
-@app.get("/status")
-async def get_processing_status():
-    """Get current processing status."""
-    return JSONResponse(content=processing_status)
-
-# Expose necessary functions and variables
+# Expose necessary functions and variables for download_api.py
 __all__ = [
     "main_processing_loop",
     "processing_status",
-    "ANALYSIS_OUTPUT_FOLDER",
     "log_message",
-    "analyze_single_frame",
     "extract_frames",
     "DEFAULT_FPS",
     "ensure_dir"
 ]
 
+
+
+