Update app.py

app.py

@@ -1,233 +1,353 @@
 import os
-import re
 import json
 import time
-from typing import Dict, Any, List
-from urllib.parse import urlparse, parse_qs
-
-from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import JSONResponse
-
-try:
-    from huggingface_hub import HfApi
-    HF_AVAILABLE = True
-except Exception:
-    HfApi = None
-    HF_AVAILABLE = False
-
-# Directory to store compiled uploads
-BASE_DIR = os.path.dirname(__file__)
-UPLOAD_DIR = os.path.join(BASE_DIR, "uploads")
-os.makedirs(UPLOAD_DIR, exist_ok=True)
-
-app = FastAPI(title="Data Collection Server", description="Receives text/URLs from captioning/image servers, groups by course, compiles JSON and optionally uploads to HuggingFace.")
-
-# In-memory store for course data
-courses: Dict[str, Dict[str, Any]] = {}
-
-URL_RE = re.compile(r"https?://[\w\-\./?%&=:@,+~#]+")
-DONE_RE = re.compile(r"\b(done|finished|completed|complete)\b", re.IGNORECASE)
-
-HF_TOKEN = os.getenv("HF_TOKEN")
-HF_DATASET_REPO = os.getenv("HF_DATASET_REPO")  # e.g. "username/dataset-name"
-
-
-def extract_urls(text: str) -> List[str]:
-    return URL_RE.findall(text or "")
-
-
-def extract_course_from_url(url: str) -> str:
+import asyncio
+import aiohttp
+from typing import Dict, List, Set
+from urllib.parse import quote, urljoin
+from datetime import datetime
+from pathlib import Path
+
+# Path for storing caption data
+CAPTIONS_DIR = Path("captions_data")
+CAPTIONS_DIR.mkdir(exist_ok=True)
+
+def get_caption_file_path(course: str) -> Path:
+    """Get the path to the JSON file for storing course captions"""
+    safe_name = quote(course, safe='')
+    return CAPTIONS_DIR / f"{safe_name}_captions.json"
+
+def save_captions_to_file(course: str, captions: List[Dict]) -> None:
+    """Save captions to a JSON file"""
     try:
-        parsed = urlparse(url)
-        qs = parse_qs(parsed.query)
-        course = qs.get("course") or qs.get("Course") or qs.get("COURSE")
-        if course:
-            return course[0]
-    except Exception:
-        pass
-    return None
-
-
-def now_ts() -> str:
-    return time.strftime("%Y%m%dT%H%M%S")
-
-
-async def parse_request(request: Request) -> Dict[str, Any]:
-    """Read incoming request in any format and return a dict with keys: text, json, form, headers"""
-    payload = {"text": "", "json": None, "form": {}, "headers": dict(request.headers)}
+        file_path = get_caption_file_path(course)
+        with open(file_path, 'w', encoding='utf-8') as f:
+            json.dump(captions, f, indent=2, ensure_ascii=False)
+    except Exception as e:
+        print(f"Error saving captions for {course}: {e}")
 
-    # Try JSON
+def load_captions_from_file(course: str) -> List[Dict]:
+    """Load existing captions from JSON file"""
     try:
