Update download_api.py

download_api.py

@@ -1,104 +1,236 @@
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import FileResponse
-from pathlib import Path
 import os
 import json
+import time
 import threading
-import requests
-from io import BytesIO
-from PIL import Image
-import numpy as np
-from cursor_tracker import track_cursor_from_images
+from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+import uvicorn
+
+# Import from cursor_tracker
+from cursor_tracker import (
+    main_processing_loop,
+    processing_status,
+    CURSOR_TRACKING_OUTPUT_FOLDER,
+    CURSOR_TEMPLATES_DIR,
+    log_message
+)
+
+# FastAPI App Definition
+app = FastAPI(title="Cursor Tracking API", description="API to access cursor tracking results", version="1.0.0")
+
+# Add CORS middleware to allow cross-origin requests
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allows all origins
+    allow_credentials=True,
+    allow_methods=["*"],  # Allows all methods
+    allow_headers=["*"],
+)
+
+# Global variable to track if processing is running
+processing_thread = None
 
-app = FastAPI()
+@app.on_event("startup")
+async def startup_event():
+    """Run the processing loop in the background when the API starts"""
+    global processing_thread
+    if not (processing_thread and processing_thread.is_alive()):
+        log_message("🚀 Starting RAR extraction, frame extraction, and cursor tracking pipeline in background...")
+        processing_thread = threading.Thread(target=main_processing_loop)
+        processing_thread.daemon = True
+        processing_thread.start()
 
-ANNOTATIONS_DIR = Path("annotations").resolve()
-CURSOR_TEMPLATES_DIR = Path("cursors").resolve()
-ANNOTATION_OUTPUT = ANNOTATIONS_DIR / "blender1.json"
+@app.get("/")
+async def root():
+    """Root endpoint with API information"""
+    return {
+        "message": "Cursor Tracking API",
+        "version": "1.0.0",
+        "endpoints": {
+            "/status": "Get processing status",
+            "/cursor-data": "List all cursor tracking JSON files",
+            "/cursor-data/{filename}": "Get specific cursor tracking data",
+            "/start-processing": "Start the RAR processing pipeline",
+            "/stop-processing": "Stop the RAR processing pipeline"
+        }
+    }
 
-REMOTE_FRAMES_LIST_URL = "https://fred808-cu2.hf.space/frames"
-REMOTE_FRAME_BASE_URL = "https://fred808-cu2.hf.space/frames"
+@app.get("/status")
+async def get_status():
+    """Get current processing status"""
+    return {
+        "processing_status": processing_status,
+        "cursor_tracking_folder": CURSOR_TRACKING_OUTPUT_FOLDER,
+        "folder_exists": os.path.exists(CURSOR_TRACKING_OUTPUT_FOLDER)
+    }
 
-ANNOTATIONS_DIR.mkdir(parents=True, exist_ok=True)
-cursor_tracking_started = False
+@app.get("/cursor-data")
+async def list_cursor_data():
+    """List all available cursor tracking JSON files"""
+    if not os.path.exists(CURSOR_TRACKING_OUTPUT_FOLDER):
+        return {"files": [], "message": "Cursor tracking output folder does not exist yet"}
+    
+    json_files = []
+    for file in os.listdir(CURSOR_TRACKING_OUTPUT_FOLDER):
+        if file.endswith(".json"):
+            file_path = os.path.join(CURSOR_TRACKING_OUTPUT_FOLDER, file)
+            file_stats = os.stat(file_path)
+            json_files.append({
+                "filename": file,
+                "size_bytes": file_stats.st_size,
+                "modified_time": time.ctime(file_stats.st_mtime),
+                "download_url": f"/cursor-data/{file}"
+            })
+    
+    return {
+        "files": json_files,
+        "total_files": len(json_files),
+        "folder_path": CURSOR_TRACKING_OUTPUT_FOLDER
+    }
 
-def append_to_annotation_file(new_data: list, output_path: str):
+@app.get("/cursor-data/{filename}")
+async def get_cursor_data(filename: str):
+    """Get specific cursor tracking data by filename"""
+    if not filename.endswith(".json"):
+        raise HTTPException(status_code=400, detail="File must be a JSON file")
+    
+    file_path = os.path.join(CURSOR_TRACKING_OUTPUT_FOLDER, filename)
+    
+    if not os.path.exists(file_path):
+        raise HTTPException(status_code=404, detail=f"File {filename} not found")
+    
     try:
-        if os.path.exists(output_path):
-            with open(output_path, "r", encoding="utf-8") as f:
-                existing = json.load(f)
-        else:
-            existing = []
-    except Exception:
-        existing = []
-
-    with open(output_path, "w", encoding="utf-8") as f:
-        json.dump(existing + new_data, f, indent=2)
-
-def run_cursor_tracker_in_batches(batch_size=5000):
-    global cursor_tracking_started
-    if cursor_tracking_started:
-        return
-    cursor_tracking_started = True
-
+        with open(file_path, "r") as f:
+            data = json.load(f)
+        
+        # Add metadata
+        file_stats = os.stat(file_path)
+        response_data = {
+            "filename": filename,
+            "file_size_bytes": file_stats.st_size,
+            "modified_time": time.ctime(file_stats.st_mtime),
+            "total_frames": len(data),
+            "cursor_active_frames": len([frame for frame in data if frame.get("cursor_active", False)]),
+            "data": data
+        }
+        
+        return response_data
+        
+    except json.JSONDecodeError:
+        raise HTTPException(status_code=500, detail=f"Invalid JSON in file {filename}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error reading file {filename}: {str(e)}")
+
+@app.post("/start-processing")
+async def start_processing(background_tasks: BackgroundTasks):
+    """Start the RAR processing pipeline in the background"""
+    global processing_thread
+    
+    if processing_thread and processing_thread.is_alive():
+        return {"message": "Processing is already running", "status": "already_running"}
+    
+    if processing_status["is_running"]:
+        return {"message": "Processing is already running", "status": "already_running"}
+    
+    # Start processing in a background thread
+    processing_thread = threading.Thread(target=main_processing_loop)
+    processing_thread.daemon = True
+    processing_thread.start()
+    
+    return {"message": "Processing started in background", "status": "started"}
+
+@app.post("/stop-processing")
+async def stop_processing():
+    """Stop the RAR processing pipeline"""
+    global processing_thread
+    
+    if not processing_status["is_running"] and (not processing_thread or not processing_thread.is_alive()):
+        return {"message": "No processing is currently running", "status": "not_running"}
+    
+    # Note: This is a graceful stop request. The actual stopping depends on the processing loop
+    # checking the processing_status["is_running"] flag
+    processing_status["is_running"] = False
+    
+    return {"message": "Stop signal sent to processing pipeline", "status": "stop_requested"}
+
+@app.get("/cursor-data/{filename}/summary")
+async def get_cursor_data_summary(filename: str):
+    """Get a summary of cursor tracking data without the full frame data"""
+    if not filename.endswith(".json"):
+        raise HTTPException(status_code=400, detail="File must be a JSON file")
+    
+    file_path = os.path.join(CURSOR_TRACKING_OUTPUT_FOLDER, filename)
+    
+    if not os.path.exists(file_path):
+        raise HTTPException(status_code=404, detail=f"File {filename} not found")
+    
     try:
