Update app.py

app.py

@@ -14,7 +14,7 @@ from ultralytics import YOLO
 import ultralytics
 import time
 import piexif
-import shutil
+import zipfile
 
 # Set YOLO config directory
 os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"
@@ -44,9 +44,9 @@ detected_issues: List[str] = []
 gps_coordinates: List[List[float]] = []
 last_metrics: Dict[str, Any] = {}
 frame_count: int = 0
-SAVE_IMAGE_INTERVAL = 1  # Save every frame with detections
+SAVE_IMAGE_INTERVAL = 1
 
-# Detection classes (aligned with model classes, excluding 'Crocodile')
+# Detection classes
 DETECTION_CLASSES = ["Longitudinal", "Pothole", "Transverse"]
 
 # Debug: Check environment
@@ -60,9 +60,24 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
 model = YOLO('./data/best.pt').to(device)
 if device == "cuda":
-    model.half()  # Use half-precision (FP16)
+    model.half()
 print(f"Model classes: {model.names}")
 
+def zip_directory(folder_path: str, zip_path: str) -> str:
+    """Zip all files in a directory."""
+    try:
+        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            for root, _, files in os.walk(folder_path):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    arcname = os.path.relpath(file_path, folder_path)
+                    zipf.write(file_path, arcname)
+        return zip_path
+    except Exception as e:
+        logging.error(f"Failed to zip {folder_path}: {str(e)}")
+        log_entries.append(f"Error: Failed to zip {folder_path}: {str(e)}")
+        return ""
+
 def generate_map(gps_coords: List[List[float]], items: List[Dict[str, Any]]) -> str:
     map_path = os.path.join(OUTPUT_DIR, "map_temp.png")
     plt.figure(figsize=(4, 4))
@@ -111,9 +126,9 @@ def write_flight_log(frame_count: int, gps_coord: List[float], timestamp: str) -
 def check_image_quality(frame: np.ndarray, input_resolution: int) -> bool:
     height, width, _ = frame.shape
     frame_resolution = width * height
-    if frame_resolution < 12_000_000:  # NHAI requires 12 MP
+    if frame_resolution < 12_000_000:
         log_entries.append(f"Frame {frame_count}: Resolution {width}x{height} ({frame_resolution/1e6:.2f}MP) below 12MP, non-compliant")
-    if frame_resolution < input_resolution:  # Ensure output is not below input
+    if frame_resolution < input_resolution:
         log_entries.append(f"Frame {frame_count}: Output resolution {width}x{height} below input resolution")
         return False
     return True
@@ -141,10 +156,6 @@ def generate_line_chart() -> Optional[str]:
     plt.close()
     return chart_path
 
-def generate_download_zip():
-    shutil.make_archive("outputs_bundle", 'zip', OUTPUT_DIR)
-    return "outputs_bundle.zip"
-
 def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
     global frame_count, last_metrics, detected_counts, detected_issues, gps_coordinates, log_entries
     frame_count = 0
@@ -157,14 +168,14 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
     if video is None:
         log_entries.append("Error: No video uploaded")
         logging.error("No video uploaded")
-        return "processed_output.mp4", json.dumps({"error": "No video uploaded"}, indent=2), "\n".join(log_entries), [], None, None, None
+        return None, json.dumps({"error": "No video uploaded"}, indent=2), "\n".join(log_entries), [], None, None, None, None, None, None
 
     start_time = time.time()
     cap = cv2.VideoCapture(video)
     if not cap.isOpened():
         log_entries.append("Error: Could not open video file")
         logging.error("Could not open video file")
-        return "processed_output.mp4", json.dumps({"error": "Could not open video file"}, indent=2), "\n".join(log_entries), [], None, None, None
+        return None, json.dumps({"error": "Could not open video file"}, indent=2), "\n".join(log_entries), [], None, None, None, None, None, None
 
     frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -172,36 +183,25 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
     fps = cap.get(cv2.CAP_PROP_FPS)
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     expected_duration = total_frames / fps if fps > 0 else 0
-    log_entries.append(f"Input video: {frame_width}x{frame_height} ({input_resolution/1e6:.2f}MP), {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds")
-    logging.info(f"Input video: {frame_width}x{frame_height} ({input_resolution/1e6:.2f}MP), {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds")
-    print(f"Input video: {frame_width}x{frame_height} ({input_resolution/1e6:.2f}MP), {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds")
+    log_entries.append(f"Input video: {frame_width}x{frame_height} ({input_resolution/1e6:.2f}MP), {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds, Frame skip: {frame_skip}")
+    logging.info(f"Input video: {frame_width}x{frame_height} ({input_resolution/1e6:.2f}MP), {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds, Frame skip: {frame_skip}")
+    print(f"Input video: {frame_width}x{frame_height} ({input_resolution/1e6:.2f}MP), {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds, Frame skip: {frame_skip}")
 
     out_width, out_height = resize_width, resize_height
     output_path = os.path.join(OUTPUT_DIR, "processed_output.mp4")