-        body = await request.json()
-        payload["json"] = body
-        # if it's a simple string payload inside JSON
-        if isinstance(body, str):
-            payload["text"] = body
-        elif isinstance(body, dict):
-            # flatten likely fields
-            for k in ["text", "caption", "message", "body", "content"]:
-                if k in body and isinstance(body[k], str):
-                    payload["text"] = body[k]
-                    break
-            # allow explicit course field
-            if "course" in body and isinstance(body["course"], str):
-                payload["course"] = body["course"]
-    except Exception:
-        # not JSON - try raw body
-        try:
-            raw = (await request.body()).decode("utf-8", errors="ignore")
-            payload["text"] = raw
-        except Exception:
-            payload["text"] = ""
-
-    # Try form (for multipart/form-data)
+        file_path = get_caption_file_path(course)
+        if file_path.exists():
+            with open(file_path, 'r', encoding='utf-8') as f:
+                return json.load(f)
+    except Exception as e:
+        print(f"Error loading captions for {course}: {e}")
+    return []
+
+# Configuration
+SOURCE_SERVER = "https://fred808-vs2.hf.space"
+CAPTION_SERVERS = [
+    "https://fred808-pill-3-1.hf.space/analyze",
+    "https://fred808-pil-4-1.hf.space/analyze",
+    "https://fred808-pil-4-2.hf.space/analyze",
+    "https://fred808-pil-4-3.hf.space/analyze",
+    "https://fred1012-fred1012-gw0j2h.hf.space/analyze",
+    "https://fred1012-fred1012-wqs6c2.hf.space/analyze",
+    "https://fred1012-fred1012-oncray.hf.space/analyze",
+    "https://fred1012-fred1012-4goge7.hf.space/analyze",
+    "https://fred1012-fred1012-z0eh7m.hf.space/analyze",
+    "https://fred1012-fred1012-u95rte.hf.space/analyze",
+    "https://fred1012-fred1012-igje22.hf.space/analyze",
+    "https://fred1012-fred1012-ibkuf8.hf.space/analyze",
+    "https://fred1012-fred1012-nwqthy.hf.space/analyze",
+    "https://fred1012-fred1012-4ldqj4.hf.space/analyze",
+    "https://fred1012-fred1012-pivlzg.hf.space/analyze",
+    "https://fred1012-fred1012-ptlc5u.hf.space/analyze",
+    "https://fred1012-fred1012-u7lh57.hf.space/analyze",
+    "https://fred1012-fred1012-q8djv1.hf.space/analyze",
+  "https://fredalone-fredalone-ozugrp.hf.space/analyze",
+  "https://fredalone-fredalone-9brxj2.hf.space/analyze",
+  "https://fredalone-fredalone-p8vq9a.hf.space/analyze",
+  "https://fredalone-fredalone-vbli2y.hf.space/analyze",
+  "https://fredalone-fredalone-uggger.hf.space/analyze",
+  "https://fredalone-fredalone-nmi7e8.hf.space/analyze",
+  "https://fredalone-fredalone-d1f26d.hf.space/analyze",
+  "https://fredalone-fredalone-461jp2.hf.space/analyze",
+  "https://fredalone-fredalone-3enfg4.hf.space/analyze",
+  "https://fredalone-fredalone-dqdbpv.hf.space/analyze",
+  "https://fredalone-fredalone-ivtjua.hf.space/analyze",
+  "https://fredalone-fredalone-6bezt2.hf.space/analyze",
+  "https://fredalone-fredalone-e0wfnk.hf.space/analyze",
+  "https://fredalone-fredalone-zu2t7j.hf.space/analyze",
+  "https://fredalone-fredalone-dqtv1o.hf.space/analyze",
+  "https://fredalone-fredalone-wclyog.hf.space/analyze",
+  "https://fredalone-fredalone-t27vig.hf.space/analyze",
+  "https://fredalone-fredalone-gahbxh.hf.space/analyze",
+  "https://fredalone-fredalone-kw2po4.hf.space/analyze",
+  "https://fredalone-fredalone-8h285h.hf.space/analyze"
+]
+MODEL_TYPE = "Florence-2-large"  # Explicitly request large model
+DATA_COLLECTION_SERVER = "https://fred808-flow.hf.space"
+
+# Tracking state
+processed_images: Dict[str, Set[str]] = {}  # {course: set(image_names)}
+course_captions: Dict[str, List[Dict]] = {}  # {course: [{image, caption, metadata}]}
+
+async def fetch_courses() -> List[str]:
+    """Fetch available courses from source server"""
+    async with aiohttp.ClientSession() as session:
+        async with session.get(f"{SOURCE_SERVER}/courses") as resp:
+            data = await resp.json()
+            if isinstance(data, dict) and 'courses' in data:
+                return [c['course_folder'] for c in data['courses'] if isinstance(c, dict)]
+            return []
+
+async def fetch_course_images(course: str) -> List[Dict]:
+    """Fetch images list for a course"""
+    course_frames = f"{course}_frames" if not course.endswith("_frames") else course
+    url = f"{SOURCE_SERVER}/images/{quote(course_frames)}"
+    async with aiohttp.ClientSession() as session:
+        async with session.get(url) as resp:
+            data = await resp.json()
+            if isinstance(data, dict) and 'images' in data:
+                return data['images']
+            return []
+
+async def get_caption(server: str, image_url: str) -> Dict:
+    """Get caption from a specific server"""
+    params = {
+        'image_url': image_url,
+        'model_choice': MODEL_TYPE
+    }
     try:
