Update app.py

app.py

@@ -8,6 +8,7 @@ from typing import List
 import io
 import gradio as gr
 from datetime import datetime
+import json
 
 # Initialize FastAPI app with increased upload limit (10MB)
 app = FastAPI()
@@ -17,25 +18,59 @@ model_id = "HuggingFaceTB/SmolVLM-Instruct"
 processor = AutoProcessor.from_pretrained(model_id, token=os.environ.get("HF_TOKEN"), padding=True)
 model = AutoModelForVision2Seq.from_pretrained(model_id, token=os.environ.get("HF_TOKEN"))
 
-# Harmful objects list for detection
-harmful_objects = ["knife", "gun", "weapon", "blood", "syringe", "bomb", "blade"]
+# Default harmful objects list
+default_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
 
+# Format results as a detailed report
+def format_detailed_report(results, timestamp):
+    report = f"Investigation Report\nGenerated on: {timestamp}\n\n"
+    for result in results:
+        report += f"Image ID: {result['image_id']} (Timestamp: {result['timestamp']})\n"
+        if "description" in result:
+            report += f"Description: {result['description']}\n"
+        if "signs" in result:
+            report += f"Signs/Number Plates: {result['signs']}\n"
+        if "harmful_objects" in result:
+            report += f"Harmful Objects: {', '.join([f'{obj['object']} ({obj['confidence']}%)' for obj in result['harmful_objects']])}\n"
+        if "similarity_to_image_1" in result:
+            report += f"Similarity to Image 1: {result['similarity_to_image_1']:.2f}\n"
+        if "faces" in result:
+            report += f"Facial Attributes: {result['faces']}\n"
+        if "objects" in result:
+            report += f"Localized Objects: {', '.join([f'{obj['object']} at {obj['bbox']}' for obj in result['objects']])}\n"
+        if "scene_context" in result:
+            report += f"Scene Context: {result['scene_context']}\n"
+        if "error" in result:
+            report += f"Error: {result['error']}\n"
+        report += "-" * 50 + "\n"
+    return report
+
 @app.post("/predict")
 async def predict(
     files: List[UploadFile] = File(...),
     description: bool = True,
     signs: bool = True,
     harmful: bool = True,
-    similarity: bool = True
+    similarity: bool = True,
+    faces: bool = False,
+    objects: bool = False,
+    scene: bool = False,
+    custom_harmful: str = ""
 ):
     if len(files) > 3:
         raise HTTPException(status_code=400, detail="Maximum 3 images allowed.")
     
