Update app.py

app.py

@@ -26,7 +26,13 @@ def resize_image(image: Image.Image, max_size: int = 1024) -> Image.Image:
     return image
 
 @app.post("/predict")
-async def predict(files: List[UploadFile] = File(...)):
+async def predict(
+    files: List[UploadFile] = File(...),
+    description: bool = True,
+    signs: bool = True,
+    harmful: bool = True,
+    similarity: bool = True
+):
     if len(files) > 3:
         raise HTTPException(status_code=400, detail="Maximum 3 images allowed.")
     
@@ -45,55 +51,59 @@ async def predict(files: List[UploadFile] = File(...)):
             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, people, and environmental context."
-            inputs = processor(text=[prompt_desc], images=[image], return_tensors="pt", padding=True)
-            outputs = model.generate(**inputs, max_new_tokens=256)
-            description = processor.decode(outputs[0], skip_special_tokens=True).replace(prompt_desc, "").strip()
+            result = {"image_id": f"Image_{idx}", "timestamp": timestamp}
             
-            # Extract signs/number plates (OCR)
-            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_new_tokens=256)
-            signs_text = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace(prompt_ocr, "").strip()
+            # Generate description if selected
+            if description:
+                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_new_tokens=256)
+                result["description"] = processor.decode(outputs[0], skip_special_tokens=True).replace(prompt_desc, "").strip() or "No description generated."
             
-            # 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_new_tokens=256)
-            detected_objects = processor.decode(detect_outputs[0], skip_special_tokens=True).lower()
-            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}]
+            # Extract signs/number plates (OCR) if selected
+            if signs:
+                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_new_tokens=256)
+                result["signs"] = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace(prompt_ocr, "").strip() or "None detected"
             
-            # Get image embedding for similarity
-            inputs_emb = processor(images=[image], return_tensors="pt", padding=True)
-            with torch.no_grad():
-                emb = model.vision_model(inputs_emb["pixel_values"]).last_hidden_state.mean(dim=1).cpu().numpy()
-            image_embeddings.append(emb)
+            # 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."
+                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()
+                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})
+                result["harmful_objects"] = harmful_detected if harmful_detected else [{"object": "None", "confidence": 0}]
             
-            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_objects": harmful_output
-            })
+            # Get image embedding for similarity if selected
+            if similarity:
+                inputs_emb = processor(images=[image], return_tensors="pt", padding=True)
+                with torch.no_grad():
+                    emb = model.vision_model(inputs_emb["pixel_values"]).last_hidden_state.mean(dim=1).cpu().numpy()
+                image_embeddings.append(emb)
+            
+            results.append(result)
         except Exception as e:
-            results.append({
+            error_result = {
                 "image_id": f"Image_{idx}",
                 "timestamp": timestamp,
-                "description": "Error processing image.",
-                "signs": "Error",
-                "harmful_objects": [{"object": "Error", "confidence": 0}],
                 "error": str(e)
-            })
+            }
+            if description:
+                error_result["description"] = "Error processing image."
+            if signs:
+                error_result["signs"] = "Error"
+            if harmful:
+                error_result["harmful_objects"] = [{"object": "Error", "confidence": 0}]
+            results.append(error_result)
     
-    # Compute similarity scores (cosine similarity to first image)
-    if len(image_embeddings) > 1:
+    # Compute similarity scores if selected
+    if similarity and len(image_embeddings) > 1:
         base_embedding = image_embeddings[0]
         for i in range(1, len(image_embeddings)):
             sim = np.dot(base_embedding, image_embeddings[i].T) / (
@@ -104,71 +114,90 @@ async def predict(files: List[UploadFile] = File(...)):
     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:
+def gradio_predict(
+    image_1, image_2, image_3,
+    description: bool,
+    signs: bool,
+    harmful: bool,
+    similarity: bool,
+    combined: bool
+):
+    images = [image_1, image_2, image_3]
+    images = [img for img in images if img is not None]
+    if not images:
+        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]):
+        return "Error: At least one output option must be selected."
+    
+    # If combined is selected, enable all outputs
+    if combined:
+        description = signs = harmful = similarity = True
     
     results = []
     image_embeddings = []
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     
     for idx, image in enumerate(images, 1):
-        if image is None:
-            continue
         try:
             # 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, people, and environmental context."
-            inputs = processor(text=[prompt_desc], images=[image], return_tensors="pt", padding=True)
-            outputs = model.generate(**inputs, max_new_tokens=256)
-            description = processor.decode(outputs[0], skip_special_tokens=True).replace(prompt_desc, "").strip()
+            result = {"image_id": f"Image_{idx}", "timestamp": timestamp}
             
