Update app.py

app.py

@@ -65,7 +65,7 @@ def validate_custom_harmful(custom_harmful: str) -> List[str]:
 def calculate_threat_score(harmful_objects: List[dict], custom_weights: dict) -> float:
     score = 0.0
     for obj in harmful_objects:
-        if obj["object"] != "None" and obj["object"] != "Error":
+        if obj["object"] not in ["None", "Error"]:
             weight = custom_weights.get(obj["object"], 1.0)
             score += obj["confidence"] / 100 * weight
     return min(score, 100.0)  # Cap at 100
@@ -115,63 +115,68 @@ def cancel_analysis():
     cancel_flag = True
     return "Analysis cancellation requested. Please wait for the current operation to stop."
 
-# Format results as a detailed report
-def format_detailed_report(results, timestamp, keyword: str = ""):
+# Format results as a detailed report (optimized)
+def format_detailed_report(results, timestamp, findings, keyword: str = ""):
+    global cancel_flag
     filtered_results = keyword_search(results, keyword)
-    report = f"# Investigation Report\n**Generated on**: {timestamp}\n\n## Key Findings\n"
-    harmful_found = any(len(result.get("harmful_objects", [])) > 0 and result["harmful_objects"][0]["object"] != "None" for result in filtered_results)
-    high_similarity = any(result.get("similarity_to_image_1", 0) > 0.8 for result in filtered_results if "similarity_to_image_1" in result)
-    faces_detected = any(result.get("faces", "No faces detected") != "No faces detected" for result in filtered_results)
-    threats_detected = any(result.get("threat_score", 0) > 50 for result in filtered_results)
-    if harmful_found:
-        report += "- Harmful objects detected in one or more images.\n"
-    if high_similarity:
-        report += "- High similarity (>0.8) detected between images.\n"
-    if faces_detected:
-        report += "- Faces detected with identifiable attributes.\n"
-    if threats_detected:
-        report += "- High threat scores (>50) detected in one or more images.\n"
-    if not (harmful_found or high_similarity or faces_detected or threats_detected):
-        report += "- No critical findings detected.\n"
-    
-    report += "\n## Analysis Details\n"
+    report_lines = [f"# Investigation Report\n**Generated on**: {timestamp}\n\n## Key Findings\n"]
+    
+    # Add key findings from precomputed dictionary
+    if findings["harmful_found"]:
+        report_lines.append("- Harmful objects detected in one or more images.\n")
+    if findings["high_similarity"]:
+        report_lines.append("- High similarity (>0.8) detected between images.\n")
+    if findings["faces_detected"]:
+        report_lines.append("- Faces detected with identifiable attributes.\n")
+    if findings["threats_detected"]:
+        report_lines.append("- High threat scores (>50) detected in one or more images.\n")
+    if not any(findings.values()):
+        report_lines.append("- No critical findings detected.\n")
+    
+    if cancel_flag:
+        return "".join(report_lines) + "\n**Report Generation Cancelled**"
+    
+    report_lines.append("\n## Analysis Details\n")
     for result in filtered_results:
-        report += f"### Image ID: {result['image_id']}\n"
+        if cancel_flag:
+            return "".join(report_lines) + "\n**Report Generation Cancelled**"
+        report_lines.append(f"### Image ID: {result['image_id']}\n")
         if "metadata" in result:
-            report += f"- **Metadata**: Timestamp: {result['metadata']['timestamp']}, GPS: {result['metadata']['gps']['latitude']}, {result['metadata']['gps']['longitude']}, Camera: {result['metadata']['camera']}\n"
+            meta = result["metadata"]
+            report_lines.append(f"- **Metadata**: Timestamp: {meta.get('timestamp', 'N/A')}, GPS: {meta.get('gps', {}).get('latitude', 'N/A')}, {meta.get('gps', {}).get('longitude', 'N/A')}, Camera: {meta.get('camera', 'N/A')}\n")
         if "description" in result:
-            report += f"- **Description**: {result['description']}\n"
+            report_lines.append(f"- **Description**: {result['description']}\n")
         if "signs" in result:
-            report += f"- **Signs**: {result['signs']}\n"
+            report_lines.append(f"- **Signs**: {result['signs']}\n")
         if "harmful_objects" in result:
             harmful_str = ", ".join([f"{obj['object']} ({obj['confidence']}%)" for obj in result['harmful_objects']])
