Update app.py

app.py

@@ -3,12 +3,13 @@ import torch
 from transformers import AutoModelForVision2Seq, AutoProcessor
 from PIL import Image
 import numpy as np
-from fastapi import FastAPI, UploadFile, File
+from fastapi import FastAPI, UploadFile, File, HTTPException
 from typing import List
 import io
 import gradio as gr
+from datetime import datetime
 
-# Initialize FastAPI app
+# Initialize FastAPI app with increased upload limit (10MB)
 app = FastAPI()
 
 # Load SmolVLM-Instruct model and processor
@@ -19,35 +20,54 @@ model = AutoModelForVision2Seq.from_pretrained(model_id, token=os.environ.get("H
 # Harmful objects list for detection
 harmful_objects = ["knife", "gun", "weapon", "blood", "syringe", "bomb", "blade"]
 
+# Resize image to max 1024x1024 while preserving aspect ratio
+def resize_image(image: Image.Image, max_size: int = 1024) -> Image.Image:
+    image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
+    return image
+
 @app.post("/predict")
 async def predict(files: List[UploadFile] = File(...)):
+    if len(files) > 3:
+        raise HTTPException(status_code=400, detail="Maximum 3 images allowed.")
+    
     results = []
     image_embeddings = []
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     
-    for file in files:
+    for idx, file in enumerate(files, 1):
         try:
-            # Read image
+            # Check file size (limit to 10MB)
             image_data = await file.read()
+            if len(image_data) > 10 * 1024 * 1024:
+                raise ValueError("Image file size exceeds 10MB limit.")
+            
+            # Read and resize image
             image = Image.open(io.BytesIO(image_data)).convert("RGB")
+            image = resize_image(image, max_size=1024)
             
             # Generate description
-            prompt_desc = "<image> Provide a detailed description of the image, including objects, colors, and context."
+            prompt_desc = "<image> Provide a detailed description of the image, including objects, colors, people, and environmental context."
             inputs = processor(text=[prompt_desc], images=[image], return_tensors="pt", padding=True)
             outputs = model.generate(**inputs, max_length=512)
             description = processor.decode(outputs[0], skip_special_tokens=True).replace(prompt_desc, "").strip()
             
             # Extract signs/number plates (OCR)
-            prompt_ocr = "<image> Extract all visible text in the image, such as signs or license plates."
+            prompt_ocr = "<image> Extract all visible text in the image, such as road signs, license plates, or billboards, with exact wording."
             inputs_ocr = processor(text=[prompt_ocr], images=[image], return_tensors="pt", padding=True)
             ocr_outputs = model.generate(**inputs_ocr, max_length=512)
             signs_text = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace(prompt_ocr, "").strip()
             
-            # Detect harmful objects/blood
-            prompt_detect = "<image> Identify any harmful objects (e.g., knife, gun, blood, syringe, bomb, blade) in the image. List them explicitly."
+            # Detect harmful objects/blood with confidence simulation
+            prompt_detect = "<image> Identify any harmful objects (e.g., knife, gun, blood, syringe, bomb, blade) in the image. List them explicitly and estimate confidence (0-100%) for each detection based on clarity."
             inputs_detect = processor(text=[prompt_detect], images=[image], return_tensors="pt", padding=True)
             detect_outputs = model.generate(**inputs_detect, max_length=512)
             detected_objects = processor.decode(detect_outputs[0], skip_special_tokens=True).lower()
-            harmful_detected = any(obj in detected_objects for obj in harmful_objects) and "Detected: " + ", ".join([obj for obj in harmful_objects if obj in detected_objects]) or "None detected"
+            harmful_detected = []
+            for obj in harmful_objects:
+                if obj in detected_objects:
+                    confidence = 90 if obj in detected_objects.split() else 60
+                    harmful_detected.append({"object": obj, "confidence": confidence})
+            harmful_output = harmful_detected if harmful_detected else [{"object": "None", "confidence": 0}]
             
             # Get image embedding for similarity
             inputs_emb = processor(images=[image], return_tensors="pt", padding=True)
@@ -56,15 +76,19 @@ async def predict(files: List[UploadFile] = File(...)):
             image_embeddings.append(emb)
             
             results.append({
+                "image_id": f"Image_{idx}",
+                "timestamp": timestamp,
                 "description": description if description else "No description generated.",
                 "signs": signs_text if signs_text else "None detected",
-                "harmful": harmful_detected
+                "harmful_objects": harmful_output
             })
         except Exception as e:
             results.append({
+                "image_id": f"Image_{idx}",
+                "timestamp": timestamp,
                 "description": "Error processing image.",
                 "signs": "Error",
-                "harmful": "Error",
+                "harmful_objects": [{"object": "Error", "confidence": 0}],
                 "error": str(e)
             })
     
@@ -75,40 +99,50 @@ async def predict(files: List[UploadFile] = File(...)):
             sim = np.dot(base_embedding, image_embeddings[i].T) / (
                 np.linalg.norm(base_embedding) * np.linalg.norm(image_embeddings[i])
             )
-            results[i]["similarity"] = float(sim[0][0])
+            results[i]["similarity_to_image_1"] = float(sim[0][0])
     
-    return results
+    return {"results": results, "analysis_timestamp": timestamp}
 
 # Gradio interface for Hugging Face Spaces
 def gradio_predict(*images):
+    if len([img for img in images if img is not None]) > 3:
+        return "Error: Maximum 3 images allowed."
+    
     results = []
     image_embeddings = []
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     
-    for image in images:
+    for idx, image in enumerate(images, 1):
         if image is None:
             continue
         try:
-            # Convert Gradio image input to PIL
+            # Convert Gradio image input to PIL and resize
             image = Image.fromarray(image).convert("RGB")
+            image = resize_image(image, max_size=1024)
             
             # Generate description
-            prompt_desc = "<image> Provide a detailed description of the image, including objects, colors, and context."
+            prompt_desc = "<image> Provide a detailed description of the image, including objects, colors, people, and environmental context."
             inputs = processor(text=[prompt_desc], images=[image], return_tensors="pt", padding=True)
             outputs = model.generate(**inputs, max_length=512)
             description = processor.decode(outputs[0], skip_special_tokens=True).replace(prompt_desc, "").strip()
             
             # Extract signs/number plates (OCR)
-            prompt_ocr = "<image> Extract all visible text in the image, such as signs or license plates."
-            inputs_ocr = processor(text=[prompt_ocr], images=[image], return_tensors="pt", padding=True)
+            prompt_ocr = "<image> Extract all visible text in the image, such as road signs, license plates, or billboards, with exact wording."
+            inputs_ocr = processor[text=[prompt_ocr], images=[image], return_tensors="pt", padding=True)
             ocr_outputs = model.generate(**inputs_ocr, max_length=512)
             signs_text = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace(prompt_ocr, "").strip()
             
-            # Detect harmful objects/blood
-            prompt_detect = "<image> Identify any harmful objects (e.g., knife, gun, blood, syringe, bomb, blade) in the image. List them explicitly."
+            # Detect harmful objects/blood with confidence simulation
+            prompt_detect = "<image> Identify any harmful objects (e.g., knife, gun, blood, syringe, bomb, blade) in the image. List them explicitly and estimate confidence (0-100%) for each detection based on clarity."
             inputs_detect = processor(text=[prompt_detect], images=[image], return_tensors="pt", padding=True)
             detect_outputs = model.generate(**inputs_detect, max_length=512)
             detected_objects = processor.decode(detect_outputs[0], skip_special_tokens=True).lower()
-            harmful_detected = any(obj in detected_objects for obj in harmful_objects) and "Detected: " + ", ".join([obj for obj in harmful_objects if obj in detected_objects]) or "None detected"
+            harmful_detected = []
+            for obj in harmful_objects:
+                if obj in detected_objects:
+                    confidence = 90 if obj in detected_objects.split() else 60
+                    harmful_detected.append({"object": obj, "confidence": confidence})
+            harmful_output = harmful_detected if harmful_detected else [{"object": "None", "confidence": 0}]
             
             # Get image embedding for similarity
             inputs_emb = processor(images=[image], return_tensors="pt", padding=True)
@@ -117,15 +151,19 @@ def gradio_predict(*images):
             image_embeddings.append(emb)
             
             results.append({
+                "image_id": f"Image_{idx}",
+                "timestamp": timestamp,
                 "description": description if description else "No description generated.",
                 "signs": signs_text if signs_text else "None detected",
-                "harmful": harmful_detected
+                "harmful_objects": harmful_output
             })
         except Exception as e:
             results.append({
+                "image_id": f"Image_{idx}",
+                "timestamp": timestamp,
                 "description": "Error processing image.",
                 "signs": "Error",
-                "harmful": "Error",
+                "harmful_objects": [{"object": "Error", "confidence": 0}],
                 "error": str(e)
             })
     
@@ -136,30 +174,27 @@ def gradio_predict(*images):
             sim = np.dot(base_embedding, image_embeddings[i].T) / (
                 np.linalg.norm(base_embedding) * np.linalg.norm(image_embeddings[i])
             )
-            results[i]["similarity"] = float(sim[0][0])
+            results[i]["similarity_to_image_1"] = float(sim[0][0])
     
-    # Format output for Gradio
-    output = ""
-    for i, result in enumerate(results):
-        output += f"**Image {i+1}**\n"
-        output += f"- Description: {result['description']}\n"
-        output += f"- Signs/Number Plates: {result['signs']}\n"
-        output += f"- Harmful Objects/Blood: {result['harmful']}\n"
-        if "similarity" in result:
-            output += f"- Similarity to Image 1: {result['similarity']:.2f}\n"
+    # Format output for Gradio as tabulated markdown
+    output = f"**Analysis Timestamp**: {timestamp}\n\n"
+    output += "| Image ID | Description | Signs/Number Plates | Harmful Objects | Similarity to Image 1 |\n"
+    output += "|----------|-------------|---------------------|-----------------|-----------------------|\n"
+    for result in results:
+        harmful_str = ", ".join([f"{obj['object']} ({obj['confidence']}%)" for obj in result['harmful_objects']])
+        similarity = f"{result['similarity_to_image_1']:.2f}" if 'similarity_to_image_1' in result else "N/A"
+        output += f"| {result['image_id']} ({result['timestamp']}) | {result['description']} | {result['signs']} | {harmful_str} | {similarity} |\n"
         if "error" in result:
-            output += f"- Error: {result['error']}\n"
-        output += "\n"
-    
+            output += f"| **Error** | {result['error']} | - | - | - |\n"
     return output
 
 # Gradio interface
 iface = gr.Interface(
     fn=gradio_predict,
-    inputs=[gr.Image(label=f"Upload Image {i+1}") for i in range(3)],  # Allow up to 3 images
-    outputs=gr.Textbox(label="Results"),
-    title="VisionSage: Image Analysis with SmolVLM",
-    description="Upload up to 3 images to get detailed descriptions, extract signs/number plates, detect harmful objects/blood, and compute similarity to the first image."
+    inputs=[gr.Image(label=f"Upload Image {i+1} (up to 10MB)") for i in range(3)],  # Allow up to 3 images
+    outputs=gr.Markdown(label="Investigation Results"),
+    title="VisionSage: Image Analysis for Investigation",
+    description="Upload up to 3 images (up to 10MB each, any resolution) to get detailed descriptions, extract signs/number plates, detect harmful objects/blood with confidence scores, and compute similarity to the first image. Results are formatted for investigative analysis."
 )
 
 if __name__ == "__main__":