Update app.py

app.py

@@ -5,20 +5,18 @@ from groq import Groq
 import time
 from PIL import Image as PILImage
 import io
-import os
 import base64
 import torch
 
-
-class SafetyMonitor:
+class RobustSafetyMonitor:
     def __init__(self):
-        """Initialize Safety Monitor with configuration."""
+        """Initialize the robust safety detection tool with configuration."""
         self.client = Groq()
         self.model_name = "llama-3.2-90b-vision-preview"
         self.max_image_size = (800, 800)
         self.colors = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (255, 255, 0), (255, 0, 255)]
         
-        # Load YOLOv5 model for object detection
+        # Load YOLOv5 model for general object detection
         self.yolo_model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
 
     def preprocess_image(self, frame):
@@ -61,14 +59,15 @@ class SafetyMonitor:
         return bbox_data, labels
 
     def analyze_frame(self, frame):
-        """Perform safety analysis on the frame."""
+        """Perform safety analysis on the frame using Llama Vision 3.2."""
         if frame is None:
             return "No frame received", {}
     
         frame = self.preprocess_image(frame)
-        image_url = self.encode_image(frame)
+        image_base64 = self.encode_image(frame)
 
         try:
+            # Use Llama Vision 3.2 to analyze the context of the image and detect risks
             completion = self.client.chat.completions.create(
                 model=self.model_name,
                 messages=[
@@ -77,19 +76,21 @@ class SafetyMonitor:
                         "content": [
                             {
                                 "type": "text",
-                                "text": "Identify and list safety concerns in this workplace image. For each issue found, include its location and specific safety concern. Look for hazards related to PPE, ergonomics, equipment, environment, and work procedures."
+                                "text": """Analyze this workplace image and identify any potential safety risks. 
+                                           Consider the positioning of workers, the equipment, materials, and environment. 
+                                           Flag risks like improper equipment use, worker proximity to danger zones, unstable materials, and environmental hazards."""
                             },
                             {
                                 "type": "image_url",
                                 "image_url": {
-                                    "url": image_url
+                                    "url": f"data:image/jpeg;base64,{image_base64}"  # Use base64 for image
                                 }
                             }
                         ]
                     }
                 ],
                 temperature=0.7,
-                max_tokens=500,
+                max_tokens=1024,
                 stream=False
             )
             return completion.choices[0].message.content, {}
@@ -97,11 +98,12 @@ class SafetyMonitor:
             print(f"Analysis error: {str(e)}")
             return f"Analysis Error: {str(e)}", {}
 
-    def draw_bounding_boxes(self, image, bboxes, labels):
-        """Draw bounding boxes around detected objects."""
+    def draw_bounding_boxes(self, image, bboxes, labels, safety_issues):
+        """Draw bounding boxes around objects based on safety issues flagged by Llama Vision."""
         font = cv2.FONT_HERSHEY_SIMPLEX
         font_scale = 0.5
         thickness = 2
+        
         for idx, bbox in enumerate(bboxes):
             x1, y1, x2, y2, conf, class_id = bbox
             label = labels[int(class_id)]
@@ -110,68 +112,55 @@ class SafetyMonitor:
             # Draw bounding box
             cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), color, thickness)
             
-            # Draw label
-            label_text = f"{label} {conf:.2f}"
-            cv2.putText(image, label_text, (int(x1), int(y1) - 10), font, font_scale, 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)
+            else:
+                label_text = f"{label} {conf:.2f}"
+                cv2.putText(image, label_text, (int(x1), int(y1) - 10), font, font_scale, color, thickness)
+
         return image
 
     def process_frame(self, frame):
-        """Main processing pipeline for dynamic safety analysis."""
+        """Main processing pipeline for dynamic, robust safety analysis."""
         if frame is None:
             return None, "No image provided"
     
         try:
             # Detect objects dynamically in the image using YOLO
             bbox_data, labels = self.detect_objects(frame)
-            frame_with_boxes = self.draw_bounding_boxes(frame, bbox_data, labels)
-    
-            # Get dynamic safety analysis from Groq's model
+            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)
-            
-            # Dynamically parse the analysis to find any safety issues flagged
+
+            # Dynamically parse the analysis to identify safety issues flagged
             safety_issues = self.parse_safety_analysis(analysis)
             
-            # Dynamically link detected objects to safety issues
-            for issue in safety_issues:
-                if 'helmet' in issue.lower():
-                    for idx, bbox in enumerate(bbox_data):
-                        x1, y1, x2, y2, conf, class_id = bbox
-                        if labels[int(class_id)] == 'person':
-                            # Dynamically label the missing helmet issue for detected persons
-                            cv2.putText(frame_with_boxes, "No Helmet!", (int(x1), int(y1) - 20), 
-                                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
-                            cv2.rectangle(frame_with_boxes, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)
-    
-                # Add more dynamic checks here for gloves, boots, etc.
-                if 'glove' in issue.lower():
-                    for idx, bbox in enumerate(bbox_data):
-                        x1, y1, x2, y2, conf, class_id = bbox
-                        if labels[int(class_id)] == 'person':
-                            # Dynamically label missing gloves for detected persons
-                            cv2.putText(frame_with_boxes, "No Gloves!", (int(x1), int(y1) - 20), 
-                                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
-                            cv2.rectangle(frame_with_boxes, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 255), 2)
-            
-            return frame_with_boxes, 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)
+
+            return annotated_frame, analysis
     
         except Exception as e:
             print(f"Processing error: {str(e)}")
             return None, f"Error processing image: {str(e)}"
     
     def parse_safety_analysis(self, analysis):
-        """Dynamically parse the safety analysis to identify issues."""
+        """Dynamically parse the safety analysis to identify contextual issues."""
         safety_issues = []
         for line in analysis.split('\n'):
-            if "missing" in line.lower() or "no" in line.lower():
+            if "risk" in line.lower() or "hazard" in line.lower():
                 safety_issues.append(line.strip())
         return safety_issues
 
+
 def create_monitor_interface():
-    monitor = SafetyMonitor()
+    monitor = RobustSafetyMonitor()
     
     with gr.Blocks() as demo:
-        gr.Markdown("# Safety Analysis System powered by Llama 3.2 90b vision")
+        gr.Markdown("# Robust Safety Analysis System powered by Llama Vision 3.2")
         
         with gr.Row():
             input_image = gr.Image(label="Upload Image")
@@ -199,7 +188,7 @@ def create_monitor_interface():
         ## Instructions:
         1. Upload any workplace/safety-related image
         2. View identified hazards and their locations
-        3. Read detailed analysis of safety concerns
+        3. Read detailed analysis of safety concerns based on the image
         """)
 
     return demo
@@ -207,4 +196,3 @@ def create_monitor_interface():
 if __name__ == "__main__":
     demo = create_monitor_interface()
     demo.launch()
-