-            report += f"- **Harmful Objects**: {harmful_str}\n"
+            report_lines.append(f"- **Harmful Objects**: {harmful_str}\n")
         if "similarity_to_image_1" in result:
-            report += f"- **Similarity**: {result['similarity_to_image_1']:.2f}\n"
+            report_lines.append(f"- **Similarity**: {result['similarity_to_image_1']:.2f}\n")
         if "faces" in result:
-            report += f"- **Faces**: {result['faces']}\n"
+            report_lines.append(f"- **Faces**: {result['faces']}\n")
         if "objects" in result:
             objects_str = ", ".join([f"{obj['object']} at {obj['bbox']}" for obj in result['objects']])
-            report += f"- **Objects**: {objects_str}\n"
+            report_lines.append(f"- **Objects**: {objects_str}\n")
         if "clothing" in result:
-            report += f"- **Clothing/Colors**: {result['clothing']}\n"
+            report_lines.append(f"- **Clothing/Colors**: {result['clothing']}\n")
         if "scene_context" in result:
-            report += f"- **Scene**: {result['scene_context']}\n"
+            report_lines.append(f"- **Scene**: {result['scene_context']}\n")
         if "activity" in result:
-            report += f"- **Activity**: {result['activity']}\n"
+            report_lines.append(f"- **Activity**: {result['activity']}\n")
         if "threat_score" in result:
-            report += f"- **Threat Score**: {result['threat_score']:.1f}/100\n"
+            report_lines.append(f"- **Threat Score**: {result['threat_score']:.1f}/100\n")
         if "annotation" in result:
-            report += f"- **Annotation**: {result['annotation']}\n"
+            report_lines.append(f"- **Annotation**: {result['annotation']}\n")
         if "flag" in result and result["flag"]:
-            report += f"- **Flagged**: Yes\n"
+            report_lines.append(f"- **Flagged**: Yes\n")
         if "comments" in result:
-            report += f"- **Investigator Comments**: {result['comments']}\n"
+            report_lines.append(f"- **Investigator Comments**: {result['comments']}\n")
         if "error" in result:
-            report += f"- **Error**: {result['error']}\n"
-        report += "\n"
+            report_lines.append(f"- **Error**: {result['error']}\n")
+        report_lines.append("\n")
     
