Update app.py

app.py

@@ -5,6 +5,8 @@ import os
 import threading
 import time
 import urllib.parse
+import zipfile
+import shutil
 from fastapi import FastAPI, UploadFile, File, HTTPException, Form
 from fastapi.responses import JSONResponse
 import json
@@ -24,11 +26,44 @@ app = FastAPI(
 
 # --- Environment Configuration ---
 HF_TOKEN = os.getenv("HF_TOKEN", "")
-HF_DATASET_ID = os.getenv("HF_DATASET_ID", "Fred808/data")  # Dataset to store results
+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_cursors.json")
 TEMP_DATASET_DIR = Path("temp_cursor_detection")
 TEMP_DATASET_DIR.mkdir(exist_ok=True)
 
+def _get_zip_file_list_from_hf() -> list:
+    """Return sorted list of zip file paths from HF_DATASET_ID."""
+    try:
+        api = HfApi(token=HF_TOKEN)
+        files = api.list_repo_files(repo_id=HF_DATASET_ID, repo_type="dataset")
+        zip_files = sorted([f for f in files if f.startswith('frames_zips/') and f.lower().endswith('.zip')])
+        print(f"[DATASET] Found {len(zip_files)} zip files in {HF_DATASET_ID}.")
+        return zip_files
+    except Exception as e:
+        print(f"[DATASET] Error listing HF dataset files: {e}")
+        return []
+
+def _download_and_extract_zip(repo_path: str) -> Optional[Path]:
+    """Download zip from HF dataset and extract into a temp subfolder."""
+    try:
+        zip_local = hf_hub_download(repo_id=HF_DATASET_ID, filename=repo_path, repo_type="dataset", token=HF_TOKEN)
+        zip_name = Path(repo_path).name
+        extract_dir = TEMP_DATASET_DIR / zip_name.replace('.zip','')
+        if extract_dir.exists():
+            shutil.rmtree(extract_dir)
+        extract_dir.mkdir(parents=True, exist_ok=True)
+        with zipfile.ZipFile(zip_local, 'r') as z:
+            z.extractall(extract_dir)
+        try:
+            os.remove(zip_local)
+        except Exception:
+            pass
+        return extract_dir
+    except Exception as e:
+        print(f"[DATASET] Error downloading/extracting {repo_path}: {e}")
+        return None
+
 # Global variable to store loaded templates
 CURSOR_TEMPLATES: Dict[str, np.ndarray] = {}
 CURSOR_TEMPLATES_DIR = Path("cursors")
@@ -201,35 +236,25 @@ def _unlock_file_as_processed(image_name: str, state: dict, next_index: int) ->
     state['next_download_index'] = next_index
     return _upload_hf_state(state)
 
-def _get_image_list_from_hf() -> list:
-    """Return sorted list of image file paths from HF_DATASET_ID."""
-    try:
-        api = HfApi(token=HF_TOKEN)
-        files = api.list_repo_files(repo_id=HF_DATASET_ID, repo_type="dataset")
-        image_files = sorted([f for f in files if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
-        print(f"[DATASET] Found {len(image_files)} image files in {HF_DATASET_ID}.")
-        return image_files
-    except Exception as e:
-        print(f"[DATASET] Error listing HF dataset files: {e}")
-        return []
 
-def _upload_cursor_results(image_name: str, results: dict) -> bool:
-    """Upload cursor detection results JSON to dataset."""
+
+def _upload_cursor_results(zip_name: str, results: dict) -> bool:
+    """Upload cursor detection results JSON to output dataset."""
     try:
-        filename = Path(image_name).with_suffix('.json').name
+        filename = Path(zip_name).with_suffix('.json').name
         content = json.dumps(results, indent=2, ensure_ascii=False).encode('utf-8')
         api = HfApi(token=HF_TOKEN)
         api.upload_file(
             path_or_fileobj=io.BytesIO(content),
             path_in_repo=f"cursor_results/{filename}",
-            repo_id=HF_DATASET_ID,
+            repo_id=HF_OUTPUT_DATASET_ID,  # Using output dataset
             repo_type="dataset",
-            commit_message=f"Cursor detection results for {image_name}"
+            commit_message=f"Cursor detection results for {zip_name}"
         )
-        print(f"[DATASET] Uploaded results for {image_name} to {HF_DATASET_ID}.")
+        print(f"[DATASET] Uploaded results for {zip_name} to {HF_OUTPUT_DATASET_ID}.")
         return True
     except Exception as e:
-        print(f"[DATASET] Failed to upload results for {image_name}: {e}")
+        print(f"[DATASET] Failed to upload results for {zip_name}: {e}")
         return False
 
 class DatasetProgress:
@@ -281,7 +306,7 @@ async def process_image(image_bytes: bytes, threshold: float = 0.8) -> dict:
         raise ValueError(f"Error processing image: {str(e)}")
 
 async def dataset_task(start_index: int = 1):
-    """Main dataset processing loop."""
+    """Main dataset processing loop for processing zip files."""
     global dataset_progress
 
     dataset_progress = DatasetProgress()
@@ -299,88 +324,115 @@ async def dataset_task(start_index: int = 1):
 
     try:
         state = await asyncio.to_thread(_load_hf_state)
-        image_list = await asyncio.to_thread(_get_image_list_from_hf)
+        zip_list = await asyncio.to_thread(_get_zip_file_list_from_hf)
         
-        if not image_list:
-            err = "No images found in dataset"
+        if not zip_list:
+            err = "No zip files found in dataset"
             dataset_progress.status = "error"
             dataset_progress.error = err
             print(f"[DATASET] {err}")
             return False
             
-        dataset_progress.total_images = len(image_list)
+        dataset_progress.total_images = len(zip_list)
         dataset_progress.status = "processing"
 
         if start_index < 1:
             start_index = 1
 
-        for idx in range(start_index-1, len(image_list)):
+        for idx in range(start_index-1, len(zip_list)):
             try:
-                image_path = image_list[idx]
-                image_name = Path(image_path).name
-                print(f"[DATASET] Processing image {idx + 1}/{len(image_list)}: {image_name}")
+                zip_path = zip_list[idx]
+                zip_name = Path(zip_path).name
+                print(f"[DATASET] Processing zip {idx + 1}/{len(zip_list)}: {zip_name}")
 
-                file_state = state.get('file_states', {}).get(image_name)
+                file_state = state.get('file_states', {}).get(zip_name)
                 if file_state == 'processed':
-                    print(f"[DATASET] Skipping {image_name}: already processed.")
+                    print(f"[DATASET] Skipping {zip_name}: already processed.")
                     dataset_progress.processed_images += 1
                     continue
                 if file_state == 'processing':
-                    print(f"[DATASET] Skipping {image_name}: currently processing by another worker.")
+                    print(f"[DATASET] Skipping {zip_name}: currently processing by another worker.")
                     continue
 
                 # Try to lock
-                locked = await asyncio.to_thread(_lock_file_for_processing, image_name, state)
+                locked = await asyncio.to_thread(_lock_file_for_processing, zip_name, state)
                 if not locked:
-                    print(f"[DATASET] Could not lock {image_name}; skipping.")
+                    print(f"[DATASET] Could not lock {zip_name}; skipping.")
                     continue
 
                 try:
-                    # Download and process image
-                    print(f"[DATASET] Downloading {image_name}...")
-                    image_bytes = await asyncio.to_thread(
-                        lambda: hf_hub_download(
-                            repo_id=HF_DATASET_ID,
-                            filename=image_path,
-                            repo_type="dataset",
-                            token=HF_TOKEN
-                        )
-                    )
-                    
-                    with open(image_bytes, 'rb') as f:
-                        content = f.read()
-                    
-                    # Process image
-                    results = await process_image(content)
-                    results['image_name'] = image_name
-                    results['image_path'] = image_path
+                    # Download and extract zip
+                    print(f"[DATASET] Downloading and extracting {zip_name}...")
+                    extract_dir = await asyncio.to_thread(_download_and_extract_zip, zip_path)
+                    if not extract_dir:
+                        print(f"[DATASET] Failed to download/extract {zip_name}; marking failed.")
+                        state['file_states'][zip_name] = 'failed'
+                        await asyncio.to_thread(_upload_hf_state, state)
+                        continue
+
+                    # Find all images in extracted directory
+                    image_paths = [p for p in extract_dir.rglob('*') if p.is_file() and p.suffix.lower() in ('.jpg','.jpeg','.png')]
+                    print(f"[DATASET] Found {len(image_paths)} images in {zip_name}")
+
+                    # Process all images in the zip
+                    results = []
+                    for image_path in image_paths:
+                        try:
+                            with open(image_path, 'rb') as f:
+                                content = f.read()
+                            
+                            # Process image for cursor detection
+                            image_result = await process_image(content)
+                            image_result['image_name'] = image_path.name
+                            image_result['image_path'] = str(image_path.relative_to(extract_dir))
+                            results.append(image_result)
+                            
+                        except Exception as e:
+                            print(f"[DATASET] Error processing {image_path.name}: {e}")
+                            continue
+
+                    # Create combined results for the zip
+                    zip_results = {
+                        'zip_name': zip_name,
+                        'zip_path': zip_path,
+                        'total_images': len(image_paths),
+                        'processed_images': len(results),
+                        'results': results
+                    }
                     
-                    # Upload results
-                    uploaded = await asyncio.to_thread(_upload_cursor_results, image_name, results)
+                    # Upload combined results
+                    uploaded = await asyncio.to_thread(_upload_cursor_results, zip_name, zip_results)
                     
                     if uploaded:
                         next_index = idx + 2  # next 1-based index
-                        ok = await asyncio.to_thread(_unlock_file_as_processed, image_name, state, next_index)
+                        ok = await asyncio.to_thread(_unlock_file_as_processed, zip_name, state, next_index)
                         if not ok:
-                            print(f"[DATASET] Warning: processed but failed to update state for {image_name}.")
+                            print(f"[DATASET] Warning: processed but failed to update state for {zip_name}.")
                         dataset_progress.processed_images += 1
-                        print(f"[DATASET] Successfully processed {image_name}")
+                        print(f"[DATASET] Successfully processed {zip_name}")
                     else:
-                        print(f"[DATASET] Failed to upload results for {image_name}")
-                        state['file_states'][image_name] = 'failed'
+                        print(f"[DATASET] Failed to upload results for {zip_name}")
+                        state['file_states'][zip_name] = 'failed'
                         await asyncio.to_thread(_upload_hf_state, state)
-                        
+                    
                 except Exception as e:
-                    print(f"[DATASET] Error processing {image_name}: {e}")
-                    state['file_states'][image_name] = 'failed'
+                    print(f"[DATASET] Error processing zip {zip_name}: {e}")
+                    state['file_states'][zip_name] = 'failed'
                     await asyncio.to_thread(_upload_hf_state, state)
                     continue
+                finally:
+                    # Cleanup extracted directory
+                    try:
+                        if extract_dir and extract_dir.exists():
+                            shutil.rmtree(extract_dir)
+                    except Exception as e:
+                        print(f"[DATASET] Warning: Failed to clean up {extract_dir}: {e}")
 
             except Exception as e:
-                print(f"[DATASET] Error in image loop: {e}")
+                print(f"[DATASET] Error in zip processing loop: {e}")
                 continue
 
-        print(f"[DATASET] Task completed. Processed {dataset_progress.processed_images}/{len(image_list)} images.")
+        print(f"[DATASET] Task completed. Processed {dataset_progress.processed_images}/{len(zip_list)} zip files.")
         dataset_progress.status = "completed"
         return True