Update app.py

app.py

@@ -9,8 +9,9 @@ import torch
 import warnings
 from typing import Tuple, List, Dict, Optional
 
-# Suppress the CUDA autocast warning
+# Suppress warnings
 warnings.filterwarnings('ignore', category=FutureWarning)
+warnings.filterwarnings('ignore', category=UserWarning)
 
 class RobustSafetyMonitor:
     def __init__(self):
@@ -20,15 +21,59 @@ class RobustSafetyMonitor:
         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 with improved configuration
+        # Load YOLOv5 with optimized settings
         self.yolo_model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
-        self.yolo_model.conf = 0.25  # Lower confidence threshold for more detections
+        self.yolo_model.conf = 0.25  # Lower confidence threshold
         self.yolo_model.iou = 0.45   # Adjusted IOU threshold
         self.yolo_model.classes = None  # Detect all classes
         self.yolo_model.max_det = 50  # Increased maximum detections
         self.yolo_model.cpu()
         self.yolo_model.eval()
 
+        # Construction-specific keywords
+        self.construction_keywords = [
+            'person', 'worker', 'helmet', 'tool', 'machine', 'equipment',
+            'brick', 'block', 'pile', 'stack', 'surface', 'floor', 'ground',
+            'construction', 'building', 'structure'
+        ]
+
+    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:
+            frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2RGB)
+        
+        return self.resize_image(frame)
+
+    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]:
+            aspect = width / height
+            if width > height:
+                new_width = self.max_image_size[0]
+                new_height = int(new_width / aspect)
+            else:
+                new_height = self.max_image_size[1]
+                new_width = int(new_height * aspect)
+            return cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
+        return image
+
+    def encode_image(self, frame: np.ndarray) -> str:
+        """Convert image to base64 encoding."""
+        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: np.ndarray) -> Tuple[np.ndarray, Dict]:
         """Enhanced object detection using YOLOv5."""
         try:
@@ -40,7 +85,7 @@ class RobustSafetyMonitor:
 
             # Run inference with augmentation
             with torch.no_grad():
-                results = self.yolo_model(frame, augment=True)  # Enable test-time augmentation
+                results = self.yolo_model(frame, augment=True)
             
             # Get detections
             bbox_data = results.xyxy[0].cpu().numpy()
@@ -50,8 +95,7 @@ class RobustSafetyMonitor:
             processed_boxes = []
             for box in bbox_data:
                 x1, y1, x2, y2, conf, cls = box
-                # Additional filtering for construction site objects
-                if conf > 0.25:  # Keep lower confidence threshold for more detections
+                if conf > 0.25:  # Keep lower confidence threshold
                     processed_boxes.append(box)
             
             return np.array(processed_boxes), labels
@@ -59,6 +103,59 @@ class RobustSafetyMonitor:
             print(f"Error in object detection: {str(e)}")
             return np.array([]), {}
 
+    def analyze_frame(self, frame: np.ndarray) -> Tuple[List[Dict], str]:
+        """Perform safety analysis using Llama Vision."""
+        if frame is None:
+            return [], "No frame received"
+    
+        try:
+            frame = self.preprocess_image(frame)
+            image_base64 = self.encode_image(frame)
+
+            completion = self.client.chat.completions.create(
+                model=self.model_name,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": """Analyze this workplace image for safety risks. Focus on:
+                                1. Worker posture and positioning
+                                2. Equipment and tool safety
+                                3. Environmental hazards
+                                4. PPE compliance
+                                5. Material handling
+
+                                List each risk on a new line starting with 'Risk:'.
+                                Format: Risk: [Object/Area] - [Detailed description of hazard]"""
+                            },
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": image_base64
+                                }
+                            }
+                        ]
+                    }
+                ],
+                temperature=0.7,
+                max_tokens=1024,
+                stream=False
+            )
+            
+            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)}"
+
     def draw_bounding_boxes(self, image: np.ndarray, bboxes: np.ndarray, 
                           labels: Dict, safety_issues: List[Dict]) -> np.ndarray:
         """Improved bounding box visualization."""
@@ -67,28 +164,21 @@ class RobustSafetyMonitor:
         font_scale = 0.5
         thickness = 2
         
-        # Define construction-related keywords for better object association
-        construction_keywords = [
-            'person', 'worker', 'helmet', 'tool', 'machine', 'equipment',
-            'brick', 'block', 'pile', 'stack', 'surface', 'floor', 'ground',
-            'construction', 'building', 'structure'
-        ]
-        
         for idx, bbox in enumerate(bboxes):
             try:
                 x1, y1, x2, y2, conf, class_id = bbox
                 label = labels[int(class_id)]
                 
                 # Check if object is construction-related
-                is_relevant = any(keyword in label.lower() for keyword in construction_keywords)
+                is_relevant = any(keyword in label.lower() for keyword in self.construction_keywords)
                 
-                if is_relevant or conf > 0.35:  # Higher threshold for non-construction objects
+                if is_relevant or conf > 0.35:
                     color = self.colors[idx % len(self.colors)]
                     
                     # Convert coordinates to integers
                     x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
                     
-                    # Draw thicker bounding box for better visibility
+                    # Draw bounding box
                     cv2.rectangle(image_copy, (x1, y1), (x2, y2), color, thickness)
                     
                     # Check for associated safety issues
@@ -108,8 +198,8 @@ class RobustSafetyMonitor:
                         y_pos = max(y1 - 10, 20)
                         cv2.putText(image_copy, label_text, (x1, y_pos), font, 
                                   font_scale, color, thickness)
-                        
-                    # Draw additional markers for high-risk areas
+                    
+                    # Mark high-risk areas
                     if conf > 0.5 and any(risk_word in label.lower() for risk_word in 
                                         ['worker', 'person', 'equipment', 'machine']):
                         cv2.circle(image_copy, (int((x1 + x2)/2), int((y1 + y2)/2)), 
@@ -143,7 +233,7 @@ class RobustSafetyMonitor:
             return None, f"Error processing image: {str(e)}"
     
     def parse_safety_analysis(self, analysis: str) -> List[Dict]:
-        """Parse the safety analysis text into structured data."""
+        """Parse the safety analysis text."""
         safety_issues = []
         
         if not isinstance(analysis, str):
@@ -152,7 +242,6 @@ class RobustSafetyMonitor:
         for line in analysis.split('\n'):
             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)
@@ -171,10 +260,10 @@ class RobustSafetyMonitor:
 
 
 def create_monitor_interface():
-    """Create the Gradio interface for the safety monitoring system."""
+    """Create the Gradio interface."""
     monitor = RobustSafetyMonitor()
     
-    with gr.Blocks() as demo:
+    with gr.Blocks(theme=gr.themes.Base()) as demo:
         gr.Markdown("# Workplace Safety Analysis System")
         gr.Markdown("Powered by Groq LLaVA Vision and YOLOv5")
         
@@ -182,7 +271,12 @@ def create_monitor_interface():
             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)
+        with gr.Row():
+            analysis_text = gr.Textbox(
+                label="Detailed Safety Analysis",
+                lines=8,
+                placeholder="Safety analysis will appear here..."
+            )
         
         def analyze_image(image):
             if image is None: