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 ultralytics import time import piexif import zipfile # Set YOLO config directory os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics" # Set up logging logging.basicConfig( filename="app.log", level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) # Directories CAPTURED_FRAMES_DIR = "captured_frames" OUTPUT_DIR = "outputs" 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: 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}") 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)) plt.scatter([x[1] for x in gps_coords], [x[0] for x in gps_coords], c='blue', label='GPS Points') plt.title("Issue Locations Map") plt.xlabel("Longitude") plt.ylabel("Latitude") plt.legend() plt.savefig(map_path) plt.close() return map_path def write_geotag(image_path: str, gps_coord: List[float]) -> bool: try: 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, 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)}") log_entries.append(f"Error: Failed to geotag {image_path}: {str(e)}") return False 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 "" 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: 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: log_entries.append(f"Frame {frame_count}: Output resolution {width}x{height} below input resolution") return False return True 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") } 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 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), "\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_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 = [] 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 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}" 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) 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": 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: 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}") 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 }) 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() 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() 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}") 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), "\n".join(log_entries[-10:]), detected_issues, chart_path, map_path, submission_json_path, images_zip, logs_zip, output_path # For video download ) # 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(): 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( 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, json_download, images_zip_download, logs_zip_download, video_download ] ) if __name__ == "__main__": iface.launch()