Update app.py

app.py

@@ -1,360 +1,233 @@
 import os
-import sys
-import subprocess
-import importlib
-import requests
-from PIL import Image
-from io import BytesIO
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel, HttpUrl
-from transformers import AutoProcessor, AutoModelForCausalLM
-import uvicorn
-
-# ===== RUNTIME DEPENDENCY ENSURER =====
-# Hardcoded torch version to ensure compatibility at startup.
-REQUIRED_TORCH_VERSION = os.getenv("REQUIRED_TORCH_VERSION", "2.2.2")
-
-def ensure_torch_installed(required_version: str = REQUIRED_TORCH_VERSION):
-    """Ensure the required torch version is installed at runtime.
-    This will attempt to import torch and compare versions. If missing or different,
-    it will pip-install the requested version using the running Python executable.
-
-    Note: Installing torch at every start may be slow and may require build artifacts
-    specific to the platform. This helper uses a simple pip install; if your target
-    platform requires a special wheel or extra index URL, set up the environment
-    outside of this script or modify the install command accordingly.
-    """
-    try:
-        import torch as _t
-        v = getattr(_t, "__version__", "")
-        # match major.minor.patch prefix
-        if v and v.startswith(required_version):
-            print(f"[INFO] torch {v} already installed")
-            return _t
-        else:
-            print(f"[INFO] torch version {v} != {required_version}, will reinstall")
-    except Exception:
-        print("[INFO] torch not found, installing now")
+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
 
-    cmd = [sys.executable, "-m", "pip", "install", f"torch=={required_version}"]
-    print(f"[INFO] Running: {' '.join(cmd)}")
-    try:
-        subprocess.check_call(cmd)
-        importlib.invalidate_caches()
-        import torch as _t2
-        print(f"[INFO] Installed torch {_t2.__version__}")
-        return _t2
-    except subprocess.CalledProcessError as e:
-        print(f"[ERROR] pip install failed: {e}")
-        raise
-
-
-# Ensure torch is available before using the model
-torch = ensure_torch_installed()
-
-# ===== CONFIG =====
-DEVICE = "cpu"                      # Use CPU for compatibility
-RESIZE_DIM = (512, 512)             # Resize images to this resolution
-MAX_IMAGE_SIZE = 10 * 1024 * 1024   # 10MB max image size
-TASK = "<MORE_DETAILED_CAPTION>"    # Hardcoded task
-
-# URL template for frame iteration - replace with your actual URL
-BASE_URL_TEMPLATE = "https://example.com/frames/frame_{frame}.jpg"
-START_FRAME = 1          # Starting frame number
-FRAME_PADDING = 6       # Number of digits to pad frame numbers with
-
-# ===== FastAPI App =====
-app = FastAPI(
-    title="Florence-2 Image Analysis API",
-    description="Analyze images using Microsoft's Florence-2 model with detailed captions",
-    version="1.0.0"
-)
-
-# ===== Request/Response Models =====
-class ImageAnalysisRequest(BaseModel):
-    image_url: HttpUrl
-
-class ImageAnalysisResponse(BaseModel):
-    caption: str
-    success: bool
-    error_message: str = None
-
-# ===== Load Florence-2 Base Model =====
-print("[INFO] Loading Florence-2 model on CPU...")
 try:
-    MODEL_ID = "microsoft/Florence-2-base"
-    
-    # Load processor
-    processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
-    
-    # Load model
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        trust_remote_code=True,
-        torch_dtype=torch.float32,
-    )
-    
-    # Move to device manually
-    model = model.to(DEVICE)
-    model.eval()
-    
-    print("[INFO] Model loaded successfully!")
-except Exception as e:
-    print(f"[ERROR] Failed to load model: {e}")
-    processor = None
-    model = None
-
-# ===== Helper Functions =====
-def download_image(url: str) -> Image.Image:
-    """Download image from URL and return PIL Image"""
+    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:
     try:
-        headers = {
-            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
-        }
-        
-        response = requests.get(str(url), headers=headers, timeout=30)
-        response.raise_for_status()
-        
-        if len(response.content) > MAX_IMAGE_SIZE:
-            raise ValueError(f"Image too large: {len(response.content)} bytes")
-        
-        content_type = response.headers.get('content-type', '')
-        if not content_type.startswith('image/'):
-            raise ValueError(f"URL does not point to an image. Content-Type: {content_type}")
-        
-        image = Image.open(BytesIO(response.content)).convert("RGB")
-        return image
-        
-    except requests.exceptions.RequestException as e:
-        raise ValueError(f"Failed to download image: {e}")
-    except Exception as e:
-        raise ValueError(f"Failed to process image: {e}")
+        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 iterate_and_analyze(base_url_template: str, start: int = 1, padding: int = 6):
-    """Iterate over a templated frame URL and analyze images sequentially.
+def now_ts() -> str:
+    return time.strftime("%Y%m%dT%H%M%S")
 
-    base_url_template should contain a placeholder `{frame}` which will be replaced by
-    the zero-padded frame number, for example:
-      https://example.com/download?course=XYZ&file=frame%3AXYZ%2F{frame}%2Fframe_000{n}.jpg
 
-    The function yields tuples: (frame_number, url, caption_or_error)
-    Continues until a frame fails to download (e.g., 404 error)
-    """
-    if "{frame}" not in base_url_template:
-        raise ValueError("base_url_template must contain '{frame}' placeholder")
-
-    consecutive_errors = 0
-    MAX_CONSECUTIVE_ERRORS = 3  # Stop after this many consecutive errors
-    
-    i = start
-    while True:
-        frame_id = str(i).zfill(padding)
-        url = base_url_template.format(frame=frame_id)
+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)}
+
+    # Try JSON
+    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:
-            img = download_image(url)
-            caption = analyze_image(img)
-            consecutive_errors = 0  # Reset error counter on success
-            yield (i, url, {"success": True, "caption": caption})
-        except requests.exceptions.HTTPError as e:
-            if e.response.status_code == 404:
-                print(f"[INFO] No more frames found after frame {i-1}")
-                break
-            yield (i, url, {"success": False, "error": str(e)})
-            consecutive_errors += 1
-        except Exception as e:
-            yield (i, url, {"success": False, "error": str(e)})
-            consecutive_errors += 1
-        
-        if consecutive_errors >= MAX_CONSECUTIVE_ERRORS:
-            print(f"[INFO] Stopping after {MAX_CONSECUTIVE_ERRORS} consecutive errors")
-            break
-            
-        i += 1
-
-def analyze_image(image: Image.Image) -> str:
-    """Analyze image using Florence-2 model with hardcoded task"""
-    if not processor or not model:
-        raise ValueError("Model not loaded properly")
-    
+            raw = (await request.body()).decode("utf-8", errors="ignore")
+            payload["text"] = raw
+        except Exception:
+            payload["text"] = ""
+
+    # Try form (for multipart/form-data)
     try:
-        print(f"[DEBUG] Input image size: {image.size}")
-        
-        # Resize image
-        image = image.resize(RESIZE_DIM, Image.LANCZOS)
-
-        # Prepare inputs - use the same approach that worked in the test
-        inputs = processor(
-            text=TASK,
-            images=image,
-            return_tensors="pt",
-            padding=True
-        )
-        
-        print(f"[DEBUG] Input keys: {list(inputs.keys())}")
-        print(f"[DEBUG] Input IDs shape: {inputs['input_ids'].shape}")
-        print(f"[DEBUG] Pixel values shape: {inputs['pixel_values'].shape}")
-        
-        # Move to device
-        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
-        
-        # Generate caption - use the specific Florence-2 generation approach
-        print("[DEBUG] Starting generation...")
-        with torch.no_grad():
-            generated_ids = model.generate(
-                input_ids=inputs["input_ids"],
-                pixel_values=inputs["pixel_values"],
-                max_new_tokens=100,
-                num_beams=3,
-                do_sample=False,
-                early_stopping=True,
-                no_repeat_ngram_size=3,
-                length_penalty=1.0,
-            )
+        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
 
-        print("[DEBUG] Generation completed")
-        
-        # Decode and clean output
-        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        print(f"[DEBUG] Raw generated text: {repr(generated_text)}")
-        
-        # Remove the task prompt from the beginning if present
-        if generated_text.startswith(TASK):
-            generated_text = generated_text[len(TASK):].strip()
-            
-        print(f"[INFO] Final caption: {generated_text}")
-        return generated_text
+    # 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"])
 
-    except Exception as e:
-        print(f"[ERROR] Exception in analyze_image: {e}")
-        import traceback
-        print(f"[ERROR] Traceback: {traceback.format_exc()}")
-        raise ValueError(f"Failed to analyze image: {e}")
-        
-# ===== API Endpoints =====
-@app.get("/")
-async def root():
-    """Health check endpoint"""
-    return {
-        "message": "Florence-2 Image Analysis API",
-        "status": "running",
-        "model_loaded": processor is not None and model is not None,
-        "task": TASK
-    }
+    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}")
 
