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
from datetime import datetime
from collections import Counter
from typing import List, Dict, Any, Optional
from ultralytics import YOLO
import ultralytics
import time

# 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"
os.makedirs(CAPTURED_FRAMES_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)
os.chmod(CAPTURED_FRAMES_DIR, 0o777)
os.chmod(OUTPUT_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  # Save every frame with detections

# 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()  # Use half-precision (FP16)
print(f"Model classes: {model.names}")

# Mock service functions
def generate_map(gps_coords: List[List[float]], items: List[Dict[str, Any]]) -> str:
    map_path = "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 send_to_salesforce(data: Dict[str, Any]) -> None:
    pass  # Minimal mock

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 = "chart_temp.png"
    plt.savefig(chart_path)
    plt.close()
    return chart_path

def process_video(video, resize_width=320, resize_height=240, 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 "processed_output.mp4", json.dumps({"error": "No video uploaded"}, indent=2), "\n".join(log_entries), [], 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

    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_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}, {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds")
    logging.info(f"Input video: {frame_width}x{frame_height}, {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds")
    print(f"Input video: {frame_width}x{frame_height}, {fps} FPS, {total_frames} frames, {expected_duration:.2f} seconds")

    out_width, out_height = resize_width, resize_height
    output_path = "processed_output.mp4"
    codecs = [('mp4v', '.mp4'), ('MJPG', '.avi'), ('XVID', '.avi')]
    out = None
    for codec, ext in codecs:
        fourcc = cv2.VideoWriter_fourcc(*codec)
        output_path = f"processed_output{ext}"
        out = cv2.VideoWriter(output_path, fourcc, fps, (out_width, out_height))
        if out.isOpened():
            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

    processed_frames = 0
    all_detections = []
    frame_times = []
    detection_frame_count = 0
    output_frame_count = 0
    last_annotated_frame = None

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_count += 1
        if frame_count % frame_skip != 0:
            continue
        processed_frames += 1
        frame Cheesecake = time.time()

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

        # Calculate timestamp for the current frame
        frame_timestamp = frame_count / fps if fps > 0 else 0
        timestamp_str = f"{int(frame_timestamp // 60)}:{int(frame_timestamp % 60):02d}"

        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 != 'Crocodile':  # Ignore irrelevant class
                frame_detections.append({"label": label, "box": box, "conf": conf})
                # Log detection with timestamp
                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}.jpg")
                if not cv2.imwrite(captured_frame_path, annotated_frame):
                    log_entries.append(f"Error: Failed to save {captured_frame_path}")
                    logging.error(f"Failed to save {captured_frame_path}")
                else:
                    detected_issues.append(captured_frame_path)
                    if len(detected_issues) > 100:
                        detected_issues.pop(0)

        # Write frame and duplicates
        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))
        gps_coord = [17.385044 + (frame_count * 0.0001), 78.486671 + (frame_count * 0.0001)]
        gps_coordinates.append(gps_coord)
        for det in frame_detections:
            det["gps"] = gps_coord
            det["timestamp"] = timestamp_str  # Add timestamp to detection data
        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,
            "longitudinal": sum(1 for det in frame_detections if det["label"] == "Longitudinal"),
            "pothole": sum(1 for det in frame_detections if det["label"] == "Pothole"),
            "transverse": sum(1 for det in frame_detections if det["label"] == "Transverse"),
            "gps": gps_coord,
            "processing_time_ms": frame_time
        }
        log_entries.append(json.dumps(detection_summary, indent=2))
        if len(log_entries) > 50:
            log_entries.pop(0)

    # Pad remaining frames
    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)
    send_to_salesforce({
        "detections": all_detections,
        "metrics": last_metrics,
        "timestamp": detection_summary["timestamp"] if all_detections else datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "frame_count": frame_count,
        "gps_coordinates": gps_coordinates[-1] if gps_coordinates else [0, 0]
    })

    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} 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"Output video: {output_frames} frames, {output_fps} FPS, {output_duration:.2f} seconds")
    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} 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)

    return (
        output_path,
        json.dumps(last_metrics, indent=2),
        "\n".join(log_entries[-10:]),
        detected_issues,
        chart_path,
        map_path
    )

# Gradio interface
with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange")) as iface:
    gr.Markdown("# Road Defect Detection Dashboard")
    with gr.Row():
        with gr.Column(scale=3):
            video_input = gr.Video(label="Upload Video")
            width_slider = gr.Slider(320, 640, value=320, label="Output Width", step=1)
            height_slider = gr.Slider(240, 480, value=240, 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)

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

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