Update app.py

app.py

@@ -14,151 +14,107 @@ warnings.filterwarnings('ignore', category=FutureWarning)
 
 class RobustSafetyMonitor:
     def __init__(self):
-        """Initialize the robust safety detection tool with configuration."""
+        """Initialize the safety detection tool with improved configuration."""
         self.client = Groq()
-        self.model_name = "llama-3.2-11b-vision-preview"  # Updated to use the correct model
+        self.model_name = "llama-3.2-11b-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)]
         
-        # Load YOLOv5 model for general object detection
+        # Load YOLOv5 model with improved configuration
         self.yolo_model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
-        
-        # Force CPU inference if CUDA is causing issues
+        self.yolo_model.conf = 0.25  # Lower confidence threshold for more detections
+        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()
 
-    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 with proper formatting."""
-        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]:
-        """Detect objects using YOLOv5."""
+        """Enhanced object detection using YOLOv5."""
         try:
+            # Ensure proper image format
+            if len(frame.shape) == 2:
+                frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
+            elif frame.shape[2] == 4:
+                frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2RGB)
+
+            # Run inference with augmentation
             with torch.no_grad():
-                results = self.yolo_model(frame)
+                results = self.yolo_model(frame, augment=True)  # Enable test-time augmentation
+            
+            # Get detections
             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"
-    
-        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 and identify any potential safety risks. 
-                                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
-                                }
-                            }
-                        ]
-                    }
-                ],
-                temperature=0.7,
-                max_tokens=1024,
-                stream=False
-            )
             
-            # 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
-
+            # Filter and process detections
+            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
+                    processed_boxes.append(box)
+            
+            return np.array(processed_boxes), labels
         except Exception as e:
-            print(f"Analysis error: {str(e)}")
-            return [], f"Analysis Error: {str(e)}"
+            print(f"Error in object detection: {str(e)}")
+            return np.array([]), {}
 
     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."""
+        """Improved bounding box visualization."""
         image_copy = image.copy()
         font = cv2.FONT_HERSHEY_SIMPLEX
         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)]
-                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 construction-related
+                is_relevant = any(keyword in label.lower() for keyword in construction_keywords)
                 
-                # 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', '')}"
+                if is_relevant or conf > 0.35:  # Higher threshold for non-construction objects
+                    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
+                    cv2.rectangle(image_copy, (x1, y1), (x2, y2), color, thickness)
+                    
+                    # Check for associated safety issues
+                    risk_found = False
+                    for safety_issue in safety_issues:
+                        issue_keywords = safety_issue.get('object', '').lower().split()
+                        if any(keyword in label.lower() for keyword in issue_keywords):
+                            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, (0, 0, 255), thickness)
-                        risk_found = True
-                        break
+                                  font_scale, color, thickness)
                         
-                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)
+                    # Draw additional markers for 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)), 
+                                 5, (0, 0, 255), -1)
+            
             except Exception as e:
                 print(f"Error drawing box: {str(e)}")
                 continue