-@app.get("/health")
-async def health_check():
-    """Detailed health check"""
-    return {
-        "status": "healthy" if (processor and model) else "unhealthy",
-        "model_loaded": processor is not None and model is not None,
-        "device": DEVICE,
-        "task": TASK
+    data = {
+        "course": course,
+        "compiled_at": now_ts(),
+        "count": len(courses[course]["items"]),
+        "items": courses[course]["items"],
     }
 
-@app.post("/analyze", response_model=ImageAnalysisResponse)
-async def analyze_image_endpoint(request: ImageAnalysisRequest):
-    """
-    Analyze an image from a URL using Florence-2 model
-    Always uses <MORE_DETAILED_CAPTION> task for detailed image descriptions
-    """
-    try:
-        if not processor or not model:
-            raise HTTPException(
-                status_code=503,
-                detail="Model not loaded. Please check server logs."
+    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:
+        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,
             )
-        
-        print(f"[INFO] Processing image from: {request.image_url}")
-        image = download_image(request.image_url)
-        print(f"[INFO] Image downloaded successfully: {image.size}")
-        
-        caption = analyze_image(image)
-        print(f"[INFO] Analysis complete")
-        
-        return ImageAnalysisResponse(
-            caption=caption,
-            success=True
-        )
-        
-    except HTTPException:
-        raise
-    except ValueError as e:
-        print(f"[ERROR] ValueError: {e}")
-        return ImageAnalysisResponse(
-            caption="",
-            success=False,
-            error_message=str(e)
-        )
-    except Exception as e:
-        print(f"[ERROR] Unexpected error: {e}")
-        return ImageAnalysisResponse(
-            caption="",
-            success=False,
-            error_message=f"Internal server error: {str(e)}"
-        )
-
-@app.get("/analyze")
-async def analyze_image_get(image_url: str):
-    """
-    GET endpoint for quick image analysis
-    Usage: /analyze?image_url=https://example.com/image.jpg
+        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):
+        try:
+            path = compile_course(course)
+            return JSONResponse({"status": "compiled", "course": course, "path": path})
+        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:
-        request = ImageAnalysisRequest(image_url=image_url)
-        return await analyze_image_endpoint(request)
+        path = compile_course(course)
+        return {"status": "compiled", "course": course, "path": path}
     except Exception as e:
-        raise HTTPException(status_code=400, detail=str(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]
+
 
-# ===== Main Execution =====
 if __name__ == "__main__":
-    if not processor or not model:
-        print("[ERROR] Model failed to load. Cannot proceed with frame analysis.")
-        sys.exit(1)
-
-    print("[INFO] Starting frame analysis...")
-    print(f"[INFO] Using URL template: {BASE_URL_TEMPLATE}")
-    print(f"[INFO] Starting from frame {START_FRAME} with {FRAME_PADDING} digit padding")
-
-    results = []
-    for frame_num, url, result in iterate_and_analyze(
-        BASE_URL_TEMPLATE,
-        start=START_FRAME,
-        padding=FRAME_PADDING
-    ):
-        if result["success"]:
-            print(f"[SUCCESS] Frame {frame_num}: {result['caption']}")
-            results.append({
-                "frame": frame_num,
-                "url": url,
-                "caption": result["caption"]
-            })
-        else:
-            print(f"[ERROR] Frame {frame_num}: {result['error']}")
-            results.append({
-                "frame": frame_num,
-                "url": url,
-                "error": result["error"]
-            })
-
-    # Save results to a JSON file
-    import json
-    output_file = "frame_analysis_results.json"
-    with open(output_file, "w", encoding="utf-8") as f:
-        json.dump(results, f, indent=2, ensure_ascii=False)
-    print(f"[INFO] Results saved to {output_file}")
-
-    # Optional: start the API server after frame analysis
-    start_server = os.getenv("START_SERVER", "false").lower() == "true"
-    if start_server:
-        port = int(os.getenv("PORT", 7860))
-        print(f"[INFO] Starting server on port {port}")
-        print(f"[INFO] Task: {TASK}")
-        print(f"[INFO] API Documentation: http://localhost:{port}/docs")
-        
-        uvicorn.run(
-            app,
-            host="0.0.0.0",
-            port=port,
-            reload=False
-        )
+    import uvicorn
+    port = int(os.getenv("PORT", "8000"))
+    uvicorn.run(app, host="0.0.0.0", port=port)