-        form = await request.form()
-        for k, v in form.multi_items():
-            # take first text-like value
-            payload["form"][k] = str(v)
-            if k in ("text", "caption", "message", "content") and not payload["text"]:
-                payload["text"] = str(v)
-            if k == "course":
-                payload["course"] = str(v)
-    except Exception:
-        pass
-
-    # If no text yet but JSON is a list or similar, stringify (best-effort)
-    if not payload["text"] and payload.get("json") is not None:
-        try:
-            payload["text"] = json.dumps(payload["json"])
-        except Exception:
-            payload["text"] = str(payload["json"])
-
-    return payload
-
-
-def add_entry(course: str, entry: Dict[str, Any]):
-    c = courses.setdefault(course, {"items": [], "last_updated": None})
-    c["items"].append(entry)
-    c["last_updated"] = time.time()
-
-
-def compile_course(course: str) -> str:
-    """Compile course data to JSON file and optionally upload to HuggingFace. Returns path to saved file."""
-    if course not in courses:
-        raise ValueError(f"Unknown course: {course}")
+        async with aiohttp.ClientSession() as session:
+            async with session.get(server, params=params, timeout=30) as resp:
+                return await resp.json()
+    except Exception as e:
+        print(f"Error from {server}: {e}")
+        return None
+
+async def get_model_info():
+    """Get model information from caption servers"""
+    model_info = []
+    async with aiohttp.ClientSession() as session:
+        for server in CAPTION_SERVERS:
+            try:
+                health_url = server.rsplit('/analyze', 1)[0] + '/health'
+                async with session.get(health_url) as resp:
+                    info = await resp.json()
+                    model_info.append({
+                        'url': server,
+                        'model': info.get('model_choice', 'unknown')
+                    })
+            except Exception as e:
+                print(f"Couldn't get model info from {server}: {e}")
+    return model_info
+
+class CaptionServer:
+    def __init__(self, url):
+        self.url = url
+        self.busy = False
+        self.model = "unknown"
+        self.total_processed = 0
+        self.total_time = 0
+
+    @property
+    def fps(self):
+        return self.total_processed / self.total_time if self.total_time > 0 else 0
+
+async def process_image(server: CaptionServer, course: str, image: Dict) -> Dict:
+    """Process single image through one caption server"""
+    if server.busy:
+        return None
+        
+    server.busy = True
+    start_time = time.time()
+    
+    try:
+        # Structure URL correctly: /images/COURSE_NAME_frames/IMAGE.png
+        course_frames = f"{course}_frames" if not course.endswith("_frames") else course
+        image_url = urljoin(SOURCE_SERVER, f"/images/{quote(course_frames)}/{quote(image['filename'])}")
+        result = await get_caption(server.url, image_url)
+        
+        processing_time = time.time() - start_time
+        server.total_time += processing_time
+        
+        if result and result.get('success') and result.get('caption'):
+            server.total_processed += 1
+            metadata = {
+                "image": image['filename'],
+                "caption": result['caption']
+            }
+            print(f"Server {server.url} processed {image['filename']} in {processing_time:.2f}s ({server.fps:.2f} fps)")
+            return metadata
+            
+    except Exception as e:
+        print(f"Error processing {image['filename']} on {server.url}: {e}")
+        
+    finally:
+        server.busy = False
+        
+    return None
 
-    data = {
-        "course": course,
-        "compiled_at": now_ts(),
-        "count": len(courses[course]["items"]),
-        "items": courses[course]["items"],
+async def submit_to_dataset(course: str, metadata_list: List[Dict]):
+    """Submit course results to dataset collection server"""
+    # Group by parent folder
+    parent_folder = os.path.dirname(course) if '/' in course else course.split('_')[0]
+    
+    # Prepare payload
+    payload = {
+        "text": f"Completed captions for course {course}. done",
+        "course": parent_folder,
+        "metadata": {
+            "course_name": course,
+            "image_count": len(metadata_list)
+        },
+        "captions": metadata_list
     }
-
-    filename = f"{course}_{now_ts()}.json"
-    safe_filename = re.sub(r"[^a-zA-Z0-9_\-\.]+", "_", filename)
-    path = os.path.join(UPLOAD_DIR, safe_filename)
-
-    with open(path, "w", encoding="utf-8") as f:
-        json.dump(data, f, ensure_ascii=False, indent=2)
-
-    # Optionally upload to HuggingFace
-    if HF_TOKEN and HF_DATASET_REPO and HF_AVAILABLE:
+    
+    async with aiohttp.ClientSession() as session:
         try:
-            api = HfApi()
-            # upload path at root of repo with same filename
-            api.upload_file(
-                path_or_fileobj=path,
-                path_in_repo=safe_filename,
-                repo_id=HF_DATASET_REPO,
-                repo_type="dataset",
-                token=HF_TOKEN,
-            )
+            async with session.post(
+                f"{DATA_COLLECTION_SERVER}/submit",
+                json=payload
+            ) as resp:
+                result = await resp.json()
+                print(f"Dataset submission result for {course}: {result}")
+                return result
         except Exception as e:
-            # Log but don't fail the compile
-            print(f"[WARN] HuggingFace upload failed: {e}")
-
-    # After compiling, clear stored items for that course
-    courses[course]["items"] = []
-    return path
-
-
-@app.post("/submit")
-async def submit(request: Request):
-    """Receive any data (text, JSON, form). Will try to extract course and URLs and store entries.
-    If message contains 'done' or similar, it will compile the course to JSON (and upload if configured).
-    """
-    payload = await parse_request(request)
-    text = (payload.get("text") or "").strip()
-
-    # Collect urls found
-    urls = extract_urls(text)
-
-    # Determine course from payload (explicit field) or from any URL
-    course = payload.get("course")
-    if not course:
-        for u in urls:
-            c = extract_course_from_url(u)
-            if c:
-                course = c
-                break
-
-    if not course:
-        course = "unknown_course"
-
-    entry = {
-        "timestamp": now_ts(),
-        "text": text,
-        "json": payload.get("json"),
-        "form": payload.get("form"),
-        "urls": urls,
-        "headers": {k: v for k, v in payload.get("headers", {}).items() if k.lower() in ("user-agent", "host", "content-type")},
-    }
-
-    add_entry(course, entry)
-
-    # Detect completion
-    if DONE_RE.search(text):
+            print(f"Error submitting to dataset: {e}")
+            return None
+
+async def process_course(course: str, servers: List[CaptionServer]):
+    """Process all images in a course using available servers"""
+    if course not in processed_images:
+        processed_images[course] = set()
+    if course not in course_captions:
+        # Load any existing captions from file
+        course_captions[course] = load_captions_from_file(course)
+        # Update processed images set from loaded captions
+        processed_images[course].update(cap['image'] for cap in course_captions[course])
+    
+    # Get list of images
+    images = await fetch_course_images(course)
+    if not images:
+        return
+    
+    print(f"\nProcessing {len(images)} images for course {course}")
+    remaining_images = [img for img in images if img['filename'] not in processed_images[course]]
+    
+    while remaining_images:
+        # Create tasks for each available server
+        tasks = []
+        for server in servers:
+            if not server.busy and remaining_images:
+                img = remaining_images[0]
+                remaining_images = remaining_images[1:]
+                tasks.append(process_image(server, course, img))
+                
+        if not tasks:
+            await asyncio.sleep(0.1)
+            continue
+            
+        # Process images in parallel across servers
+        results = await asyncio.gather(*tasks)
+        
+        # Handle results
+        has_new_results = False
+        for result in results:
+            if result:
+                processed_images[course].add(result['image'])
+                course_captions[course].append(result)
+                has_new_results = True
+        
+        # Save progress after each batch with new results
+        if has_new_results:
+            save_captions_to_file(course, course_captions[course])
+        
+        # Show progress
+        total = len(images)
+        done = len(processed_images[course])
+        print(f"\rProgress: {done}/{total} images ({done/total*100:.1f}%)", end="")
+        
+        if not remaining_images and len(processed_images[course]) == len(images):
+            print(f"\nCourse {course} complete, submitting to dataset...")
+            await submit_to_dataset(course, course_captions[course])
+            processed_images[course].clear()
+            course_captions[course].clear()
+            break
+
+async def main():
+    # Initialize caption servers
+    servers = [CaptionServer(url) for url in CAPTION_SERVERS]
+    
+    # Check for existing caption files and report
+    existing_captions = list(CAPTIONS_DIR.glob("*_captions.json"))
+    if existing_captions:
+        print("\nFound existing caption files:")
+        for cap_file in existing_captions:
+            course = cap_file.stem.replace("_captions", "")
+            try:
+                with open(cap_file, 'r', encoding='utf-8') as f:
+                    captions = json.load(f)
+                    print(f"- {course}: {len(captions)} captions")
+            except Exception as e:
+                print(f"- Error reading {cap_file.name}: {e}")
+        print()
+    
+    # Get model information and verify Florence-2-large availability
+    model_info = await get_model_info()
+    print("\nCaption Servers:")
+    available_servers = []
+    for info, server in zip(model_info, servers):
+        server.model = info['model']
+        if MODEL_TYPE in info.get('model', ''):
+            available_servers.append(server)
+            print(f"✓ {server.url} confirmed {MODEL_TYPE}")
+        else:
+            print(f"✗ {server.url} using {server.model} - skipping (requires {MODEL_TYPE})")
+    
+    if not available_servers:
+        print(f"\nError: No servers with {MODEL_TYPE} available!")
+        return
+    
+    # Update servers list to only use those with large model
+    servers = available_servers
+    print(f"\nUsing {len(servers)} servers with {MODEL_TYPE}")
+    print()
+    
+    start_time = time.time()
+    
+    while True:
         try:
-            path = compile_course(course)
-            return JSONResponse({"status": "compiled", "course": course, "path": path})
+            # Get available courses
+            courses = await fetch_courses()
+            if not courses:
+                print("No courses found, waiting...")
+                await asyncio.sleep(10)
+                continue
+            
+            print(f"Found {len(courses)} courses")
+            
+            # Process each course with all available servers
+            for course in courses:
+                await process_course(course, servers)
+                
+                # Show server stats
+                print("\nServer Stats:")
+                total_processed = sum(s.total_processed for s in servers)
+                elapsed = time.time() - start_time
+                if elapsed > 0:
+                    print(f"Total images processed: {total_processed}")
+                    print(f"Overall speed: {total_processed/elapsed:.2f} fps")
+                    for s in servers:
+                        print(f"- {s.url}: {s.total_processed} images, {s.fps:.2f} fps")
+                print()
+            
+            # Wait before next check
+            await asyncio.sleep(5)
+            
         except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    # Detect explicit 'course change' in URLs (if a URL contains a different course than stored) -- best-effort
-    # If a URL indicates a different course and there were previous items, compile previous course first
-    # Example: previous stored course is same; we don't track per-source last course, so skip this more complex behavior for now
-
-    return JSONResponse({"status": "stored", "course": course, "count": len(courses[course]["items"])})
-
-
-@app.get("/status")
-async def status():
-    summary = {c: {"count": len(v["items"]), "last_updated": v["last_updated"]} for c, v in courses.items()}
-    return {"courses": summary}
-
-
-@app.post("/compile")
-async def compile_endpoint(course: str = None):
-    """Force compile a course. If course is not provided and only one exists, compile that one."""
-    if not course:
-        if len(courses) == 1:
-            course = next(iter(courses.keys()))
-        else:
-            raise HTTPException(status_code=400, detail="Provide course query parameter when multiple courses exist.")
-
-    try:
-        path = compile_course(course)
-        return {"status": "compiled", "course": course, "path": path}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-
-@app.get("/debug/{course}")
-async def debug_course(course: str):
-    if course not in courses:
-        raise HTTPException(status_code=404, detail="Course not found")
-    return courses[course]
-
+            print(f"Error in main loop: {e}")
+            await asyncio.sleep(10)
 
 if __name__ == "__main__":
-    import uvicorn
-    port = int(os.getenv("PORT", "8000"))
-    uvicorn.run(app, host="0.0.0.0", port=port)
+    print("Starting caption coordinator...")
+    print(f"Source server: {SOURCE_SERVER}")
+    print(f"Caption servers: {CAPTION_SERVERS}")
+    print(f"Dataset server: {DATA_COLLECTION_SERVER}")
+    
+    asyncio.run(main())