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
import shutil

# 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 = {}

    # Clear previous outputs
    for dir in [CAPTURED_FRAMES_DIR, FLIGHT_LOG_DIR, OUTPUT_DIR]:
        if os.path.exists(dir):
            shutil.rmtree(dir)
        os.makedirs(dir, exist_ok=True)
        os.chmod(dir, 0o777)

    if video is None:
        log_entries.append("Error: No video uploaded")
        logging.error("No video uploaded")
        return "processed_output.mp4", 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 "processed_output.mp4", 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")
    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")
    codecs = [('mp4v', '.mp4'), ('XVID', '.avi'), ('MJPG', '.avi')]
    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 "processed_output.mp4", json.dumps({"error": "All codecs failed"}, indent=2), "\n".join(log_entries), [], None, None, None, None, None, None

    processed_frames = 0
    all_detections = []
    frame_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()

        frame = cv2.resize(frame, (out_width, out_height))
        if not check_image_quality(frame, input_resolution):
            log_entries.append(f"Frame {frame_count}: Skipped due to low resolution")
            continue

        results = model(frame, verbose=False, conf=0.5, iou=0.7)
        annotated_frame = results[0].plot()

        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)

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

        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)

        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)

    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
    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, 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, 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 captured_frames and flight_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  # Duplicate 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", visible=False)
        images_zip_download = gr.File(label="Download Geotagged Images (ZIP)", visible=False)
        logs_zip_download = gr.File(label="Download Flight Logs (ZIP)", visible=False)
        video_download = gr.File(label="Download Processed Video", visible=False)

    def update_download_buttons(*outputs):
        video_path, metrics, logs, issues, chart, map, json_path, images_zip, logs_zip, video_file = outputs
        return (
            video_path,
            metrics,
            logs,
            issues,
            chart,
            map,
            gr.File.update(value=json_path, visible=True),
            gr.File.update(value=images_zip, visible=True),
            gr.File.update(value=logs_zip, visible=True),
            gr.File.update(value=video_file, visible=True)
        )

    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
        ]
    ).then(
        update_download_buttons,
        inputs=[
            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,
            json_download,
            images_zip_download,
            logs_zip_download,
            video_download
        ]
    )

if __name__ == "__main__":
    iface.launch()