Update app.py

app.py

@@ -1,16 +1,21 @@
-import os
-import zipfile
-import json
-import logging
 import cv2
 import torch
+import gradio as gr
 import numpy as np
+import os
+import json
+import logging
 import matplotlib.pyplot as plt
+import csv
 from datetime import datetime
 from collections import Counter
+from typing import List, Dict, Any, Optional
 from ultralytics import YOLO
-import piexif
+import ultralytics
 import time
+import piexif
+import zipfile
+import shutil
 
 # Set YOLO config directory
 os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"
@@ -29,19 +34,37 @@ FLIGHT_LOG_DIR = "flight_logs"
 os.makedirs(CAPTURED_FRAMES_DIR, exist_ok=True)
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 os.makedirs(FLIGHT_LOG_DIR, exist_ok=True)
+os.chmod(CAPTURED_FRAMES_DIR, 0o777)
+os.chmod(OUTPUT_DIR, 0o777)
+os.chmod(FLIGHT_LOG_DIR, 0o777)
 
 # Global variables
-detected_counts = []
-detected_issues = []
-gps_coordinates = []
-log_entries = []
-frame_count = 0
+log_entries: List[str] = []
+detected_counts: List[int] = []
+detected_issues: List[str] = []
+gps_coordinates: List[List[float]] = []
+last_metrics: Dict[str, Any] = {}
+frame_count: int = 0
+SAVE_IMAGE_INTERVAL = 1
+
+# Detection classes
 DETECTION_CLASSES = ["Longitudinal", "Pothole", "Transverse"]
 
+# Debug: Check environment
+print(f"Torch version: {torch.__version__}")
+print(f"Gradio version: {gr.__version__}")
+print(f"Ultralytics version: {ultralytics.__version__}")
+print(f"CUDA available: {torch.cuda.is_available()}")
+
 # Load custom YOLO model
 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()
+print(f"Model classes: {model.names}")
 
+# Function to zip all files in a directory
 def zip_directory(folder_path: str, zip_path: str) -> str:
     """Zip all files in a directory."""
     try:
@@ -54,10 +77,11 @@ def zip_directory(folder_path: str, zip_path: str) -> str:
         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 ""
 
+# Function to generate the map of detected issues' locations
 def generate_map(gps_coords: List[List[float]], items: List[Dict[str, Any]]) -> str:
-    """Generate and save map of detected issue locations."""
     map_path = os.path.join(OUTPUT_DIR, "map_temp.png")
     plt.figure(figsize=(4, 4))
     plt.scatter([x[1] for x in gps_coords], [x[0] for x in gps_coords], c='blue', label='GPS Points')
@@ -69,11 +93,13 @@ def generate_map(gps_coords: List[List[float]], items: List[Dict[str, Any]]) ->
     plt.close()
     return map_path
 
+# Function to geotag the image with GPS coordinates
 def write_geotag(image_path: str, gps_coord: List[float]) -> bool:
-    """Add GPS coordinates as EXIF data to an image."""
     try:
