Update app.py

app.py

@@ -14,7 +14,7 @@ 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"
@@ -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_IMAGE_INTERVAL = 1  # Save every frame with detections
 
-# Detection classes
+# Detection classes (aligned with model classes, excluding 'Crocodile')
 DETECTION_CLASSES = ["Longitudinal", "Pothole", "Transverse"]
 
 # Debug: Check environment
@@ -60,24 +60,9 @@ 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()
+    model.half()  # Use half-precision (FP16)
 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))
@@ -126,9 +111,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:
+    if frame_resolution < 12_000_000:  # NHAI requires 12 MP
         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:
+    if frame_resolution < input_resolution:  # Ensure output is not below input
         log_entries.append(f"Frame {frame_count}: Output resolution {width}x{height} below input resolution")
         return False
     return True
@@ -156,6 +141,10 @@ 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
@@ -168,14 +157,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 None, json.dumps({"error": "No video uploaded"}, indent=2), "\n".join(log_entries), [], None, None, None, None, None
+        return "processed_output.mp4", json.dumps({"error": "No video uploaded"}, indent=2), "\n".join(log_entries), [], 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 None, json.dumps({"error": "Could not open video file"}, indent=2), "\n".join(log_entries), [], None, None, None, None, None
+        return "processed_output.mp4", json.dumps({"error": "Could not open video file"}, indent=2), "\n".join(log_entries), [], None, None, None
 
     frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -183,18 +172,32 @@ 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, 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}")
+    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")
 
     out_width, out_height = resize_width, resize_height
     output_path = os.path.join(OUTPUT_DIR, "processed_output.mp4")
-    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")
+    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")
         cap.release()
-        return None, json.dumps({"error": "mp4v codec failed"}, indent=2), "\n".join(log_entries), [], None, None, None, None, None
+        return "processed_output.mp4", json.dumps({"error": "All codecs failed"}, indent=2), "\n".join(log_entries), [], None, None, None
 
     processed_frames = 0
     all_detections = []
@@ -292,16 +295,20 @@ 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)
-        log_entries.append(f"Frame {frame_count}: Processed in {frame_time:.2f} ms")
+
+        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))
         if len(log_entries) > 50:
             log_entries.pop(0)
 
-        # Timeout check
-        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
@@ -342,10 +349,6 @@ 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 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),
@@ -353,9 +356,7 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
         detected_issues,
         chart_path,
         map_path,
-        submission_json_path,
-        images_zip,
-        logs_zip
+        generate_download_zip()  # Provide the zip link for all outputs
     )
 
 # Gradio interface
@@ -378,30 +379,13 @@ 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")
-        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")
+        zip_download = gr.File(label="Download All Outputs (ZIP)")
 
     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,
-            json_download,
-            images_zip_download,
-            logs_zip_download,
-            video_download
-        ]
+        outputs=[video_output, metrics_output, logs_output, issue_gallery, chart_output, map_output, zip_download]
     )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()