-    codecs = [('mp4v', '.mp4'), ('XVID', '.avi'), ('MJPG', '.avi')]  # Prioritize mp4v
-    out = None
-    for codec, ext in codecs:
-        fourcc = cv2.VideoWriter_fourcc(*codec)
-        temp_output_path = os.path.join(OUTPUT_DIR, f"processed_output{ext}")
-        out = cv2.VideoWriter(temp_output_path, fourcc, fps, (out_width, out_height))
-        if out.isOpened():
-            output_path = temp_output_path
-            log_entries.append(f"Using codec: {codec}, output: {output_path}")
-            logging.info(f"Using codec: {codec}, output: {output_path}")
-            break
-        else:
-            log_entries.append(f"Failed to initialize codec: {codec}")
-            logging.warning(f"Failed to initialize codec: {codec}")
-
-    if not out or not out.isOpened():
-        log_entries.append("Error: All codecs failed to initialize video writer")
-        logging.error("All codecs failed to initialize video writer")
+    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (out_width, out_height))
+    if not out.isOpened():
+        log_entries.append("Error: Failed to initialize mp4v codec")
+        logging.error("Failed to initialize mp4v codec")
         cap.release()
-        return "processed_output.mp4", json.dumps({"error": "All codecs failed"}, indent=2), "\n".join(log_entries), [], None, None, None
+        return None, json.dumps({"error": "mp4v codec failed"}, indent=2), "\n".join(log_entries), [], None, None, None, None, None, None
 
     processed_frames = 0
     all_detections = []
     frame_times = []
+    inference_times = []
+    resize_times = []
+    io_times = []
     detection_frame_count = 0
     output_frame_count = 0
     last_annotated_frame = None
@@ -222,13 +222,20 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
         processed_frames += 1
         frame_start = time.time()
 
+        # Resize
+        resize_start = time.time()
         frame = cv2.resize(frame, (out_width, out_height))
+        resize_times.append((time.time() - resize_start) * 1000)
+
         if not check_image_quality(frame, input_resolution):
             log_entries.append(f"Frame {frame_count}: Skipped due to low resolution")
             continue
 
+        # Inference
+        inference_start = time.time()
         results = model(frame, verbose=False, conf=0.5, iou=0.7)
         annotated_frame = results[0].plot()
+        inference_times.append((time.time() - inference_start) * 1000)
 
         frame_timestamp = frame_count / fps if fps > 0 else 0
         timestamp_str = f"{int(frame_timestamp // 60)}:{int(frame_timestamp % 60):02d}"
@@ -236,6 +243,7 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
         gps_coord = [17.385044 + (frame_count * 0.0001), 78.486671 + (frame_count * 0.0001)]
         gps_coordinates.append(gps_coord)
 
+        io_start = time.time()
         frame_detections = []
         for detection in results[0].boxes:
             cls = int(detection.cls)
@@ -282,6 +290,8 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
                 "frame": frame_count
             })
 
+        io_times.append((time.time() - io_start) * 1000)
+
         out.write(annotated_frame)
         output_frame_count += 1
         last_annotated_frame = annotated_frame
@@ -295,20 +305,15 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
 
         frame_time = (time.time() - frame_start) * 1000
         frame_times.append(frame_time)
-
-        detection_summary = {
-            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-            "video_timestamp": timestamp_str,
-            "frame": frame_count,
-            "gps": gps_coord,
-            "processing_time_ms": frame_time,
-            "detections": {label: sum(1 for det in frame_detections if det["label"] == label) for label in DETECTION_CLASSES}
-        }
-        data_lake_submission["analytics"].append(detection_summary)
-        log_entries.append(json.dumps(detection_summary, indent=2))
+        log_entries.append(f"Frame {frame_count}: Processed in {frame_time:.2f} ms (Resize: {resize_times[-1]:.2f} ms, Inference: {inference_times[-1]:.2f} ms, I/O: {io_times[-1]:.2f} ms)")
         if len(log_entries) > 50:
             log_entries.pop(0)
 
