app.py

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
from fpdf import FPDF  # PDF generation

# 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)

# Force CPU mode for Hugging Face Spaces (as GPUs are not available by default)
device = 'cpu'  
logging.info(f"Using device: {device}")

# Load custom YOLO model
model = YOLO('./data/best.pt').to(device)
model.float()  # Ensure the model is using full precision (on CPU)

# Global variables
log_entries: List[str] = []
detected_counts: List[int] = []
detected_issues: List[str] = []
gps_coordinates: List[List[float]] = []
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()}")

# Check image resolution quality
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:  # Minimum resolution of 12MP
        log_entries.append(f"Frame {frame_count}: Resolution {width}x{height} ({frame_resolution / 1e6:.2f}MP) below 12MP, non-compliant")
        return False
    if frame_resolution < input_resolution:
        log_entries.append(f"Frame {frame_count}: Output resolution {width}x{height} below input resolution")
        return False
    return True

# Write GPS geotag to the image
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:
        log_entries.append(f"Error in geotagging {image_path}: {str(e)}")
        logging.error(f"Failed to geotag {image_path}: {str(e)}")
        return False

# Write flight log data into a CSV file
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 ""

# Generate PDF report with Top 5 or Top 10 Images
def generate_pdf_report(log_entries, detected_issues, chart_path, map_path, metrics, top_images):
    pdf = FPDF()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.add_page()

    # Add Title
    pdf.set_font("Arial", "B", 16)
    pdf.cell(200, 10, txt="Road Defect Detection Report", ln=True, align="C")
    pdf.ln(10)

    # Add Log Entries
    pdf.set_font("Arial", size=12)
    pdf.cell(200, 10, txt="Log Entries:", ln=True)
    pdf.multi_cell(0, 10, txt="\n".join(log_entries))
    pdf.ln(5)

    # Add Detected Issues
    pdf.cell(200, 10, txt="Detected Issues:", ln=True)
    for issue in detected_issues:
        pdf.cell(200, 10, txt=issue, ln=True)
    pdf.ln(5)

    # Add Metrics
    pdf.cell(200, 10, txt="Detection Metrics:", ln=True)
    pdf.multi_cell(0, 10, txt=json.dumps(metrics, indent=2))
    pdf.ln(5)

    # Add Top 5 or Top 10 Images
    pdf.cell(200, 10, txt="Top 5 Detected Images:", ln=True)
    for image_path in top_images:
        if os.path.exists(image_path):
            pdf.image(image_path, x=10, y=pdf.get_y(), w=180)
            pdf.ln(60)  # Space out the images

    # Add Chart Image
    if chart_path:
        pdf.cell(200, 10, txt="Detection Trend Chart:", ln=True)
        pdf.image(chart_path, x=10, y=pdf.get_y(), w=180)
        pdf.ln(80)

    # Add Map Image
    if map_path:
        pdf.cell(200, 10, txt="Issue Locations Map:", ln=True)
        pdf.image(map_path, x=10, y=pdf.get_y(), w=180)
        pdf.ln(80)

    # Save PDF
    pdf_output_path = os.path.join(OUTPUT_DIR, "detection_report.pdf")
    pdf.output(pdf_output_path)
    return pdf_output_path

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 process_video(video, resize_width=4000, resize_height=3000, frame_skip=5, top_n=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()
    
    # Gradio uploads the video as a path string, so we need to directly use that path
    video_path = video.name if hasattr(video, 'name') else video  # Check if it's a file-like object

    cap = cv2.VideoCapture(video_path)  # Access the uploaded video file correctly
    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": {}
    }

    top_images = []  # Track top 5 images

    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 len(top_images) < top_n:  # Limit the number of images to top_n
                        top_images.append(captured_frame_path)
                    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"))

    # Generate PDF report with top images
    pdf_path = generate_pdf_report(log_entries, detected_issues, chart_path, map_path, last_metrics, top_images)

    return (
        output_path,
        json.dumps(last_metrics, indent=2),
        "\n".join(log_entries[-10:]),
        detected_issues,
        chart_path,
        map_path,
        pdf_path,  # PDF report 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()