+    # Process custom harmful objects
+    harmful_objects = default_harmful_objects.copy()
+    if custom_harmful.strip():
+        custom_objects = [obj.strip().lower() for obj in custom_harmful.split(",")]
+        harmful_objects.extend(custom_objects)
+    
     results = []
     image_embeddings = []
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
@@ -69,7 +104,7 @@ async def predict(
             
             # Detect harmful objects/blood with confidence simulation if selected
             if harmful:
-                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."
+                prompt_detect = f"<image> Identify any harmful objects ({', '.join(harmful_objects)}) 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_new_tokens=256)
                 detected_objects = processor.decode(detect_outputs[0], skip_special_tokens=True).lower()
@@ -80,6 +115,27 @@ async def predict(
                         harmful_detected.append({"object": obj, "confidence": confidence})
                 result["harmful_objects"] = harmful_detected if harmful_detected else [{"object": "None", "confidence": 0}]
             
+            # Detect facial attributes if selected
+            if faces:
+                prompt_faces = "<image> Detect faces and estimate attributes such as age range (e.g., child, adult, senior), gender (male, female, unknown), and emotional cues (e.g., neutral, angry, scared)."
+                inputs_faces = processor(text=[prompt_faces], images=[image], return_tensors="pt", padding=True)
+                faces_outputs = model.generate(**inputs_faces, max_new_tokens=256)
+                result["faces"] = processor.decode(faces_outputs[0], skip_special_tokens=True).replace(prompt_faces, "").strip() or "No faces detected"
+            
+            # Localize objects if selected
+            if objects:
+                prompt_objects = "<image> Identify and localize key objects (e.g., vehicles, weapons, bags) in the image. Provide object names and approximate bounding box coordinates (x_min, y_min, x_max, y_max) in the image."
+                inputs_objects = processor(text=[prompt_objects], images=[image], return_tensors="pt", padding=True)
+                objects_outputs = model.generate(**inputs_objects, max_new_tokens=256)
+                result["objects"] = [{"object": obj.strip(), "bbox": "(unknown)"} for obj in processor.decode(objects_outputs[0], skip_special_tokens=True).replace(prompt_objects, "").split(",") if obj.strip()] or [{"object": "None", "bbox": "N/A"}]
+            
+            # Analyze scene context if selected
+            if scene:
+                prompt_scene = "<image> Classify the scene type (e.g., indoor, outdoor, urban, rural) and estimate the time of day (e.g., day, night, dusk)."
+                inputs_scene = processor(text=[prompt_scene], images=[image], return_tensors="pt", padding=True)
+                scene_outputs = model.generate(**inputs_scene, max_new_tokens=256)
+                result["scene_context"] = processor.decode(scene_outputs[0], skip_special_tokens=True).replace(prompt_scene, "").strip() or "No context determined"
+            
             # Get image embedding for similarity if selected
             if similarity:
                 inputs_emb = processor(images=[image], return_tensors="pt", padding=True)
@@ -100,6 +156,12 @@ async def predict(
                 error_result["signs"] = "Error"
             if harmful:
                 error_result["harmful_objects"] = [{"object": "Error", "confidence": 0}]
+            if faces:
+                error_result["faces"] = "Error"
+            if objects:
+                error_result["objects"] = [{"object": "Error", "bbox": "N/A"}]
+            if scene:
+                error_result["scene_context"] = "Error"
             results.append(error_result)
     
     # Compute similarity scores if selected
@@ -120,7 +182,13 @@ def gradio_predict(
     signs: bool,
     harmful: bool,
     similarity: bool,
-    combined: bool
+    faces: bool,
+    objects: bool,
+    scene: bool,
+    combined: bool,
+    json_export: bool,
+    detailed_report: bool,
+    custom_harmful: str
 ):
     images = [image_1, image_2, image_3]
     images = [img for img in images if img is not None]
@@ -128,12 +196,18 @@ def gradio_predict(
         return "Error: At least one image must be uploaded."
     if len(images) > 3:
         return "Error: Maximum 3 images allowed."
-    if not any([description, signs, harmful, similarity, combined]):
+    if not any([description, signs, harmful, similarity, faces, objects, scene, combined]):
         return "Error: At least one output option must be selected."
     
     # If combined is selected, enable all outputs
     if combined:
-        description = signs = harmful = similarity = True
+        description = signs = harmful = similarity = faces = objects = scene = True
+    
+    # Process custom harmful objects
+    harmful_objects = default_harmful_objects.copy()
+    if custom_harmful.strip():
+        custom_objects = [obj.strip().lower() for obj in custom_harmful.split(",")]
+        harmful_objects.extend(custom_objects)
     
     results = []
     image_embeddings = []
@@ -163,7 +237,7 @@ def gradio_predict(
             
             # Detect harmful objects/blood with confidence simulation if selected
             if harmful:
-                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."
+                prompt_detect = f"<image> Identify any harmful objects ({', '.join(harmful_objects)}) 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_new_tokens=256)
                 detected_objects = processor.decode(detect_outputs[0], skip_special_tokens=True).lower()
@@ -174,6 +248,27 @@ def gradio_predict(
                         harmful_detected.append({"object": obj, "confidence": confidence})
                 result["harmful_objects"] = harmful_detected if harmful_detected else [{"object": "None", "confidence": 0}]
             
+            # Detect facial attributes if selected
+            if faces:
+                prompt_faces = "<image> Detect faces and estimate attributes such as age range (e.g., child, adult, senior), gender (male, female, unknown), and emotional cues (e.g., neutral, angry, scared)."
+                inputs_faces = processor(text=[prompt_faces], images=[image], return_tensors="pt", padding=True)
+                faces_outputs = model.generate(**inputs_faces, max_new_tokens=256)
+                result["faces"] = processor.decode(faces_outputs[0], skip_special_tokens=True).replace(prompt_faces, "").strip() or "No faces detected"
+            
+            # Localize objects if selected
+            if objects:
+                prompt_objects = "<image> Identify and localize key objects (e.g., vehicles, weapons, bags) in the image. Provide object names and approximate bounding box coordinates (x_min, y_min, x_max, y_max) in the image."
+                inputs_objects = processor(text=[prompt_objects], images=[image], return_tensors="pt", padding=True)
+                objects_outputs = model.generate(**inputs_objects, max_new_tokens=256)
+                result["objects"] = [{"object": obj.strip(), "bbox": "(unknown)"} for obj in processor.decode(objects_outputs[0], skip_special_tokens=True).replace(prompt_objects, "").split(",") if obj.strip()] or [{"object": "None", "bbox": "N/A"}]
+            
+            # Analyze scene context if selected
+            if scene:
+                prompt_scene = "<image> Classify the scene type (e.g., indoor, outdoor, urban, rural) and estimate the time of day (e.g., day, night, dusk)."
+                inputs_scene = processor(text=[prompt_scene], images=[image], return_tensors="pt", padding=True)
+                scene_outputs = model.generate(**inputs_scene, max_new_tokens=256)
+                result["scene_context"] = processor.decode(scene_outputs[0], skip_special_tokens=True).replace(prompt_scene, "").strip() or "No context determined"
+            
             # Get image embedding for similarity if selected
             if similarity:
                 inputs_emb = processor(images=[image], return_tensors="pt", padding=True)
@@ -194,6 +289,12 @@ def gradio_predict(
                 error_result["signs"] = "Error"
             if harmful:
                 error_result["harmful_objects"] = [{"object": "Error", "confidence": 0}]
+            if faces:
+                error_result["faces"] = "Error"
+            if objects:
+                error_result["objects"] = [{"object": "Error", "bbox": "N/A"}]
+            if scene:
+                error_result["scene_context"] = "Error"
             results.append(error_result)
     
     # Compute similarity scores if selected
@@ -216,6 +317,12 @@ def gradio_predict(
         headers.append("Harmful Objects")
     if similarity:
         headers.append("Similarity to Image 1")
+    if faces:
+        headers.append("Facial Attributes")
+    if objects:
+        headers.append("Localized Objects")
+    if scene:
+        headers.append("Scene Context")
     
     output += "| " + " | ".join(headers) + " |\n"
     output += "| " + " | ".join(["---"] * len(headers)) + " |\n"
@@ -232,11 +339,27 @@ def gradio_predict(
         if similarity:
             similarity_val = f"{result.get('similarity_to_image_1', 0):.2f}" if 'similarity_to_image_1' in result else "N/A"
             row.append(similarity_val)
+        if faces:
+            row.append(result.get("faces", "N/A"))
+        if objects:
+            objects_str = ", ".join([f"{obj['object']} at {obj['bbox']}" for obj in result.get("objects", [{"object": "N/A", "bbox": "N/A"}])])
+            row.append(objects_str)
+        if scene:
+            row.append(result.get("scene_context", "N/A"))
         
         output += "| " + " | ".join(row) + " |\n"
         if "error" in result:
             output += f"| **Error** | {result['error']}" + " | " * (len(headers) - 1) + " |\n"
     
+    # Handle JSON export
+    if json_export:
+        json_output = {"results": results, "analysis_timestamp": timestamp}
+        output += "\n**JSON Export**:\n```json\n" + json.dumps(json_output, indent=2) + "\n```"
+    
+    # Handle detailed report
+    if detailed_report:
+        output += "\n**Detailed Report**:\n" + format_detailed_report(results, timestamp)
+    
     return output
 
 # Gradio interface
@@ -250,11 +373,17 @@ iface = gr.Interface(
         gr.Checkbox(label="Signs/Number Plates", value=True),
         gr.Checkbox(label="Harmful Objects/Blood", value=True),
         gr.Checkbox(label="Similarity to Image 1", value=True),
-        gr.Checkbox(label="Combined (All Outputs)", value=False)
+        gr.Checkbox(label="Facial Attributes", value=False),
+        gr.Checkbox(label="Localized Objects", value=False),
+        gr.Checkbox(label="Scene Context", value=False),
+        gr.Checkbox(label="Combined (All Outputs)", value=False),
+        gr.Checkbox(label="JSON Export", value=False),
+        gr.Checkbox(label="Detailed Report", value=False),
+        gr.Textbox(label="Custom Harmful Objects (comma-separated, e.g., handgun,crowbar)", placeholder="Enter custom objects to detect")
     ],
     outputs=gr.Textbox(label="Investigation Results", placeholder="Results will appear here..."),
-    title="VisionSage: Image Analysis for Investigation",
-    description="Upload up to 3 images (up to 10MB each, any resolution) and select desired outputs: Description, Signs/Number Plates, Harmful Objects/Blood, Similarity, or Combined. Results are displayed in a structured table for investigative analysis."
+    title="VisionSage: Image Analysis for Crime Investigation",
+    description="Upload up to 3 images (up to 10MB each, any resolution) and select desired outputs for investigative analysis. Features include descriptions, text extraction, harmful object detection, similarity analysis, facial attributes, object localization, and scene context. Export results as JSON or a detailed report."
 )
 
 if __name__ == "__main__":