-            # Extract signs/number plates (OCR)
-            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_new_tokens=256)
-            signs_text = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace(prompt_ocr, "").strip()
+            # Generate description if selected
+            if description:
+                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_new_tokens=256)
+                result["description"] = processor.decode(outputs[0], skip_special_tokens=True).replace(prompt_desc, "").strip() or "No description generated."
             
-            # 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_new_tokens=256)
-            detected_objects = processor.decode(detect_outputs[0], skip_special_tokens=True).lower()
-            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}]
+            # Extract signs/number plates (OCR) if selected
+            if signs:
+                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_new_tokens=256)
+                result["signs"] = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace(prompt_ocr, "").strip() or "None detected"
             
-            # Get image embedding for similarity
-            inputs_emb = processor(images=[image], return_tensors="pt", padding=True)
-            with torch.no_grad():
-                emb = model.vision_model(inputs_emb["pixel_values"]).last_hidden_state.mean(dim=1).cpu().numpy()
-            image_embeddings.append(emb)
+            # 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."
+                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()
+                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})
+                result["harmful_objects"] = harmful_detected if harmful_detected else [{"object": "None", "confidence": 0}]
             
-            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_objects": harmful_output
-            })
+            # Get image embedding for similarity if selected
+            if similarity:
+                inputs_emb = processor(images=[image], return_tensors="pt", padding=True)
+                with torch.no_grad():
+                    emb = model.vision_model(inputs_emb["pixel_values"]).last_hidden_state.mean(dim=1).cpu().numpy()
+                image_embeddings.append(emb)
+            
+            results.append(result)
         except Exception as e:
-            results.append({
+            error_result = {
                 "image_id": f"Image_{idx}",
                 "timestamp": timestamp,
-                "description": "Error processing image.",
-                "signs": "Error",
-                "harmful_objects": [{"object": "Error", "confidence": 0}],
                 "error": str(e)
-            })
+            }
+            if description:
+                error_result["description"] = "Error processing image."
+            if signs:
+                error_result["signs"] = "Error"
+            if harmful:
+                error_result["harmful_objects"] = [{"object": "Error", "confidence": 0}]
+            results.append(error_result)
     
-    # Compute similarity scores
-    if len(image_embeddings) > 1:
+    # Compute similarity scores if selected
+    if similarity and len(image_embeddings) > 1:
         base_embedding = image_embeddings[0]
         for i in range(1, len(image_embeddings)):
             sim = np.dot(base_embedding, image_embeddings[i].T) / (
@@ -178,25 +207,56 @@ def gradio_predict(*images):
     
     # 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"
+    headers = ["Image ID"]
+    if description:
+        headers.append("Description")
+    if signs:
+        headers.append("Signs/Number Plates")
+    if harmful:
+        headers.append("Harmful Objects")
+    if similarity:
+        headers.append("Similarity to Image 1")
+    
+    output += "| " + " | ".join(headers) + " |\n"
+    output += "| " + " | ".join(["---"] * len(headers)) + " |\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"
+        row = [f"{result['image_id']} ({result['timestamp']})"]
+        if description:
+            row.append(result.get("description", "N/A"))
+        if signs:
+            row.append(result.get("signs", "N/A"))
+        if harmful:
+            harmful_str = ", ".join([f"{obj['object']} ({obj['confidence']}%)" for obj in result.get("harmful_objects", [{"object": "N/A", "confidence": 0}])])
+            row.append(harmful_str)
+        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)
+        
+        output += "| " + " | ".join(row) + " |\n"
         if "error" in result:
-            output += f"| **Error** | {result['error']} | - | - | - |\n"
+            output += f"| **Error** | {result['error']}" + " | " * (len(headers) - 1) + " |\n"
+    
     return output
 
 # Gradio interface
 iface = gr.Interface(
     fn=gradio_predict,
-    inputs=[gr.Image(label=f"Upload Image {i+1} (up to 10MB)") for i in range(3)],  # Allow up to 3 images
+    inputs=[
+        gr.Image(label="Upload Image 1 (up to 10MB)"),
+        gr.Image(label="Upload Image 2 (up to 10MB)"),
+        gr.Image(label="Upload Image 3 (up to 10MB)"),
+        gr.Checkbox(label="Description", value=True),
+        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)
+    ],
     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."
+    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 formatted for investigative analysis."
 )
 
 if __name__ == "__main__":
-   
+    # Launch Gradio interface for Hugging Face Spaces
     iface.launch(server_name="0.0.0.0", server_port=7860)