-    return report
+    return "".join(report_lines)
 
 @app.post("/predict")
 async def predict(
@@ -417,11 +422,19 @@ def gradio_predict(
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     errors = []
     
+    # Precompute key findings
+    findings = {
+        "harmful_found": False,
+        "high_similarity": False,
+        "faces_detected": False,
+        "threats_detected": False
+    }
+    
     for idx, image in enumerate(images, 1):
         if cancel_flag:
             output = f"**Analysis Cancelled at**: {timestamp}\n"
             if results:
-                output += "\n**Partial Results**:\n" + format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+                output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
             return output
         try:
             # Convert Gradio image input to PIL and resize
@@ -438,7 +451,7 @@ def gradio_predict(
                 if cancel_flag:
                     output = f"**Analysis Cancelled at**: {timestamp}\n"
                     if results:
-                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                     return output
                 result["metadata"] = extract_metadata(image_data.getvalue())
             
@@ -446,7 +459,7 @@ def gradio_predict(
                 if cancel_flag:
                     output = f"**Analysis Cancelled at**: {timestamp}\n"
                     if results:
-                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                     return output
                 prompt_desc = "<image> Provide a " + ("brief description of the image." if description_level == "basic" else "detailed description of the image, including objects, colors, people, and environmental context.")
                 inputs = processor(text=[prompt_desc], images=[image], return_tensors="pt", padding=True)
@@ -457,7 +470,7 @@ def gradio_predict(
                 if cancel_flag:
                     output = f"**Analysis Cancelled at**: {timestamp}\n"
                     if results:
-                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                     return output
                 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)
@@ -468,7 +481,7 @@ def gradio_predict(
                 if cancel_flag:
                     output = f"**Analysis Cancelled at**: {timestamp}\n"
                     if results:
-                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                     return output
                 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)
@@ -480,23 +493,27 @@ def gradio_predict(
                         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}]
+                if harmful_detected and result["harmful_objects"][0]["object"] != "None":
+                    findings["harmful_found"] = True
             
             if faces:
                 if cancel_flag:
                     output = f"**Analysis Cancelled at**: {timestamp}\n"
                     if results:
-                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                     return output
                 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"
+                if result["faces"] != "No faces detected":
+                    findings["faces_detected"] = True
             
             if objects:
                 if cancel_flag:
                     output = f"**Analysis Cancelled at**: {timestamp}\n"
                     if results:
-                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                     return output
                 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)
@@ -507,7 +524,7 @@ def gradio_predict(
                 if cancel_flag:
                     output = f"**Analysis Cancelled at**: {timestamp}\n"
                     if results:
-                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                     return output
                 prompt_activity = "<image> Describe any activities or actions occurring in the image, such as walking, running, or driving."
                 inputs_activity = processor(text=[prompt_activity], images=[image], return_tensors="pt", padding=True)
@@ -519,7 +536,7 @@ def gradio_predict(
                 if cancel_flag:
                     output = f"**Analysis Cancelled at**: {timestamp}\n"
                     if results:
-                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                     return output
                 prompt_clothing = "<image> Identify clothing items and their colors worn by people in the image."
                 inputs_clothing = processor(text=[prompt_clothing], images=[image], return_tensors="pt", padding=True)
@@ -530,7 +547,7 @@ def gradio_predict(
                 if cancel_flag:
                     output = f"**Analysis Cancelled at**: {timestamp}\n"
                     if results:
-                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                     return output
                 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)
@@ -541,9 +558,11 @@ def gradio_predict(
                 if cancel_flag:
                     output = f"**Analysis Cancelled at**: {timestamp}\n"
                     if results:
-                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                     return output
                 result["threat_score"] = calculate_threat_score(result.get("harmful_objects", []), weights)
+                if result["threat_score"] > 50:
+                    findings["threats_detected"] = True
             
             # User Options & Controls
             if annotation.strip():
@@ -557,7 +576,7 @@ def gradio_predict(
                 if cancel_flag:
                     output = f"**Analysis Cancelled at**: {timestamp}\n"
                     if results:
-                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+                        output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                     return output
                 inputs_emb = processor(images=[image], return_tensors="pt", padding=True)
                 with torch.no_grad():
@@ -597,12 +616,14 @@ def gradio_predict(
             if cancel_flag:
                 output = f"**Analysis Cancelled at**: {timestamp}\n"
                 if results:
-                    output += "\n**Partial Results**:\n" + format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+                    output += "\n**Partial Results**:\n" + format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
                 return output
             sim = np.dot(base_embedding, image_embeddings[i].T) / (
                 np.linalg.norm(base_embedding) * np.linalg.norm(image_embeddings[i])
             )
             results[i]["similarity_to_image_1"] = float(sim[0][0])
+            if results[i]["similarity_to_image_1"] > 0.8:
+                findings["high_similarity"] = True
     
     # Filter by attributes
     if filter_attributes.strip():
@@ -621,14 +642,14 @@ def gradio_predict(
         case_output = manage_case_folder(case_folder, results, case_action)
     
     # Format output
-    output = format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
+    output = format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors)
     
     # Reset cancellation flag
     cancel_flag = False
     return output
 
 # Helper function to format results
-def format_results(results, timestamp, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors):
+def format_results(results, timestamp, findings, keyword_search, json_export, detailed_report, case_folder, case_action, filter_attributes, annotation, flag_images, comments, errors):
     output = f"**Analysis Timestamp**: {timestamp}\n\n"
     headers = ["Image ID"]
     if metadata:
@@ -715,7 +736,7 @@ def format_results(results, timestamp, keyword_search, json_export, detailed_rep
     
     # Handle detailed report
     if detailed_report:
-        output += "\n**Detailed Report**:\n" + format_detailed_report(results, timestamp, keyword_search)
+        output += "\n**Detailed Report**:\n" + format_detailed_report(results, timestamp, findings, keyword_search)
     
     # Append case folder output
     if case_folder.strip() and case_output:
@@ -729,9 +750,9 @@ with gr.Blocks() as iface:
     gr.Markdown("Upload up to 3 images (up to 10MB each, any resolution) ")
     
     with gr.Row():
-        image_1 = gr.Image(label="Upload Image 1 ")
-        image_2 = gr.Image(label="Upload Image 2 ")
-        image_3 = gr.Image(label="Upload Image 3 ")
+        image_1 = gr.Image(label="Upload Image 1 (up to 10MB)")
+        image_2 = gr.Image(label="Upload Image 2 (up to 10MB)")
+        image_3 = gr.Image(label="Upload Image 3 (up to 10MB)")
     
     with gr.Row():
         with gr.Column():