Update app.py

app.py

@@ -1,16 +1,16 @@
-import cv2
-import torch
-import gradio as gr
-import numpy as np
 import os
+import zipfile
 import json
 import logging
+import cv2
+import torch
+import numpy as np
 import matplotlib.pyplot as plt
-import csv
 from datetime import datetime
 from collections import Counter
-import zipfile
-from jinja2 import Template
+from ultralytics import YOLO
+import piexif
+import time
 
 # Set YOLO config directory
 os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"
@@ -29,35 +29,19 @@ 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
-log_entries = []
 detected_counts = []
 detected_issues = []
 gps_coordinates = []
-last_metrics = {}
+log_entries = []
 frame_count = 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__}")
-
 # Load custom YOLO model
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {device}")
-model = torch.hub.load("ultralytics/yolov5", "custom", path='./data/best.pt').to(device)
-if device == "cuda":
-    model.half()
-print(f"Model classes: {model.names}")
+model = YOLO('./data/best.pt').to(device)
 
-# Helper functions for video processing, geotagging, flight logs, and quality checks
 def zip_directory(folder_path: str, zip_path: str) -> str:
     """Zip all files in a directory."""
     try:
@@ -72,7 +56,8 @@ def zip_directory(folder_path: str, zip_path: str) -> str:
         logging.error(f"Failed to zip {folder_path}: {str(e)}")
         return ""
 
-def generate_map(gps_coords: list, items: list) -> str:
+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')
@@ -84,50 +69,25 @@ def generate_map(gps_coords: list, items: list) -> str:
     plt.close()
     return map_path
 
-def write_geotag(image_path: str, gps_coord: list) -> bool:
+def write_geotag(image_path: str, gps_coord: List[float]) -> bool:
+    """Add GPS coordinates as EXIF data to an image."""
     try:
-        # Geotagging logic
-        # (code to add EXIF data here...)
+        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")
+        exif_dict = piexif.load(image_path) if os.path.exists(image_path) else {"GPS": {}}
+        exif_dict["GPS"] = {
+            piexif.GPSIFD.GPSLatitudeRef: lat_ref,
+            piexif.GPSIFD.GPSLatitude: ((int(lat), 1), (0, 1), (0, 1)),
+            piexif.GPSIFD.GPSLongitudeRef: lon_ref,
+            piexif.GPSIFD.GPSLongitude: ((int(lon), 1), (0, 1), (0, 1))
+        }
+        piexif.insert(piexif.dump(exif_dict), image_path)
         return True
     except Exception as e:
         logging.error(f"Failed to geotag {image_path}: {str(e)}")
         return False
 
-def write_flight_log(frame_count: int, gps_coord: list, 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)}")
-        return ""
-
-# Generate HTML report using Jinja2 template
-def generate_report(detections, video_path, issue_images, flight_logs, chart_path, map_path, submission_json):
-    with open("report_template.html", "r") as file:
-        template = Template(file.read())
-    
-    report_content = template.render(
-        detections=detections,
-        video_path=video_path,
-        issue_images=issue_images,
-        flight_logs=flight_logs,
-        chart_path=chart_path,
-        map_path=map_path,
-        submission_json=submission_json
-    )
-    
-    report_path = "output_report.html"
-    with open(report_path, "w") as report_file:
-        report_file.write(report_content)
-    
-    return report_path
-
-# Function to process video and generate outputs
-def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
+def process_video(video):
     global frame_count, detected_counts, detected_issues, gps_coordinates, log_entries
     frame_count = 0
     detected_counts.clear()
@@ -135,91 +95,88 @@ def process_video(video, resize_width=4000, resize_height=3000, frame_skip=5):
     gps_coordinates.clear()
     log_entries.clear()
 
-    if video is None:
-        log_entries.append("Error: No video uploaded")
-        return None, json.dumps({"error": "No video uploaded"}, indent=2)
-
     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)
 
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    out_path = os.path.join(OUTPUT_DIR, "processed_output.mp4")
+    out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (4000, 3000))
+
+    all_detections = []
+    data_lake_submission = {"images": [], "flight_logs": [], "metrics": {}}
+
     while True:
         ret, frame = cap.read()
         if not ret:
             break
         frame_count += 1
-        if frame_count % frame_skip != 0:
-            continue
+        results = model(frame)
+        annotated_frame = results[0].plot()
+
+        # 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)
+
+        frame_detections = []
+        for detection in results[0].boxes:
+            label = model.names[int(detection.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}")
+
+        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)
+
+            detected_issues.append(captured_frame_path)
+            data_lake_submission["images"].append({"path": captured_frame_path, "frame": frame_count, "gps": gps_coord})
 
-        # Process frame, detect issues, and log results
-        # (process frames and update logs here...)
+        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")])
         
-        # Add detected issues, save images, etc.
-
-    # Generate detection trend chart
-    chart_path = generate_line_chart()
-
-    # Generate map of GPS coordinates
-    map_path = generate_map(gps_coordinates, detected_issues)
-
-    # Prepare data for submission to Data Lake
-    data_lake_submission = {
-        "images": detected_issues,
-        "flight_logs": [],
-        "analytics": detections,
-        "metrics": last_metrics
-    }
-    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)
-    except Exception as e:
-        log_entries.append(f"Error: Failed to save submission JSON: {str(e)}")
-
-    # Zip files for download
-    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"))
-
-    # Generate final report
-    report_path = generate_report(
-        detections=detections,
-        video_path="processed_video.mp4",
-        issue_images=detected_issues,
-        flight_logs=logs_zip,
-        chart_path=chart_path,
-        map_path=map_path,
-        submission_json=submission_json_path
-    )
-
-    # Create the final zip file containing all report components
-    final_zip = zipfile.ZipFile("final_report.zip", 'w', zipfile.ZIP_DEFLATED)
-    final_zip.write(report_path)
-    final_zip.write(images_zip)
-    final_zip.write(logs_zip)
-    final_zip.write("processed_video.mp4")
-    final_zip.close()
-
-    return final_zip
-
-# Gradio Interface
+        data_lake_submission["flight_logs"].append({"path": log_path, "frame": frame_count})
+
+        out.write(annotated_frame)
+
+    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()
+
+    # 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
+
+    return zip_path
+
+
+# Gradio interface (keep unchanged)
+import gradio as gr
 with gr.Blocks() as iface:
-    gr.Markdown("# NHAI Road Defect Detection Dashboard")
-    with gr.Row():
-        with gr.Column(scale=3):
-            video_input = gr.Video(label="Upload Video")
-            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)
-
-    process_btn.click(
-        fn=process_video,
-        inputs=[video_input, width_slider, height_slider, skip_slider],
-        outputs=[metrics_output]
-    )
-
-if __name__ == "__main__":
-    iface.launch()
+    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)")
+
+    process_btn.click(fn=process_video, inputs=[video_input], outputs=[zip_output])
+
+iface.launch()