Update app.py

app.py

@@ -209,14 +209,14 @@ CONFIG = {
     "WORKER_TRACKING_DURATION": 5.0,
     "MAX_PROCESSING_TIME": 60,
     "FRAME_SKIP": 1,
-    "BATCH_SIZE": 4,  # Reduced to 4 to lower memory usage
+    "BATCH_SIZE": 4,
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
     "TRACK_BUFFER": 90,
     "TRACK_THRESH": 0.3,
     "MATCH_THRESH": 0.5,
     "SNAPSHOT_QUALITY": 95,
     "MAX_WORKER_DISTANCE": 300,
-    "TARGET_RESOLUTION": (640, 360)  # Added to resize frames
+    "MODEL_INPUT_SIZE": (640, 640)  # Updated to match YOLO input requirements
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -235,7 +235,6 @@ def load_model():
                 torch.hub.download_url_to_file('https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt', model_path)
                 
         model = YOLO(model_path).to(device)
-        # Enable FP16 inference if on GPU
         if device.type == "cuda":
             model.model.half()
         logger.info(f"Model classes: {model.names}")
@@ -247,12 +246,56 @@ def load_model():
 model = load_model()
 
 # ========================== # Helper Functions # ==========================
-def preprocess_frame(frame):
-    # Resize frame to target resolution
-    target_res = CONFIG["TARGET_RESOLUTION"]
-    frame = cv2.resize(frame, target_res, interpolation=cv2.INTER_LINEAR)
-    frame = cv2.convertScaleAbs(frame, alpha=1.2, beta=20)
-    return frame
+def preprocess_frame(frame, original_shape):
+    # Resize while preserving aspect ratio, then pad to MODEL_INPUT_SIZE (640x640)
+    target_size = CONFIG["MODEL_INPUT_SIZE"]  # (640, 640)
+    h, w = frame.shape[:2]
+    scale = min(target_size[0] / w, target_size[1] / h)
+    new_w, new_h = int(w * scale), int(h * scale)
+    
+    # Resize the frame
+    frame_resized = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
+    
+    # Create a new 640x640 image with padding
+    padded_frame = np.zeros((target_size[1], target_size[0], 3), dtype=np.uint8)
+    top = (target_size[1] - new_h) // 2
+    left = (target_size[0] - new_w) // 2
+    padded_frame[top:top+new_h, left:left+new_w] = frame_resized
+    
+    # Apply contrast adjustment
+    padded_frame = cv2.convertScaleAbs(padded_frame, alpha=1.2, beta=20)
+    
+    # Store padding info to adjust bounding boxes later
+    padding_info = {
+        "scale": scale,
+        "top": top,
+        "left": left,
+        "original_shape": original_shape
+    }
+    
+    return padded_frame, padding_info
+
+def adjust_bbox(bbox, padding_info):
+    # Adjust bounding box coordinates from padded 640x640 space back to original frame space
+    scale = padding_info["scale"]
+    top = padding_info["top"]
+    left = padding_info["left"]
+    
+    x, y, w, h = bbox
+    # Remove padding offset and scale back
+    x = (x - left) / scale
+    y = (y - top) / scale
+    w = w / scale
+    h = h / scale
+    
+    # Ensure coordinates are within original frame bounds
+    orig_h, orig_w = padding_info["original_shape"][:2]
+    x = max(0, min(x, orig_w))
+    y = max(0, min(y, orig_h))
+    w = max(0, min(w, orig_w - x))
+    h = max(0, min(h, orig_h - y))
+    
+    return [x, y, w, h]
 
 def draw_detections(frame, detections):
     result_frame = frame.copy()
@@ -453,7 +496,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
         
         try:
             record = sf.Safety_Video_Report__c.create(record_data)
-            logger.info(f"Created Safety_Video_Report__c record: {record['id']}")
+            logger.info(f"Created Safety_Violation_Report__c record: {record['id']}")
         except Exception as e:
             logger.error(f"Failed to create Safety_Video_Report__c: {e}")
             record = sf.Account.create({"Name": f"Safety_Report_{int(time.time())}"})
@@ -558,6 +601,8 @@ def process_video(video_data, temp_dir):
         while processed_frames < total_frames:
             batch_frames = []
             batch_indices = []
+            batch_padding_info = []
+            batch_original_frames = []
             
             for _ in range(CONFIG["BATCH_SIZE"]):
                 frame_idx = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
@@ -569,7 +614,9 @@ def process_video(video_data, temp_dir):
                     logger.warning(f"Failed to read frame {frame_idx}. Skipping.")
                     break
                 
-                frame = preprocess_frame(frame)
+                # Keep a copy of the original frame for drawing detections
+                original_frame = frame.copy()
+                frame, padding_info = preprocess_frame(frame, original_shape=frame.shape)
                 
                 for _ in range(frame_skip - 1):
                     if not cap.grab():
@@ -577,6 +624,8 @@ def process_video(video_data, temp_dir):
                 
                 batch_frames.append(frame)
                 batch_indices.append(frame_idx)
+                batch_padding_info.append(padding_info)
+                batch_original_frames.append(original_frame)
                 processed_frames += 1
 
             if not batch_frames:
@@ -595,15 +644,14 @@ def process_video(video_data, temp_dir):
                 logger.error(f"Model inference failed: {e}")
                 raise ValueError(f"Failed to process video frames with YOLO model: {str(e)}")
             finally:
-                # Clear memory
-                batch_frames = []  # Clear the list to free memory
+                batch_frames = []
                 if device.type == "cuda":
                     torch.cuda.empty_cache()
 
-            for i, (result, frame_idx) in enumerate(zip(results, batch_indices)):
+            for i, (result, frame_idx, padding_info, original_frame) in enumerate(zip(results, batch_indices, batch_padding_info, batch_original_frames)):
                 current_time = frame_idx / fps
                 
-                if time.time() - start_time > 0.5:  # Yield more frequently
+                if time.time() - start_time > 0.5:
                     progress = (processed_frames / total_frames) * 100
                     elapsed_time = time.time() - start_time
                     fps_processed = processed_frames / elapsed_time if elapsed_time > 0 else 0
@@ -625,6 +673,8 @@ def process_video(video_data, temp_dir):
                         continue
 
                     bbox = box.xywh.cpu().numpy()[0]
+                    # Adjust bounding box coordinates to original frame space
+                    bbox = adjust_bbox(bbox, padding_info)
                     track_inputs.append({
                         "bbox": bbox,
                         "conf": conf,
@@ -669,8 +719,8 @@ def process_video(video_data, temp_dir):
                             "timestamp": current_time
                         }
                         
-                        snapshot_frame = batch_frames_tensor[i].cpu().numpy().transpose(1, 2, 0) * 255
-                        snapshot_frame = snapshot_frame.astype(np.uint8)
+                        # Use the original frame for drawing detections
+                        snapshot_frame = original_frame.copy()
                         snapshot_frame = draw_detections(snapshot_frame, [detection])
                         
                         cv2.putText(
@@ -703,9 +753,8 @@ def process_video(video_data, temp_dir):
                         
                         logger.info(f"Captured snapshot for {label} violation by worker {worker_id} at {current_time:.2f}s")
 
-                # Clear snapshots periodically to reduce memory usage
                 if len(snapshots) > 100:
-                    snapshots = snapshots[-10:]  # Keep only the last 10 snapshots in memory
+                    snapshots = snapshots[-10:]
 
         cap.release()
         processing_time = time.time() - start_time