Update app.py

app.py

@@ -260,55 +260,254 @@ class RobustSafetyMonitor:
 
 
 def create_monitor_interface():
-    """Create the Gradio interface."""
-    monitor = RobustSafetyMonitor()
+    api_key = os.getenv("GROQ_API_KEY")
     
-    with gr.Blocks(theme=gr.themes.Base()) as demo:
-        gr.Markdown("# Workplace Safety Analysis System")
-        gr.Markdown("Powered by Groq LLaVA Vision and YOLOv5")
-        
-        with gr.Row():
-            input_image = gr.Image(label="Upload Workplace Image", type="numpy")
-            output_image = gr.Image(label="Safety Analysis Visualization")
+    class SafetyMonitor:
+        def __init__(self):
+            """Initialize Safety Monitor 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)]
+
+        def resize_image(self, image):
+            """Resize image while maintaining aspect ratio."""
+            height, width = image.shape[:2]
+            aspect = width / height
+            
+            if width > height:
+                new_width = min(self.max_image_size[0], width)
+                new_height = int(new_width / aspect)
+            else:
+                new_height = min(self.max_image_size[1], height)
+                new_width = int(new_height * aspect)
+                
+            return cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
+
+        def analyze_frame(self, frame: np.ndarray) -> str:
+            """Analyze frame for safety concerns."""
+            if frame is None:
+                return "No frame received"
+                
+            # Convert and resize image
+            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)
+            
+            frame = self.resize_image(frame)
+            frame_pil = PILImage.fromarray(frame)
+            
+            # Convert to base64
+            buffered = io.BytesIO()
+            frame_pil.save(buffered, 
+                         format="JPEG", 
+                         quality=95,  # High quality for better analysis
+                         optimize=True)
+            img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
+            image_url = f"data:image/jpeg;base64,{img_base64}"
+            
+            try:
+                completion = self.client.chat.completions.create(
+                    model=self.model_name,
+                    messages=[
+                        {
+                            "role": "user",
+                            "content": [
+                                {
+                                    "type": "text",
+                                    "text": """Analyze this workplace image for safety hazards. For each hazard:
+                                    1. Specify the exact location (e.g., center, top-left, bottom-right)
+                                    2. Describe the safety concern in detail
+
+                                    Format each finding as:
+                                    - <location>position:detailed safety description</location>
+
+                                    Consider:
+                                    - PPE usage and compliance
+                                    - Ergonomic risks
+                                    - Equipment safety
+                                    - Environmental hazards
+                                    - Work procedures
+                                    - Material handling
+                                    """
+                                },
+                                {
+                                    "type": "image_url",
+                                    "image_url": {
+                                        "url": image_url
+                                    }
+                                }
+                            ]
+                        }
+                    ],
+                    temperature=0.5,
+                    max_tokens=500,
+                    stream=False
+                )
+                return completion.choices[0].message.content
+            except Exception as e:
+                print(f"Analysis error: {str(e)}")
+                return f"Analysis Error: {str(e)}"
+
+        def draw_observations(self, image, observations):
+            """Draw safety observations with accurate locations."""
+            height, width = image.shape[:2]
+            font = cv2.FONT_HERSHEY_SIMPLEX
+            font_scale = 0.5
+            thickness = 2
+            
+            def get_region_coordinates(location_text):
+                """Get coordinates based on location description."""
+                location_text = location_text.lower()
+                regions = {
+                    # Basic positions
+                    'center': (width//3, height//3, 2*width//3, 2*height//3),
+                    'top': (width//4, 0, 3*width//4, height//3),
+                    'bottom': (width//4, 2*height//3, 3*width//4, height),
+                    'left': (0, height//4, width//3, 3*height//4),
+                    'right': (2*width//3, height//4, width, 3*height//4),
+                    'top-left': (0, 0, width//3, height//3),
+                    'top-right': (2*width//3, 0, width, height//3),
+                    'bottom-left': (0, 2*height//3, width//3, height),
+                    'bottom-right': (2*width//3, 2*height//3, width, height),
+                    
+                    # Work areas
+                    'workspace': (width//4, height//4, 3*width//4, 3*height//4),
+                    'machine': (2*width//3, 0, width, height),
+                    'equipment': (2*width//3, height//3, width, 2*height//3),
+                    'material': (0, 2*height//3, width//3, height),
+                    'ground': (0, 2*height//3, width, height),
+                    'floor': (0, 3*height//4, width, height),
+                    
+                    # Body regions
+                    'body': (width//3, height//3, 2*width//3, 2*height//3),
+                    'hands': (width//2, height//2, 3*width//4, 2*height//3),
+                    'head': (width//3, 0, 2*width//3, height//4),
+                    'feet': (width//3, 3*height//4, 2*width//3, height),
+                    'back': (width//3, height//3, 2*width//3, 2*height//3),
+                    'knees': (width//3, 2*height//3, 2*width//3, height),
+                    
+                    # Special areas
+                    'workspace': (width//4, height//4, 3*width//4, 3*height//4),
+                    'working-area': (width//4, height//4, 3*width//4, 3*height//4),
+                    'surrounding': (0, 0, width, height),
+                    'background': (0, 0, width, height)
+                }
+
+                # Find best matching region
+                best_match = 'center'  # default
+                max_match_length = 0
+                
+                for region_name in regions.keys():
+                    if region_name in location_text and len(region_name) > max_match_length:
+                        best_match = region_name
+                        max_match_length = len(region_name)
+                
+                return regions[best_match]
+
+            for idx, obs in enumerate(observations):
+                color = self.colors[idx % len(self.colors)]
+                
+                # Split location and description if available
+                parts = obs.split(':')
+                if len(parts) >= 2:
+                    location = parts[0]
+                    description = ':'.join(parts[1:])
+                else:
+                    location = 'center'
+                    description = obs
+                
+                # Get region coordinates
+                x1, y1, x2, y2 = get_region_coordinates(location)
+                
+                # Draw rectangle
+                cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
+                
+                # Add label
+                label = description[:50] + "..." if len(description) > 50 else description
+                label_size = cv2.getTextSize(label, font, font_scale, thickness)[0]
+                
+                # Position text above box
+                text_x = max(0, x1)
+                text_y = max(20, y1 - 5)
+                
+                # Draw text background
+                cv2.rectangle(image, 
+                            (text_x, text_y - label_size[1] - 5),
+                            (text_x + label_size[0], text_y),
+                            color, -1)
+                
+                # Draw text
+                cv2.putText(image, label, (text_x, text_y - 5),
+                           font, font_scale, (255, 255, 255), thickness)
+            
+            return image
+
+        def process_frame(self, frame: np.ndarray) -> tuple[np.ndarray, str]:
+            """Process frame and generate safety analysis."""
+            if frame is None:
+                return None, "No image provided"
+                
+            analysis = self.analyze_frame(frame)
+            display_frame = self.resize_image(frame.copy())
+            
+            # Parse observations
+            observations = []
+            for line in analysis.split('\n'):
+                line = line.strip()
+                if line.startswith('-'):
+                    if '<location>' in line and '</location>' in line:
+                        start = line.find('<location>') + len('<location>')
+                        end = line.find('</location>')
+                        observation = line[start:end].strip()
+                        if observation:
+                            observations.append(observation)
+            
+            # Draw observations
+            if observations:
+                annotated_frame = self.draw_observations(display_frame, observations)
+                return annotated_frame, analysis
+            
+            return display_frame, analysis
+
+    # Create interface
+    monitor = SafetyMonitor()
+    
+    with gr.Blocks() as demo:
+        gr.Markdown("# Safety Analysis System powered by Llama 3.2 90b vision")
         
         with gr.Row():
-            analysis_text = gr.Textbox(
-                label="Detailed Safety Analysis",
-                lines=8,
-                placeholder="Safety analysis will appear here..."
-            )
+            input_image = gr.Image(label="Upload Image")
+            output_image = gr.Image(label="Safety Analysis")
         
+        analysis_text = gr.Textbox(label="Detailed Analysis", lines=5)
+            
         def analyze_image(image):
             if image is None:
-                return None, "Please upload an image"
+                return None, "No image provided"
             try:
                 processed_frame, analysis = monitor.process_frame(image)
                 return processed_frame, analysis
             except Exception as e:
-                print(f"Analysis error: {str(e)}")
-                return None, f"Error analyzing image: {str(e)}"
+                print(f"Processing error: {str(e)}")
+                return None, f"Error processing image: {str(e)}"
             
-        input_image.upload(
+        input_image.change(
             fn=analyze_image,
             inputs=input_image,
             outputs=[output_image, analysis_text]
         )
 
         gr.Markdown("""
-        ## 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
+        ## Instructions:
+        1. Upload a workplace image
+        2. View detected safety concerns
+        3. Check detailed analysis
         """)
 
     return demo
 
 if __name__ == "__main__":
     demo = create_monitor_interface()
-    demo.launch(share=True)
+    demo.launch()