Update app.py

app.py

@@ -6,7 +6,7 @@ import os
 import json
 import logging
 import matplotlib.pyplot as plt
-import csv  # Added for flight logs
+import csv
 from datetime import datetime
 from collections import Counter
 from typing import List, Dict, Any, Optional
@@ -14,6 +14,8 @@ from ultralytics import YOLO
 import ultralytics
 import time
 import piexif
+import zipfile
+import shutil
 
 # Set YOLO config directory
 os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"
@@ -43,9 +45,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
@@ -59,9 +61,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))
@@ -110,9 +127,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
@@ -149,17 +166,24 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
     log_entries.clear()
     last_metrics = {}
 
+    # Clear previous outputs
+    for dir in [CAPTURED_FRAMES_DIR, FLIGHT_LOG_DIR, OUTPUT_DIR]:
+        if os.path.exists(dir):
+            shutil.rmtree(dir)
+        os.makedirs(dir, exist_ok=True)
+        os.chmod(dir, 0o777)
+
     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
+        return "processed_output.mp4", 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
+        return "processed_output.mp4", 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))
@@ -173,7 +197,7 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
 
     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
+    codecs = [('mp4v', '.mp4'), ('XVID', '.avi'), ('MJPG', '.avi')]
     out = None
     for codec, ext in codecs:
         fourcc = cv2.VideoWriter_fourcc(*codec)
@@ -192,7 +216,7 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
         log_entries.append("Error: All codecs failed to initialize video writer")
         logging.error("All codecs failed to initialize video writer")
         cap.release()
-        return "processed_output.mp4", json.dumps({"error": "All codecs failed"}, indent=2), "\n".join(log_entries), [], None, None
+        return "processed_output.mp4", json.dumps({"error": "All codecs failed"}, indent=2), "\n".join(log_entries), [], None, None, None, None, None, None
 
     processed_frames = 0
     all_detections = []
@@ -344,13 +368,21 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
     chart_path = generate_line_chart()
     map_path = generate_map(gps_coordinates[-5:], all_detections)
 
+    # Zip captured_frames and flight_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),
         "\n".join(log_entries[-10:]),
         detected_issues,
         chart_path,
-        map_path
+        map_path,
+        submission_json_path,
+        images_zip,
+        logs_zip,
+        output_path  # Duplicate for video download
     )
 
 # Gradio interface
@@ -373,11 +405,70 @@ 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)
+    with gr.Row():
+        gr.Markdown("## Download Results")
+    with gr.Row():
+        json_download = gr.File(label="Download Data Lake JSON", visible=False)
+        images_zip_download = gr.File(label="Download Geotagged Images (ZIP)", visible=False)
+        logs_zip_download = gr.File(label="Download Flight Logs (ZIP)", visible=False)
+        video_download = gr.File(label="Download Processed Video", visible=False)
+
+    def update_download_buttons(*outputs):
+        video_path, metrics, logs, issues, chart, map, json_path, images_zip, logs_zip, video_file = outputs
+        return (
+            video_path,
+            metrics,
+            logs,
+            issues,
+            chart,
+            map,
+            gr.File.update(value=json_path, visible=True),
+            gr.File.update(value=images_zip, visible=True),
+            gr.File.update(value=logs_zip, visible=True),
+            gr.File.update(value=video_file, visible=True)
+        )
 
     process_btn.click(
         process_video,
         inputs=[video_input, width_slider, height_slider, skip_slider],
-        outputs=[video_output, metrics_output, logs_output, issue_gallery, chart_output, map_output]
+        outputs=[
+            video_output,
+            metrics_output,
+            logs_output,
+            issue_gallery,
+            chart_output,
+            map_output,
+            json_download,
+            images_zip_download,
+            logs_zip_download,
+            video_download
+        ]
+    ).then(
+        update_download_buttons,
+        inputs=[
+            video_output,
+            metrics_output,
+            logs_output,
+            issue_gallery,
+            chart_output,
+            map_output,
+            json_download,
+            images_zip_download,
+            logs_zip_download,
+            video_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__":