Update app.py

app.py

@@ -6,21 +6,32 @@ from PIL import Image as PILImage
 import io
 import base64
 import torch
+import warnings
+from typing import Tuple, List, Dict, Optional
 
+# Suppress the CUDA autocast warning
+warnings.filterwarnings('ignore', category=FutureWarning)
 
 class RobustSafetyMonitor:
     def __init__(self):
         """Initialize the robust safety detection tool with configuration."""
         self.client = Groq()
-        self.model_name = "llama-3.2-90b-vision-preview"
+        self.model_name = "llama-3.2-11b-vision-preview"  # Updated to use the correct model
         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 general object detection
         self.yolo_model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
+        
+        # Force CPU inference if CUDA is causing issues
+        self.yolo_model.cpu()
+        self.yolo_model.eval()
 
-    def preprocess_image(self, frame):
+    def preprocess_image(self, frame: np.ndarray) -> np.ndarray:
         """Process image for analysis."""
+        if frame is None:
+            raise ValueError("No image provided")
+            
         if len(frame.shape) == 2:
             frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
         elif len(frame.shape) == 3 and frame.shape[2] == 4:
@@ -28,7 +39,7 @@ class RobustSafetyMonitor:
         
         return self.resize_image(frame)
 
-    def resize_image(self, image):
+    def resize_image(self, image: np.ndarray) -> np.ndarray:
         """Resize image while maintaining aspect ratio."""
         height, width = image.shape[:2]
         if height > self.max_image_size[1] or width > self.max_image_size[0]:
@@ -42,35 +53,37 @@ class RobustSafetyMonitor:
             return cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
         return image
 
-    def encode_image(self, frame):
+    def encode_image(self, frame: np.ndarray) -> str:
         """Convert image to base64 encoding with proper formatting."""
-        frame_pil = PILImage.fromarray(frame)
-        buffered = io.BytesIO()
-        frame_pil.save(buffered, format="JPEG", quality=95)  # Ensure JPEG format
-        img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
-        
-        # Prepend the required image data type to base64 string
-        img_base64_str = f"data:image/jpeg;base64,{img_base64}"
-        return img_base64_str
+        try:
+            frame_pil = PILImage.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+            buffered = io.BytesIO()
+            frame_pil.save(buffered, format="JPEG", quality=95)
+            img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
+            return f"data:image/jpeg;base64,{img_base64}"
+        except Exception as e:
+            raise ValueError(f"Error encoding image: {str(e)}")
 
-    def detect_objects(self, frame):
+    def detect_objects(self, frame: np.ndarray) -> Tuple[np.ndarray, Dict]:
         """Detect objects using YOLOv5."""
-        results = self.yolo_model(frame)
-        # Extract bounding boxes, class labels, and confidence scores
-        bbox_data = results.xyxy[0].numpy()  # Bounding box coordinates
-        labels = results.names  # Class names
-        return bbox_data, labels
-
-    def analyze_frame(self, frame):
-        """Perform safety analysis on the frame using Llama Vision 3.2."""
+        try:
+            with torch.no_grad():
+                results = self.yolo_model(frame)
+            bbox_data = results.xyxy[0].cpu().numpy()
+            labels = results.names
+            return bbox_data, labels
+        except Exception as e:
+            raise ValueError(f"Error detecting objects: {str(e)}")
+
+    def analyze_frame(self, frame: np.ndarray) -> Tuple[List[Dict], str]:
+        """Perform safety analysis on the frame using Llama Vision."""
         if frame is None:
-            return "No frame received", {}
+            return [], "No frame received"
     
-        frame = self.preprocess_image(frame)
-        image_base64 = self.encode_image(frame)
-
         try:
