app.py

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 os
import base64

def create_monitor_interface():
    api_key = os.getenv("GROQ_API_KEY")
    
    class SafetyMonitor:
        def __init__(self):
            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):
            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 analyze_frame(self, frame: np.ndarray) -> str:
            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)
            
            # High quality image for better analysis
            buffered = io.BytesIO()
            frame_pil.save(buffered, 
                         format="JPEG", 
                         quality=95,
                         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 conditions and hazards. Focus on:

1. Work posture and ergonomics
2. PPE and safety equipment usage
3. Tool handling and techniques
4. Environmental conditions
5. Equipment and machinery safety
6. Ground conditions and hazards

Describe each safety condition observed, using this exact format:
- <location>position</location>: detailed safety observation

Examples:
- <location>center</location>: Improper kneeling posture without knee protection, risking joint injury
- <location>background</location>: Heavy machinery operating in close proximity creating hazard zone
- <location>ground</location>: Uneven surface and debris creating trip hazards

Be specific about locations and safety concerns."""
                                },
                                {
                                    "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 process_frame(self, frame: np.ndarray) -> tuple[np.ndarray, str]:
            if frame is None:
                return None, "No image provided"
                
            analysis = self.analyze_frame(frame)
            display_frame = frame.copy()
            
            # Parse observations from the formatted response
            observations = []
            lines = analysis.split('\n')
            for line in lines:
                if '<location>' in line and '</location>' in line:
                    start = line.find('<location>') + len('<location>')
                    end = line.find('</location>')
                    location = line[start:end].strip()
                    
                    # Get the description that follows the location tags
                    desc_start = line.find('</location>') + len('</location>:')
                    description = line[desc_start:].strip()
                    
                    if location and description:
                        observations.append({
                            'location': location,
                            'description': description
                        })
            
            # Draw observations if we found any
            if observations:
                annotated_frame = self.draw_observations(display_frame, observations)
                return annotated_frame, analysis
            
            return display_frame, analysis

        def draw_observations(self, image, observations):
            """Draw accurate bounding boxes based on safety issue locations."""
            height, width = image.shape[:2]
            font = cv2.FONT_HERSHEY_SIMPLEX
            font_scale = 0.5
            thickness = 2
            padding = 10
            
            def get_region_coordinates(position: str) -> tuple:
                """Get coordinates based on position description."""
                regions = {
                    'center': (width//3, height//3, 2*width//3, 2*height//3),
                    'background': (0, 0, width, height),
                    'top-left': (0, 0, width//3, height//3),
                    'top': (width//3, 0, 2*width//3, height//3),
                    'top-right': (2*width//3, 0, width, height//3),
                    'left': (0, height//3, width//3, 2*height//3),
                    'right': (2*width//3, height//3, width, 2*height//3),
                    'bottom-left': (0, 2*height//3, width//3, height),
                    'bottom': (width//3, 2*height//3, 2*width//3, height),
                    'bottom-right': (2*width//3, 2*height//3, width, height),
                    'ground': (0, 2*height//3, width, height),
                    'machinery': (0, 0, width//2, height),
                    'work-area': (width//4, height//4, 3*width//4, 3*height//4)
                }
                
                # Find best matching region
                position = position.lower()
                for key in regions.keys():
                    if key in position:
                        return regions[key]
                
                return regions['center']
            
            for idx, obs in enumerate(observations):
                color = self.colors[idx % len(self.colors)]
                
                # Get coordinates for this observation
                x1, y1, x2, y2 = get_region_coordinates(obs['location'])
                
                # Draw rectangle
                cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
                
                # Add label with background
                label = obs['description'][:50] + "..." if len(obs['description']) > 50 else obs['description']
                label_size, _ = cv2.getTextSize(label, font, font_scale, thickness)
                
                # Position text above the box
                text_x = max(0, x1)
                text_y = max(label_size[1] + padding, y1 - padding)
                
                # Draw text background
                cv2.rectangle(image, 
                            (text_x, text_y - label_size[1] - padding),
                            (text_x + label_size[0] + padding, text_y),
                            color, -1)
                
                # Draw text
                cv2.putText(image, label,
                           (text_x + padding//2, text_y - padding//2),
                           font, font_scale, (255, 255, 255), thickness)
            
            return image

    # Create the main interface
    monitor = SafetyMonitor()
    
    with gr.Blocks() as demo:
        gr.Markdown("# Safety Analysis System powered by Llama 3.2 90b vision")
        
        with gr.Row():
            input_image = gr.Image(label="Upload Image")
            output_image = gr.Image(label="Annotated Results")
        
        analysis_text = gr.Textbox(label="Detailed Analysis", lines=5)
            
        def analyze_image(image):
            if image is None:
                return None, "No image provided"
            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)}"
            
        input_image.change(
            fn=analyze_image,
            inputs=input_image,
            outputs=[output_image, analysis_text]
        )

        gr.Markdown("""
        ## Instructions:
        1. Upload an image to analyze safety conditions
        2. View annotated results showing safety concerns
        3. Read detailed analysis of identified issues
        """)

    return demo

demo = create_monitor_interface()
demo.launch()