-        lat, lon = abs(gps_coord[0]), abs(gps_coord[1])
-        lat_ref, lon_ref = ("N" if gps_coord[0] >= 0 else "S"), ("E" if gps_coord[1] >= 0 else "W")
+        lat = abs(gps_coord[0])
+        lon = abs(gps_coord[1])
+        lat_ref = "N" if gps_coord[0] >= 0 else "S"
+        lon_ref = "E" if gps_coord[1] >= 0 else "W"
         exif_dict = piexif.load(image_path) if os.path.exists(image_path) else {"GPS": {}}
         exif_dict["GPS"] = {
             piexif.GPSIFD.GPSLatitudeRef: lat_ref,
@@ -85,98 +111,333 @@ def write_geotag(image_path: str, gps_coord: List[float]) -> bool:
         return True
     except Exception as e:
         logging.error(f"Failed to geotag {image_path}: {str(e)}")
+        log_entries.append(f"Error: Failed to geotag {image_path}: {str(e)}")
         return False
 
-def process_video(video):
-    global frame_count, detected_counts, detected_issues, gps_coordinates, log_entries
+# Function to write flight logs
+def write_flight_log(frame_count: int, gps_coord: List[float], timestamp: str) -> str:
+    log_path = os.path.join(FLIGHT_LOG_DIR, f"flight_log_{frame_count:06d}.csv")
+    try:
+        with open(log_path, 'w', newline='') as csvfile:
+            writer = csv.writer(csvfile)
+            writer.writerow(["Frame", "Timestamp", "Latitude", "Longitude", "Speed_ms", "Satellites", "Altitude_m"])
+            writer.writerow([frame_count, timestamp, gps_coord[0], gps_coord[1], 5.0, 12, 60])
+        return log_path
+    except Exception as e:
+        logging.error(f"Failed to write flight log {log_path}: {str(e)}")
+        log_entries.append(f"Error: Failed to write flight log {log_path}: {str(e)}")
+        return ""
+
+# Function to update the metrics for detections
+def update_metrics(detections: List[Dict[str, Any]]) -> Dict[str, Any]:
+    counts = Counter([det["label"] for det in detections])
+    return {
+        "items": [{"type": k, "count": v} for k, v in counts.items()],
+        "total_detections": len(detections),
+        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    }
+
+# Function to generate detection trend chart
+def generate_line_chart() -> Optional[str]:
+    if not detected_counts:
+        return None
+    plt.figure(figsize=(4, 2))
+    plt.plot(detected_counts[-50:], marker='o', color='#FF8C00')
+    plt.title("Detections Over Time")
+    plt.xlabel("Frame")
+    plt.ylabel("Count")
+    plt.grid(True)
+    plt.tight_layout()
+    chart_path = os.path.join(OUTPUT_DIR, "chart_temp.png")
+    plt.savefig(chart_path)
+    plt.close()
+    return chart_path
+
+# Function to generate a single ZIP report containing all results
+def generate_single_report(output_path, detected_issues, flight_logs, metrics, chart_path, map_path):
+    try:
+        # Create a directory for the report files
+        report_dir = os.path.join(OUTPUT_DIR, "final_report")
+        os.makedirs(report_dir, exist_ok=True)
+
+        # Copy the processed video
+        shutil.copy(output_path, os.path.join(report_dir, "processed_video.mp4"))
+        
+        # Save the metrics JSON
+        metrics_json_path = os.path.join(report_dir, "metrics.json")
+        with open(metrics_json_path, 'w') as json_file:
+            json.dump(metrics, json_file, indent=2)
+        
+        # Zip all captured frames
+        images_zip_path = zip_directory(CAPTURED_FRAMES_DIR, os.path.join(report_dir, "captured_frames.zip"))
+        
+        # Zip the flight logs
+        logs_zip_path = zip_directory(FLIGHT_LOG_DIR, os.path.join(report_dir, "flight_logs.zip"))
+        
+        # Save the detection trend chart
+        if chart_path:
+            shutil.copy(chart_path, os.path.join(report_dir, "detection_trend_chart.png"))
+        
+        # Save the issue locations map
+        if map_path:
+            shutil.copy(map_path, os.path.join(report_dir, "issue_locations_map.png"))
+        
+        # Create a ZIP of the entire report folder
+        zip_path = os.path.join(OUTPUT_DIR, "final_report.zip")
+        shutil.make_archive(zip_path.replace('.zip', ''), 'zip', report_dir)
+        
+        # Clean up the report directory after zipping
+        shutil.rmtree(report_dir)
+
+        return zip_path
+    except Exception as e:
+        logging.error(f"Error generating single report: {str(e)}")
+        log_entries.append(f"Error generating single report: {str(e)}")
+        return ""
+
+# Video processing function
+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
     detected_counts.clear()
     detected_issues.clear()
     gps_coordinates.clear()
     log_entries.clear()
+    last_metrics = {}
+
+    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, 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)
+        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))
+    input_resolution = frame_width * frame_height
     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}")
 
-    out_path = os.path.join(OUTPUT_DIR, "processed_output.mp4")
-    out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (4000, 3000))
+    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")
+        cap.release()
+        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 = []
-    data_lake_submission = {"images": [], "flight_logs": [], "metrics": {}}
+    frame_times = []
+    inference_times = []
+    resize_times = []
+    io_times = []
+    detection_frame_count = 0
+    output_frame_count = 0
+    last_annotated_frame = None
+    data_lake_submission = {
+        "images": [],
+        "flight_logs": [],
+        "analytics": [],
+        "metrics": {}
+    }
 
     while True:
         ret, frame = cap.read()
         if not ret:
             break
         frame_count += 1
-        results = model(frame)
+        if frame_count % frame_skip != 0:
+            continue
+        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}"
 
-        # Simulate GPS coordinates for each frame
         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:
-            label = model.names[int(detection.cls)]
+            cls = int(detection.cls)
+            conf = float(detection.conf)
+            box = detection.xyxy[0].cpu().numpy().astype(int).tolist()
+            label = model.names[cls]
             if label in DETECTION_CLASSES:
-                frame_detections.append({"label": label, "box": detection.xyxy[0].cpu().numpy().tolist()})
-                log_entries.append(f"Detected {label} in frame {frame_count}")
+                frame_detections.append({
+                    "label": label,
+                    "box": box,
+                    "conf": conf,
+                    "gps": gps_coord,
+                    "timestamp": timestamp_str
+                })
+                log_message = f"Frame {frame_count} at {timestamp_str}: Detected {label} with confidence {conf:.2f}"
+                log_entries.append(log_message)
+                logging.info(log_message)
 
         if frame_detections:
-            captured_frame_path = os.path.join(CAPTURED_FRAMES_DIR, f"detected_{frame_count:06d}.jpg")
-            cv2.imwrite(captured_frame_path, annotated_frame)
-            write_geotag(captured_frame_path, gps_coord)
+            detection_frame_count += 1
+            if detection_frame_count % SAVE_IMAGE_INTERVAL == 0:
+                captured_frame_path = os.path.join(CAPTURED_FRAMES_DIR, f"detected_{frame_count:06d}.jpg")
+                if cv2.imwrite(captured_frame_path, annotated_frame):
+                    if write_geotag(captured_frame_path, gps_coord):
+                        detected_issues.append(captured_frame_path)
+                        data_lake_submission["images"].append({
+                            "path": captured_frame_path,
+                            "frame": frame_count,
+                            "gps": gps_coord,
+                            "timestamp": timestamp_str
+                        })
+                        if len(detected_issues) > 100:
+                            detected_issues.pop(0)
+                    else:
+                        log_entries.append(f"Frame {frame_count}: Geotagging failed")
+                else:
+                    log_entries.append(f"Error: Failed to save {captured_frame_path}")
+                    logging.error(f"Failed to save {captured_frame_path}")
 
-            detected_issues.append(captured_frame_path)
-            data_lake_submission["images"].append({"path": captured_frame_path, "frame": frame_count, "gps": gps_coord})
+        flight_log_path = write_flight_log(frame_count, gps_coord, timestamp_str)
+        if flight_log_path:
+            data_lake_submission["flight_logs"].append({
+                "path": flight_log_path,
+                "frame": frame_count
+            })
 
-        log_path = os.path.join(FLIGHT_LOG_DIR, f"flight_log_{frame_count:06d}.csv")
-        with open(log_path, 'w', newline='') as csvfile:
-            writer = csv.writer(csvfile)
-            writer.writerow(["Frame", "Latitude", "Longitude", "Timestamp"])
-            writer.writerow([frame_count, gps_coord[0], gps_coord[1], datetime.now().strftime("%Y-%m-%d %H:%M:%S")])
-        
-        data_lake_submission["flight_logs"].append({"path": log_path, "frame": frame_count})
+        io_times.append((time.time() - io_start) * 1000)
 
         out.write(annotated_frame)
+        output_frame_count += 1
+        last_annotated_frame = annotated_frame
+        if frame_skip > 1:
+            for _ in range(frame_skip - 1):
+                out.write(annotated_frame)
+                output_frame_count += 1
+
+        detected_counts.append(len(frame_detections))
+        all_detections.extend(frame_detections)
+
+        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 (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
+
+    last_metrics = update_metrics(all_detections)
+    data_lake_submission["metrics"] = last_metrics
+    data_lake_submission["frame_count"] = frame_count
+    data_lake_submission["gps_coordinates"] = gps_coordinates[-1] if gps_coordinates else [0, 0]
+
+    submission_json_path = os.path.join(OUTPUT_DIR, "data_lake_submission.json")
+    try:
+        with open(submission_json_path, 'w') as f:
+            json.dump(data_lake_submission, f, indent=2)
+        log_entries.append(f"Submission JSON saved: {submission_json_path}")
+        logging.info(f"Submission JSON saved: {submission_json_path}")
+    except Exception as e:
+        log_entries.append(f"Error: Failed to save submission JSON: {str(e)}")
+        logging.error(f"Failed to save submission JSON: {str(e)}")
 
     cap.release()
     out.release()
 
-    # Generate the map and trend chart
-    map_path = generate_map(gps_coordinates, all_detections)
-    trend_chart_path = os.path.join(OUTPUT_DIR, "detection_trend.png")
-    plt.plot(detected_counts)
-    plt.savefig(trend_chart_path)
-    plt.close()
+    cap = cv2.VideoCapture(output_path)
+    output_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    output_fps = cap.get(cv2.CAP_PROP_FPS)
+    output_duration = output_frames / output_fps if output_fps > 0 else 0
+    cap.release()
 
-    # Compile everything into a single ZIP file
-    zip_path = os.path.join(OUTPUT_DIR, "final_report.zip")
-    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
-        zipf.write(out_path, os.path.basename(out_path))  # Add processed video
-        zipf.write(map_path, os.path.basename(map_path))  # Add map
-        zipf.write(trend_chart_path, os.path.basename(trend_chart_path))  # Add trend chart
-        zipf.write("data_lake_submission.json", "data_lake_submission.json")  # Add submission JSON
-        zipf = zip_directory(CAPTURED_FRAMES_DIR, zip_path)  # Add captured frames
-        zipf = zip_directory(FLIGHT_LOG_DIR, zip_path)  # Add flight logs
+    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 (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}")
 
-    return zip_path
+    chart_path = generate_line_chart()
+    map_path = generate_map(gps_coordinates[-5:], all_detections)
 
+    # Generate the single ZIP report
+    final_report_zip = generate_single_report(
+        output_path,
+        detected_issues,
+        flight_logs,
+        last_metrics,
+        chart_path,
+        map_path
+    )
 
-# Gradio interface (keep unchanged)
-import gradio as gr
-with gr.Blocks() as iface:
-    gr.Markdown("# Drone Analysis Report")
-    video_input = gr.Video(label="Upload Video")
-    process_btn = gr.Button("Generate Report")
-    zip_output = gr.File(label="Download Final Report (ZIP)")
+    return final_report_zip
+# Gradio interface
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange")) as iface:
+    gr.Markdown("# NHAI Road Defect Detection Dashboard")
+    with gr.Row():
+        with gr.Column(scale=3):
+            video_input = gr.Video(label="Upload Video (12MP recommended for NHAI compliance)")
+            width_slider = gr.Slider(320, 4000, value=4000, label="Output Width", step=1)
+            height_slider = gr.Slider(240, 3000, value=3000, label="Output Height", step=1)
+            skip_slider = gr.Slider(1, 10, value=5, label="Frame Skip", step=1)
+            process_btn = gr.Button("Process Video", variant="primary")
+        with gr.Column(scale=1):
+            metrics_output = gr.Textbox(label="Detection Metrics", lines=5, interactive=False)
+    with gr.Row():
+        video_output = gr.Video(label="Processed Video")
+        issue_gallery = gr.Gallery(label="Detected Issues", columns=4, height="auto", object_fit="contain")
+    with gr.Row():
+        chart_output = gr.Image(label="Detection Trend")
+        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():
+        zip_download = gr.File(label="Download Report (ZIP)")
 
-    process_btn.click(fn=process_video, inputs=[video_input], outputs=[zip_output])
+    process_btn.click(
+        fn=process_video,
+        inputs=[video_input, width_slider, height_slider, skip_slider],
+        outputs=[zip_download]
+    )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()