+        if time.time() - start_time > 600:
+            log_entries.append("Error: Processing timeout after 600 seconds")
+            logging.error("Processing timeout after 600 seconds")
+            break
+
     while output_frame_count < total_frames and last_annotated_frame is not None:
         out.write(last_annotated_frame)
         output_frame_count += 1
@@ -339,16 +344,23 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
 
     total_time = time.time() - start_time
     avg_frame_time = sum(frame_times) / len(frame_times) if frame_times else 0
+    avg_resize_time = sum(resize_times) / len(resize_times) if resize_times else 0
+    avg_inference_time = sum(inference_times) / len(inference_times) if inference_times else 0
+    avg_io_time = sum(io_times) / len(io_times) if io_times else 0
     log_entries.append(f"Output video: {output_frames} frames, {output_fps:.2f} FPS, {output_duration:.2f} seconds")
     logging.info(f"Output video: {output_frames} frames, {output_fps:.2f} FPS, {output_duration:.2f} seconds")
-    log_entries.append(f"Total processing time: {total_time:.2f} seconds, Avg frame time: {avg_frame_time:.2f} ms, Detection frames: {detection_frame_count}, Output frames: {output_frame_count}")
-    logging.info(f"Total processing time: {total_time:.2f} seconds, Avg frame time: {avg_frame_time:.2f} ms, Detection frames: {detection_frame_count}, Output frames: {output_frame_count}")
+    log_entries.append(f"Total processing time: {total_time:.2f} seconds, Avg frame time: {avg_frame_time:.2f} ms (Avg Resize: {avg_resize_time:.2f} ms, Avg Inference: {avg_inference_time:.2f} ms, Avg I/O: {avg_io_time:.2f} ms), Detection frames: {detection_frame_count}, Output frames: {output_frame_count}")
+    logging.info(f"Total processing time: {total_time:.2f} seconds, Avg frame time: {avg_frame_time:.2f} ms (Avg Resize: {avg_resize_time:.2f} ms, Avg Inference: {avg_inference_time:.2f} ms, Avg I/O: {avg_io_time:.2f} ms), Detection frames: {detection_frame_count}, Output frames: {output_frame_count}")
     print(f"Output video: {output_frames} frames, {output_fps:.2f} FPS, {output_duration:.2f} seconds")
     print(f"Total processing time: {total_time:.2f} seconds, Avg frame time: {avg_frame_time:.2f} ms, Detection frames: {detection_frame_count}, Output frames: {output_frame_count}")
 
     chart_path = generate_line_chart()
     map_path = generate_map(gps_coordinates[-5:], all_detections)
 
+    # Zip images and logs
+    images_zip = zip_directory(CAPTURED_FRAMES_DIR, os.path.join(OUTPUT_DIR, "captured_frames.zip"))
+    logs_zip = zip_directory(FLIGHT_LOG_DIR, os.path.join(OUTPUT_DIR, "flight_logs.zip"))
+
     return (
         output_path,
         json.dumps(last_metrics, indent=2),
@@ -356,7 +368,10 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
         detected_issues,
         chart_path,
         map_path,
-        generate_download_zip()  # Provide the zip link for all outputs
+        submission_json_path,
+        images_zip,
+        logs_zip,
+        output_path  # For video download
     )
 
 # Gradio interface
@@ -379,13 +394,30 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange")) as iface:
         map_output = gr.Image(label="Issue Locations Map")
     with gr.Row():
         logs_output = gr.Textbox(label="Logs", lines=5, interactive=False)
-        zip_download = gr.File(label="Download All Outputs (ZIP)")
+    with gr.Row():
+        gr.Markdown("## Download Results")
+    with gr.Row():
+        json_download = gr.File(label="Download Data Lake JSON")
+        images_zip_download = gr.File(label="Download Geotagged Images (ZIP)")
+        logs_zip_download = gr.File(label="Download Flight Logs (ZIP)")
+        video_download = gr.File(label="Download Processed Video")
 
     process_btn.click(
-        process_video,
+        fn=process_video,
         inputs=[video_input, width_slider, height_slider, skip_slider],
-        outputs=[video_output, metrics_output, logs_output, issue_gallery, chart_output, map_output, zip_download]
+        outputs=[
+            video_output,
+            metrics_output,
+            logs_output,
+            issue_gallery,
+            chart_output,
+            map_output,
+            json_download,
+            images_zip_download,
+            logs_zip_download,
+            video_download
+        ]
     )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()