Update app.py

app.py

@@ -132,7 +132,7 @@ class YOLOv8Model:
     def predict(self, image):
         try:
             results = self.model(image)
-            return results  # Return full results for bounding box plotting
+            return results  # Return full results for violation detection
         except Exception as e:
             logger.error(f"Prediction error: {e}")
             raise
@@ -149,29 +149,6 @@ def preprocess_frame(frame):
         logger.error(f"Frame preprocessing error: {e}")
         raise
 
-def draw_bounding_boxes(frame, results):
-    try:
-        logger.info("Drawing bounding boxes on frame")
-        violation_detected = False
-        for result in results:
-            for box in result.boxes:
-                # Extract bounding box coordinates
-                x1, y1, x2, y2 = map(int, box.xyxy[0])
-                conf = box.conf[0]
-                cls = int(box.cls[0])
-                label = f"{result.names[cls]} {conf:.2f}"
-                # Draw bounding box and label
-                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
-                cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
-                violation_detected = True
-                logger.info(f"Violation detected: {label} at coordinates ({x1}, {y1}, {x2}, {y2})")
-        if not violation_detected:
-            logger.info("No violations detected in the frame")
-        return frame
-    except Exception as e:
-        logger.error(f"Error drawing bounding boxes: {e}")
-        raise
-
 def capture_rtsp_frames(rtsp_url: str, max_frames=10):
     try:
         logger.info(f"Attempting to connect to RTSP stream: {rtsp_url}")
@@ -368,8 +345,6 @@ def process_image(image):
         processed_frame = preprocess_frame(frame)
         start_time = time.time()
         results = yolo_model.predict(processed_frame)
-        # Skip drawing bounding boxes as per request; return the original frame
-        # annotated_frame = draw_bounding_boxes(frame.copy(), results)
         processing_time = time.time() - start_time
         logger.info(f"Image processing completed in {processing_time:.2f} seconds")
         
@@ -383,12 +358,14 @@ def process_image(image):
                     violations_detected = True
                     cls = int(box.cls[0])
                     violation_type = result.names[cls]
-                    # Temporary workaround: Simulate all violation types for testing by randomly assigning
+                    # Temporary workaround: Simulate detection of all violation types for testing
                     # TODO: Fine-tune the YOLOv8 model with custom classes (no_helmet, unsafe_distance, unauthorized_area)
                     if violation_type == "person":
-                        logger.warning("Temporary workaround: Simulating violation types for testing. Fine-tune the model with custom safety violation classes for accurate detection.")
+                        logger.warning("Temporary workaround: Simulating safety violation detection. Fine-tune the model with custom safety violation classes for accurate detection.")
+                        # Randomly assign one of the three violation types for testing
                         possible_violations = ["no_helmet", "unsafe_distance", "unauthorized_area"]
                         violation_type = random.choice(possible_violations)
+                        logger.info(f"Simulated violation type: {violation_type}")
                     # Map YOLO labels to Salesforce picklist values (used when Salesforce is enabled)
                     violation_mapping = {
                         "no_helmet": "No Helmet",
@@ -448,9 +425,7 @@ def process_rtsp_stream():
         for frame, timestamp in capture_rtsp_frames(RTSP_URL, max_frames=10):
             processed_frame = preprocess_frame(frame)
             results = yolo_model.predict(processed_frame)
-            # Skip drawing bounding boxes as per request; append the original frame
-            # annotated_frame = draw_bounding_boxes(frame.copy(), results)
-            frames.append(frame)
+            frames.append(frame)  # Store raw frame without annotations
             
             # Log violations if detected
             for result in results:
@@ -459,12 +434,14 @@ def process_rtsp_stream():
                     if conf > 0.5:  # Confidence threshold for >90% accuracy
                         cls = int(box.cls[0])
                         violation_type = result.names[cls]
-                        # Temporary workaround: Simulate all violation types for testing by randomly assigning
+                        # Temporary workaround: Simulate detection of all violation types for testing
                         # TODO: Fine-tune the YOLOv8 model with custom classes (no_helmet, unsafe_distance, unauthorized_area)
                         if violation_type == "person":
-                            logger.warning("Temporary workaround: Simulating violation types for testing. Fine-tune the model with custom safety violation classes for accurate detection.")
+                            logger.warning("Temporary workaround: Simulating safety violation detection. Fine-tune the model with custom safety violation classes for accurate detection.")
+                            # Randomly assign one of the three violation types for testing
                             possible_violations = ["no_helmet", "unsafe_distance", "unauthorized_area"]
                             violation_type = random.choice(possible_violations)
+                            logger.info(f"Simulated violation type: {violation_type}")
                         # Map YOLO labels to Salesforce picklist values (used when Salesforce is enabled)
                         violation_mapping = {
                             "no_helmet": "No Helmet",