Update app.py

app.py

@@ -28,7 +28,7 @@ app = FastAPI(
 HF_TOKEN = os.getenv("HF_TOKEN", "")
 HF_DATASET_ID = os.getenv("HF_DATASET_ID", "Fred808/BG3")  # Source dataset with zips
 HF_OUTPUT_DATASET_ID = os.getenv("HF_OUTPUT_DATASET_ID", "Fred808/data")  # Results dataset
-HF_STATE_FILE = os.getenv("HF_STATE_FILE", "processing_state_cursors2.json")
+HF_STATE_FILE = os.getenv("HF_STATE_FILE", "processing_state_cursors.json")
 TEMP_DATASET_DIR = Path("temp_cursor_detection")
 TEMP_DATASET_DIR.mkdir(exist_ok=True)
 
@@ -279,17 +279,26 @@ class DatasetProgress:
 # Global progress tracker
 dataset_progress = DatasetProgress()
 
-async def process_image(image_bytes: bytes, threshold: float = 0.8) -> dict:
+async def process_image(image_path: Path, threshold: float = 0.8) -> dict:
     """Process a single image and return cursor detection results."""
     try:
-        np_array = np.frombuffer(image_bytes, np.uint8)
-        frame = cv2.imdecode(np_array, cv2.IMREAD_UNCHANGED)
-        
+        # Read image with OpenCV directly
+        frame = cv2.imread(str(image_path), cv2.IMREAD_UNCHANGED)
         if frame is None:
-            raise ValueError("Could not decode image")
+            raise ValueError(f"Could not read image: {image_path}")
 
+        print(f"[DETECT] Processing image {image_path.name}, shape: {frame.shape}")
+        
+        # Run cursor detection
         pos, conf, template_name = detect_cursor_in_frame_multi(frame, CURSOR_TEMPLATES, threshold)
+        
+        # Log detection result
+        if pos is not None:
+            print(f"[DETECT] Found cursor in {image_path.name} at {pos} using template {template_name} (conf: {conf:.3f})")
+        else:
+            print(f"[DETECT] No cursor found in {image_path.name} (best conf: {conf:.3f})")
 
+        # Handle confidence values
         confidence = float(conf)
         if confidence == float('inf') or confidence == float('-inf'):
             confidence = 1.0 if confidence > 0 else 0.0
@@ -303,7 +312,8 @@ async def process_image(image_bytes: bytes, threshold: float = 0.8) -> dict:
             'image_shape': list(frame.shape)
         }
     except Exception as e:
-        raise ValueError(f"Error processing image: {str(e)}")
+        print(f"[ERROR] Failed to process {image_path.name}: {str(e)}")
+        raise ValueError(f"Error processing image {image_path.name}: {str(e)}")
 
 async def dataset_task(start_index: int = 1):
     """Main dataset processing loop for processing zip files."""
@@ -376,17 +386,26 @@ async def dataset_task(start_index: int = 1):
 
                     # Process all images in the zip
                     results = []
-                    for image_path in image_paths:
+                    print(f"[DATASET] Starting cursor detection on {len(image_paths)} images...")
+                    
+                    for i, image_path in enumerate(image_paths, 1):
                         try:
-                            with open(image_path, 'rb') as f:
-                                content = f.read()
-                            
                             # Process image for cursor detection
-                            image_result = await process_image(content)
+                            print(f"[DATASET] Processing image {i}/{len(image_paths)}: {image_path.name}")
+                            image_result = await process_image(image_path)
+                            
+                            # Add file information
                             image_result['image_name'] = image_path.name
                             image_result['image_path'] = str(image_path.relative_to(extract_dir))
+                            
                             results.append(image_result)
                             
+                            # Log result
+                            if image_result['cursor_active']:
+                                print(f"[DATASET] ✓ Found cursor in {image_path.name}")
+                            else:
+                                print(f"[DATASET] ✗ No cursor found in {image_path.name}")
+                            
                         except Exception as e:
                             print(f"[DATASET] Error processing {image_path.name}: {e}")
                             continue
@@ -446,6 +465,15 @@ async def dataset_task(start_index: int = 1):
 @app.on_event("startup")
 async def startup_event():
     """Load templates when the application starts."""
+    if not CURSOR_TEMPLATES_DIR.exists():
+        print(f"Creating cursor templates directory: {CURSOR_TEMPLATES_DIR}")
+        CURSOR_TEMPLATES_DIR.mkdir(parents=True, exist_ok=True)
+    
+    if not list(CURSOR_TEMPLATES_DIR.glob('*.png')):
+        print("WARNING: No cursor template files found in cursors directory!")
+        print(f"Please add cursor template PNG files to: {CURSOR_TEMPLATES_DIR}")
+        print("The server will start but cursor detection will not work without templates.")
+    
     load_cursor_templates()
 
 @app.post('/start_dataset')
@@ -500,9 +528,30 @@ async def track_cursor_endpoint(
             detail="Cursor templates are not loaded."
         )
 
-    content = await file.read()
-    results = await process_image(content, threshold)
-    return JSONResponse(content=results)
+    try:
+        # Save uploaded file to temporary path
+        temp_file = TEMP_DATASET_DIR / "temp_image"
+        temp_file.parent.mkdir(parents=True, exist_ok=True)
+        
+        content = await file.read()
+        with open(temp_file, 'wb') as f:
+            f.write(content)
+        
+        # Process image
+        results = await process_image(temp_file, threshold)
+        
+        # Cleanup
+        try:
+            os.remove(temp_file)
+        except Exception:
+            pass
+            
+        return JSONResponse(content=results)
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error processing image: {str(e)}"
+        )
 
 @app.post("/track_cursor_url")
 async def track_cursor_url_endpoint(