Update app.py

app.py

@@ -1,100 +1,214 @@
 import os
-import cv2
 import torch
-from pathlib import Path
+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
 
 # ===== CONFIG =====
-VIDEO_PATH = "How.mp4"               # Local video file in root
-FRAMES_DIR = "extracted"            # Where frames are stored
-FPS = 3                             # Frames to extract per second
 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
 
-# ===== Ensure Output Directory =====
-def ensure_dir(path):
-    Path(path).mkdir(parents=True, exist_ok=True)
+# ===== FastAPI App =====
+app = FastAPI(
+    title="Florence-2 Image Analysis API",
+    description="Analyze images using Microsoft's Florence-2 model",
+    version="1.0.0"
+)
 
-# ===== Frame Extraction Function =====
-def extract_frames(video_path, output_dir, fps=3):
-    ensure_dir(output_dir)
-    cap = cv2.VideoCapture(str(video_path))
-    if not cap.isOpened():
-        print(f"[ERROR] Failed to open video file: {video_path}")
-        return []
+# ===== Request/Response Models =====
+class ImageAnalysisRequest(BaseModel):
+    image_url: HttpUrl
+    task: str = "<MORE_DETAILED_CAPTION>"  # Default task
 
-    video_fps = cap.get(cv2.CAP_PROP_FPS)
-    if not video_fps or video_fps <= 0:
-        print("[WARN] Using fallback FPS: 30")
-        video_fps = 30
-    frame_interval = int(round(video_fps / fps))
-
-    frame_idx = 0
-    saved_idx = 1
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    frame_paths = []
-
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-        if frame_idx % frame_interval == 0:
-            frame_name = f"{saved_idx:04d}.png"
-            output_path = Path(output_dir) / frame_name
-            cv2.imwrite(str(output_path), frame)
-            frame_paths.append(str(output_path))
-            print(f"[INFO] Saved frame {frame_idx} -> {frame_name}")
-            saved_idx += 1
-        frame_idx += 1
-    cap.release()
-    return frame_paths
+class ImageAnalysisResponse(BaseModel):
+    caption: str
+    success: bool
+    error_message: str = None
 
 # ===== Load Florence-2 Base Model =====
 print("[INFO] Loading Florence-2-base model on CPU...")
-processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    "microsoft/Florence-2-base",
-    trust_remote_code=True,
-    attn_implementation="eager"
-).to(DEVICE).eval()
+try:
+    processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        "microsoft/Florence-2-base",
+        trust_remote_code=True,
+        attn_implementation="eager"
+    ).to(DEVICE).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"""
+    try:
+        # Set headers to mimic browser request
+        headers = {
+            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
+        }
+        
+        response = requests.get(str(url), headers=headers, timeout=30)
+        response.raise_for_status()
+        
+        # Check content length
+        if len(response.content) > MAX_IMAGE_SIZE:
+            raise ValueError(f"Image too large: {len(response.content)} bytes (max: {MAX_IMAGE_SIZE})")
+        
+        # Check if content is actually an image
+        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}")
 
-# ===== Analyze a Frame =====
-def analyze_frame(image_path):
-    image = Image.open(image_path).convert("RGB")
-    image = image.resize(RESIZE_DIM, Image.BILINEAR)  # Resize for speed
-    inputs = processor(
-        text="<MORE_DETAILED_CAPTION>",
-        images=image,
-        return_tensors="pt"
-    ).to(DEVICE)
-    with torch.no_grad():
-        generated_ids = model.generate(
-            input_ids=inputs["input_ids"],
-            pixel_values=inputs["pixel_values"],
-            max_new_tokens=1024,
-            num_beams=3,
-            do_sample=False
+def analyze_image(image: Image.Image, task: str = "<MORE_DETAILED_CAPTION>") -> str:
+    """Analyze image using Florence-2 model"""
+    if not processor or not model:
+        raise ValueError("Model not loaded properly")
+    
+    try:
+        # Resize image for faster processing
+        image = image.resize(RESIZE_DIM, Image.BILINEAR)
+        
+        # Prepare inputs
+        inputs = processor(
+            text=task,
+            images=image,
+            return_tensors="pt"
+        ).to(DEVICE)
+        
+        # Generate caption
+        with torch.no_grad():
+            generated_ids = model.generate(
+                input_ids=inputs["input_ids"],
+                pixel_values=inputs["pixel_values"],
+                max_new_tokens=1024,
+                num_beams=3,
+                do_sample=False
+            )
+        
+        # Decode and post-process
+        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+        result = processor.post_process_generation(
+            generated_text,
+            task=task,
+            image_size=RESIZE_DIM
         )
-    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-    result = processor.post_process_generation(
-        generated_text,
-        task="<MORE_DETAILED_CAPTION>",
-        image_size=RESIZE_DIM
-    )
-    return result["<MORE_DETAILED_CAPTION>"]
+        
+        return result[task]
+        
+    except Exception as e:
+        raise ValueError(f"Failed to analyze image: {e}")
 
-# ===== Main Execution =====
-if __name__ == "__main__":
-    frame_list = extract_frames(VIDEO_PATH, FRAMES_DIR, FPS)
-    print(f"[INFO] Extracted {len(frame_list)} frames.")
+# ===== 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
+    }
 
-    for idx, frame_path in enumerate(frame_list):
-        print(f"\n[FRAME {idx+1}] Analyzing: {frame_path}")
-        caption = analyze_frame(frame_path)
-        print(f"[RESULT] {caption}")
+@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,
+        "available_tasks": [
+            "<MORE_DETAILED_CAPTION>",
+            "<DETAILED_CAPTION>",
+            "<CAPTION>",
+            "<OD>",  # Object Detection
+            "<DENSE_REGION_CAPTION>",
+            "<REGION_PROPOSAL>"
+        ]
+    }
 
-    # Optional: Start a dummy Uvicorn server (if you want to expand into an API later)
-    import uvicorn
-    port = int(os.getenv("PORT", 7860))  # for Gradio Spaces compatibility
-    uvicorn.run("main:app", host="0.0.0.0", port=port) if os.getenv("RUN_SERVER") else None
+@app.post("/analyze", response_model=ImageAnalysisResponse)
+async def analyze_image_endpoint(request: ImageAnalysisRequest):
+    """
+    Analyze an image from a URL using Florence-2 model
+    
+    Available tasks:
+    - <MORE_DETAILED_CAPTION>: Generate detailed image captions
+    - <DETAILED_CAPTION>: Generate detailed captions
+    - <CAPTION>: Generate basic captions
+    - <OD>: Object detection
+    - <DENSE_REGION_CAPTION>: Dense region captioning
+    - <REGION_PROPOSAL>: Region proposal
+    """
+    try:
+        # Validate task
+        valid_tasks = [
+            "<MORE_DETAILED_CAPTION>", "<DETAILED_CAPTION>", "<CAPTION>",
+            "<OD>", "<DENSE_REGION_CAPTION>", "<REGION_PROPOSAL>"
+        ]
+        if request.task not in valid_tasks:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Invalid task. Available tasks: {valid_tasks}"
+            )
+        
+        # Download and process image
+        print(f"[INFO] Processing image from: {request.image_url}")
+        image = download_image(request.image_url)
+        print(f"[INFO] Image downloaded successfully: {image.size}")
+        
+        # Analyze image
+        caption = analyze_image(image, request.task)
+        print(f"[INFO] Analysis complete: {caption}")
+        
+        return ImageAnalysisResponse(
+            caption=caption,
+            success=True
+        )
+        
+    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}")
+        raise HTTPException(status_code=500, detail=f"Internal server error: {e}")
+
+@app.get("/analyze")
+async def analyze_image_get(image_url: str, task: str = "<MORE_DETAILED_CAPTION>"):
+    """
+    GET endpoint for quick image analysis
+    Usage: /analyze?image_url=https://example.com/image.jpg&task=<MORE_DETAILED_CAPTION>
+    """
+    request = ImageAnalysisRequest(image_url=image_url, task=task)
+    return await analyze_image_endpoint(request)
+
+# ===== Main Execution =====
+if __name__ == "__main__":
+    port = int(os.getenv("PORT", 7860))
+    print(f"[INFO] Starting server on port {port}")
+    print(f"[INFO] Model status: {'Loaded' if (processor and model) else 'Failed to load'}")
+    print(f"[INFO] API Documentation: http://localhost:{port}/docs")
+    
+    uvicorn.run(
+        "main:app",
+        host="0.0.0.0",
+        port=port,
+        reload=False  # Set to True for development
+    )