-        res = requests.get(REMOTE_FRAMES_LIST_URL)
-        res.raise_for_status()
-        frame_list = res.json().get("frames", [])
-        total = len(frame_list)
-        print(f"[INFO] Found {total} frames.")
-
-        for start in range(0, total, batch_size):
-            end = min(start + batch_size, total)
-            batch = frame_list[start:end]
-            print(f"[BATCH] Fetching frames {start} to {end - 1}")
-
-            images = []
-            for fname in batch:
-                try:
-                    url = f"{REMOTE_FRAME_BASE_URL}/{fname}"
-                    r = requests.get(url)
-                    if r.status_code == 429:
-                        import time; time.sleep(1.0)
-                        r = requests.get(url)
-                    r.raise_for_status()
-                    img = Image.open(BytesIO(r.content)).convert("RGB")
-                    images.append((fname, np.array(img)))
-                    import time; time.sleep(0.3)
-                except Exception as e:
-                    print(f"[WARN] Could not fetch {fname}: {e}")
-
-            if not images:
-                continue
-
-            results = track_cursor_from_images(
-                images=images,
-                cursor_templates_dir=str(CURSOR_TEMPLATES_DIR),
-                output_json_path=None,
-                threshold=0.8,
-                return_results=True
-            )
-            append_to_annotation_file(results, str(ANNOTATION_OUTPUT))
-            print(f"[✅] Appended {len(results)} annotations to JSON.")
-
+        with open(file_path, "r") as f:
+            data = json.load(f)
+        
+        # Calculate summary statistics
+        total_frames = len(data)
+        cursor_active_frames = len([frame for frame in data if frame.get("cursor_active", False)])
+        cursor_inactive_frames = total_frames - cursor_active_frames
+        
+        # Get unique templates used
+        templates_used = set()
+        confidence_scores = []
+        
+        for frame in data:
+            if frame.get("cursor_active", False) and frame.get("template"):
+                templates_used.add(frame["template"])
+                if frame.get("confidence") is not None:
+                    confidence_scores.append(frame["confidence"])
+        
+        # Calculate confidence statistics
+        avg_confidence = sum(confidence_scores) / len(confidence_scores) if confidence_scores else 0
+        max_confidence = max(confidence_scores) if confidence_scores else 0
+        min_confidence = min(confidence_scores) if confidence_scores else 0
+        
+        file_stats = os.stat(file_path)
+        
+        summary = {
+            "filename": filename,
+            "file_size_bytes": file_stats.st_size,
+            "modified_time": time.ctime(file_stats.st_mtime),
+            "total_frames": total_frames,
+            "cursor_active_frames": cursor_active_frames,
+            "cursor_inactive_frames": cursor_inactive_frames,
+            "cursor_detection_rate": cursor_active_frames / total_frames if total_frames > 0 else 0,
+            "templates_used": list(templates_used),
+            "confidence_stats": {
+                "average": avg_confidence,
+                "maximum": max_confidence,
+                "minimum": min_confidence,
+                "total_measurements": len(confidence_scores)
+            }
+        }
+        
+        return summary
+        
+    except json.JSONDecodeError:
+        raise HTTPException(status_code=500, detail=f"Invalid JSON in file {filename}")
     except Exception as e:
-        print(f"[ERROR] Cursor tracking failed: {e}")
+        raise HTTPException(status_code=500, detail=f"Error reading file {filename}: {str(e)}")
+
+if __name__ == "__main__":
+    # Start the FastAPI server
+    print("Starting Cursor Tracking FastAPI Server...")
+    print("API Documentation will be available at: http://localhost:8000/docs")
+    print("API Root endpoint: http://localhost:8000/")
+    
+    # Ensure the cursor tracking output folder exists
+    os.makedirs(CURSOR_TRACKING_OUTPUT_FOLDER, exist_ok=True)
+    
+    uvicorn.run(
+        app, 
+        host="0.0.0.0", 
+        port=8000, 
+        log_level="info",
+        reload=False  # Set to False for production
+    )
 
-@app.on_event("startup")
-def startup_event():
-    threading.Thread(target=run_cursor_tracker_in_batches, daemon=True).start()
 
-@app.get("/download/{filename}")
-def download_file(filename: str):
-    file_path = ANNOTATIONS_DIR / filename
-    if not file_path.exists() or not file_path.is_file():
-        raise HTTPException(status_code=404, detail="File not found")
-    return FileResponse(str(file_path), filename=filename)
 
-@app.get("/")
-def root():
-    files = [f.name for f in ANNOTATIONS_DIR.glob("*.json")]
-    return {
-        "message": "Use /download/{filename} to get cursor tracking results.",
-        "available_files": files
-    }