Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 import cv2
 import numpy as np
 from groq import Groq
-import time
 from PIL import Image as PILImage
 import io
 import base64
@@ -49,7 +48,7 @@ class RobustSafetyMonitor:
         buffered = io.BytesIO()
         frame_pil.save(buffered, format="JPEG", quality=95)  # Ensure JPEG format
         img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
-        return img_base64  # Return only the base64 string, no "data:image/jpeg;base64,"
+        return img_base64  # Return only the base64 string
 
     def detect_objects(self, frame):
         """Detect objects using YOLOv5."""
@@ -94,7 +93,10 @@ class RobustSafetyMonitor:
                 max_tokens=1024,
                 stream=False
             )
-            return completion.choices[0].message.content, {}
+            # Process and parse the response correctly
+            response = completion.choices[0].message.content
+            return self.parse_safety_analysis(response), response  # Return parsed analysis and full response
+
         except Exception as e:
             print(f"Analysis error: {str(e)}")
             return f"Analysis Error: {str(e)}", {}
@@ -114,9 +116,11 @@ class RobustSafetyMonitor:
             cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), color, thickness)
             
             # Link detected object to potential risks based on Llama Vision analysis
-            if any(safety_issue.lower() in label.lower() for safety_issue in safety_issues):
-                label_text = f"Risk: {label}"
-                cv2.putText(image, label_text, (int(x1), int(y1) - 10), font, font_scale, (0, 0, 255), thickness)
+            for safety_issue in safety_issues:
+                if safety_issue['object'].lower() in label.lower():
+                    label_text = f"Risk: {safety_issue['description']}"
+                    cv2.putText(image, label_text, (int(x1), int(y1) - 10), font, font_scale, (0, 0, 255), thickness)
+                    break
             else:
                 label_text = f"{label} {conf:.2f}"
                 cv2.putText(image, label_text, (int(x1), int(y1) - 10), font, font_scale, color, thickness)
@@ -134,11 +138,8 @@ class RobustSafetyMonitor:
             frame_with_boxes = self.draw_bounding_boxes(frame, bbox_data, labels, [])
 
             # Get dynamic safety analysis from Llama Vision 3.2
-            analysis, _ = self.analyze_frame(frame)
+            safety_issues, analysis = self.analyze_frame(frame)
 
-            # Dynamically parse the analysis to identify safety issues flagged
-            safety_issues = self.parse_safety_analysis(analysis)
-            
             # Update the frame with bounding boxes based on safety issues flagged
             annotated_frame = self.draw_bounding_boxes(frame_with_boxes, bbox_data, labels, safety_issues)
 
@@ -149,11 +150,17 @@ class RobustSafetyMonitor:
             return None, f"Error processing image: {str(e)}"
     
     def parse_safety_analysis(self, analysis):
-        """Dynamically parse the safety analysis to identify contextual issues."""
+        """Parse the safety analysis to identify contextual issues and link to objects."""
         safety_issues = []
         for line in analysis.split('\n'):
             if "risk" in line.lower() or "hazard" in line.lower():
-                safety_issues.append(line.strip())
+                # Extract object involved and description
+                parts = line.split(':', 1)
+                if len(parts) == 2:
+                    safety_issues.append({
+                        "object": parts[0].strip(),
+                        "description": parts[1].strip()
+                    })
         return safety_issues