-            # Use Llama Vision 3.2 to analyze the context of the image and detect risks
+            frame = self.preprocess_image(frame)
+            image_base64 = self.encode_image(frame)
+
             completion = self.client.chat.completions.create(
                 model=self.model_name,
                 messages=[
@@ -80,13 +93,13 @@ class RobustSafetyMonitor:
                             {
                                 "type": "text",
                                 "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."""
+                                List each risk on a new line starting with 'Risk:'.
+                                Format: Risk: [Object/Area] - [Description of hazard]"""
                             },
                             {
                                 "type": "image_url",
                                 "image_url": {
-                                    "url": image_base64  # Corrected: Now sending the proper data format
+                                    "url": image_base64
                                 }
                             }
                         ]
@@ -96,98 +109,134 @@ class RobustSafetyMonitor:
                 max_tokens=1024,
                 stream=False
             )
-            # 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
+            
+            # Get the response content safely
+            try:
+                response = completion.choices[0].message.content
+            except AttributeError:
+                response = str(completion.choices[0].message)
+                
+            safety_issues = self.parse_safety_analysis(response)
+            return safety_issues, response
 
         except Exception as e:
             print(f"Analysis error: {str(e)}")
-            return f"Analysis Error: {str(e)}", {}
+            return [], f"Analysis Error: {str(e)}"
 
-    def draw_bounding_boxes(self, image, bboxes, labels, safety_issues):
-        """Draw bounding boxes around objects based on safety issues flagged by Llama Vision."""
+    def draw_bounding_boxes(self, image: np.ndarray, bboxes: np.ndarray, 
+                          labels: Dict, safety_issues: List[Dict]) -> np.ndarray:
+        """Draw bounding boxes around objects based on safety issues."""
+        image_copy = image.copy()
         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)]
-            color = self.colors[idx % len(self.colors)]
-            
-            # Draw bounding box
-            cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), color, thickness)
-            
-            # Link detected object to potential risks based on Llama Vision analysis
-            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)
+            try:
+                x1, y1, x2, y2, conf, class_id = bbox
+                label = labels[int(class_id)]
+                color = self.colors[idx % len(self.colors)]
+                
+                # Convert coordinates to integers
+                x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
+                
+                # Draw bounding box
+                cv2.rectangle(image_copy, (x1, y1), (x2, y2), color, thickness)
+                
+                # Check if object is associated with any safety issues
+                risk_found = False
+                for safety_issue in safety_issues:
+                    if safety_issue.get('object', '').lower() in label.lower():
+                        label_text = f"Risk: {safety_issue.get('description', '')}"
+                        y_pos = max(y1 - 10, 20)
+                        cv2.putText(image_copy, label_text, (x1, y_pos), font, 
+                                  font_scale, (0, 0, 255), thickness)
+                        risk_found = True
+                        break
+                        
+                if not risk_found:
+                    label_text = f"{label} {conf:.2f}"
+                    y_pos = max(y1 - 10, 20)
+                    cv2.putText(image_copy, label_text, (x1, y_pos), font, 
+                              font_scale, color, thickness)
+            except Exception as e:
+                print(f"Error drawing box: {str(e)}")
+                continue
 
-        return image
+        return image_copy
 
-    def process_frame(self, frame):
-        """Main processing pipeline for dynamic, robust safety analysis."""
+    def process_frame(self, frame: np.ndarray) -> Tuple[Optional[np.ndarray], str]:
+        """Main processing pipeline for safety analysis."""
         if frame is None:
             return None, "No image provided"
     
         try:
-            # Detect objects dynamically in the image using YOLO
+            # Detect objects
             bbox_data, labels = self.detect_objects(frame)
-            frame_with_boxes = self.draw_bounding_boxes(frame, bbox_data, labels, [])
-
-            # Get dynamic safety analysis from Llama Vision 3.2
+            
+            # Get safety analysis
             safety_issues, analysis = self.analyze_frame(frame)
-
-            # 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)
-
+            
+            # Draw annotations
+            annotated_frame = self.draw_bounding_boxes(frame, 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):
-        """Parse the safety analysis to identify contextual issues and link to objects."""
+    def parse_safety_analysis(self, analysis: str) -> List[Dict]:
+        """Parse the safety analysis text into structured data."""
         safety_issues = []
+        
+        if not isinstance(analysis, str):
+            return safety_issues
+            
         for line in analysis.split('\n'):
-            if "risk" in line.lower() or "hazard" in line.lower():
-                # Extract object involved and description
-                parts = line.split(':', 1)
-                if len(parts) == 2:
+            if "risk:" in line.lower():
+                try:
+                    # Extract object and description
+                    parts = line.lower().split('risk:', 1)[1].strip()
+                    if '-' in parts:
+                        obj, desc = parts.split('-', 1)
+                    else:
+                        obj, desc = parts, parts
+                        
                     safety_issues.append({
-                        "object": parts[0].strip(),
-                        "description": parts[1].strip()
+                        "object": obj.strip(),
+                        "description": desc.strip()
                     })
+                except Exception as e:
+                    print(f"Error parsing line: {line}, Error: {str(e)}")
+                    continue
+                    
         return safety_issues
 
 
 def create_monitor_interface():
+    """Create the Gradio interface for the safety monitoring system."""
     monitor = RobustSafetyMonitor()
     
     with gr.Blocks() as demo:
-        gr.Markdown("# Robust Safety Analysis System powered by Llama Vision 3.2")
+        gr.Markdown("# Workplace Safety Analysis System")
+        gr.Markdown("Powered by Groq LLaVA Vision and YOLOv5")
         
         with gr.Row():
-            input_image = gr.Image(label="Upload Image")
-            output_image = gr.Image(label="Safety Analysis")
+            input_image = gr.Image(label="Upload Workplace Image", type="numpy")
+            output_image = gr.Image(label="Safety Analysis Visualization")
+        
+        analysis_text = gr.Textbox(label="Detailed Safety Analysis", lines=5)
         
-        analysis_text = gr.Textbox(label="Detailed Analysis", lines=5)
-            
         def analyze_image(image):
             if image is None:
-                return None, "No image provided"
+                return None, "Please upload an image"
             try:
                 processed_frame, analysis = monitor.process_frame(image)
                 return processed_frame, analysis
             except Exception as e:
-                print(f"Processing error: {str(e)}")
-                return None, f"Error processing image: {str(e)}"
+                print(f"Analysis error: {str(e)}")
+                return None, f"Error analyzing image: {str(e)}"
             
         input_image.upload(
             fn=analyze_image,
@@ -196,14 +245,20 @@ def create_monitor_interface():
         )
 
         gr.Markdown("""
-        ## Instructions:
-        1. Upload any workplace/safety-related image
-        2. View identified hazards and their locations
-        3. Read detailed analysis of safety concerns based on the image
+        ## Instructions
+        1. Upload a workplace image for safety analysis
+        2. View detected hazards and their locations in the visualization
+        3. Read the detailed safety analysis below the images
+        
+        ## Features
+        - Real-time object detection
+        - AI-powered safety risk analysis
+        - Visual risk highlighting
+        - Detailed safety recommendations
         """)
 
     return demo
 
 if __name__ == "__main__":
     demo = create_monitor_interface()
-    demo.launch()
+    demo.launch(share=True)