Update app.py

app.py

@@ -7,15 +7,17 @@ from PIL import Image as PILImage
 import io
 import os
 import base64
+import random
 
 def create_monitor_interface():
     api_key = os.getenv("GROQ_API_KEY")
     
     class SafetyMonitor:
         def __init__(self):
-            self.client = Groq()  # API key handled by environment variable
+            self.client = Groq()
             self.model_name = "llama-3.2-90b-vision-preview"
-            self.max_image_size = (128, 128)
+            self.max_image_size = (640, 640)  # Increased size for better visibility
+            self.colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)]
             
         def resize_image(self, image):
             height, width = image.shape[:2]
@@ -47,7 +49,7 @@ def create_monitor_interface():
             buffered = io.BytesIO()
             frame_pil.save(buffered, 
                          format="JPEG", 
-                         quality=20, 
+                         quality=30, 
                          optimize=True)
             img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
             image_url = f"data:image/jpeg;base64,{img_base64}"
@@ -61,7 +63,9 @@ def create_monitor_interface():
                             "content": [
                                 {
                                     "type": "text",
-                                    "text": "Analyze this workplace for safety concerns. List any safety issues, missing PPE, or hazards you observe."
+                                    "text": """Analyze this workplace image and describe each safety concern in this format:
+                                    - <location>Description</location>
+                                    Use one line per issue, starting with a dash and location in tags."""
                                 },
                                 {
                                     "type": "image_url",
@@ -77,7 +81,7 @@ def create_monitor_interface():
                         }
                     ],
                     temperature=0.1,
-                    max_tokens=100,
+                    max_tokens=150,
                     top_p=1,
                     stream=False,
                     stop=None
@@ -87,30 +91,59 @@ def create_monitor_interface():
                 print(f"Detailed error: {str(e)}")
                 return f"Analysis Error: {str(e)}"
 
+        def draw_observations(self, image, observations):
+            height, width = image.shape[:2]
+            font = cv2.FONT_HERSHEY_SIMPLEX
+            font_scale = 0.5
+            thickness = 2
+            
+            # Generate random positions for each observation
+            for idx, obs in enumerate(observations):
+                color = self.colors[idx % len(self.colors)]
+                
+                # Generate random box position
+                box_width = width // 3
+                box_height = height // 3
+                x = random.randint(0, width - box_width)
+                y = random.randint(0, height - box_height)
+                
+                # Draw rectangle
+                cv2.rectangle(image, (x, y), (x + box_width, y + box_height), color, 2)
+                
+                # Add label with background
+                label = obs[:40] + "..." if len(obs) > 40 else obs
+                label_size = cv2.getTextSize(label, font, font_scale, thickness)[0]
+                cv2.rectangle(image, (x, y - 20), (x + label_size[0], y), color, -1)
+                cv2.putText(image, label, (x, y - 5), font, font_scale, (255, 255, 255), thickness)
+            
+            return image
+
         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()
+            display_frame = self.resize_image(frame.copy())
             
-            # Add text overlay
-            overlay = display_frame.copy()
-            height, width = display_frame.shape[:2]
-            cv2.rectangle(overlay, (5, 5), (width-5, 100), (0, 0, 0), -1)
-            cv2.addWeighted(overlay, 0.3, display_frame, 0.7, 0, display_frame)
+            # Parse observations from the analysis
+            observations = []
+            for line in analysis.split('\n'):
+                line = line.strip()
+                if line.startswith('-'):
+                    # Extract text between <location> tags if present
+                    if '<location>' in line and '</location>' in line:
+                        start = line.find('<location>') + len('<location>')
+                        end = line.find('</location>')
+                        observation = line[end + len('</location>'):].strip()
+                    else:
+                        observation = line[1:].strip()  # Remove the dash
+                    if observation:
+                        observations.append(observation)
             
-            # Add analysis text
-            y_position = 30
-            lines = analysis.split('\n')
-            for line in lines:
-                cv2.putText(display_frame, line[:80], (10, y_position),
-                           cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
-                y_position += 30
-                if y_position >= 90:
-                    break
-
-            return display_frame, analysis
+            # Draw observations on the image
+            annotated_frame = self.draw_observations(display_frame, observations)
+            
+            return annotated_frame, analysis
 
     # Create the main interface
     monitor = SafetyMonitor()
@@ -120,9 +153,9 @@ def create_monitor_interface():
         
         with gr.Row():
             input_image = gr.Image(label="Upload Image")
-            output_image = gr.Image(label="Results")
+            output_image = gr.Image(label="Annotated Results")
         
-        analysis_text = gr.Textbox(label="Analysis", lines=5)
+        analysis_text = gr.Textbox(label="Detailed Analysis", lines=5)
             
         def analyze_image(image):
             if image is None: