app.py

import cv2
import torch
import gradio as gr
import numpy as np
from ultralytics import YOLO
import time
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

# 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

# Debug: Check environment
print(f"Torch version: {torch.__version__}")
print(f"Gradio version: {gr.__version__}")
print(f"Ultralytics version: {YOLO.__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 (replace with actual implementations if available)
def generate_map(gps_coords: List[List[float]], items: List[Dict[str, Any]]) -> str:
    """Mock map generation: returns a placeholder image path."""
    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("Mock 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:
    """Mock Salesforce dispatch: logs data."""
    logging.info(f"Mock Salesforce dispatch: {json.dumps(data, indent=2)}")

def update_metrics(detections: List[Dict[str, Any]]) -> Dict[str, Any]:
    """Compute detection metrics."""
    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]:
    """Generate detection trend chart."""
    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
    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"
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (out_width, out_height))

    processed_frames = 0
    all_detections = []

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_count += 1
        if frame_count % frame_skip != 0:
            continue
        processed_frames += 1
        print(f"Processing frame {frame_count}/{total_frames}")

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

        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()  # [x_min, y_min, x_max, y_max]
            label = model.names[cls]
            frame_detections.append({"label": label, "box": box, "conf": conf})
            log_entries.append(f"Frame {frame_count}: Detected {label} with confidence {conf:.2f}")
            logging.info(f"Frame {frame_count}: Detected {label} with confidence {conf:.2f}")

        if frame_detections:
            captured_frame_path = os.path.join(CAPTURED_FRAMES_DIR, f"detected_{frame_count}.jpg")
            cv2.imwrite(captured_frame_path, annotated_frame)
            detected_issues.append(captured_frame_path)
            if len(detected_issues) > 100:
                detected_issues.pop(0)

        frame_path = os.path.join(OUTPUT_DIR, f"frame_{frame_count:04d}.jpg")
        cv2.imwrite(frame_path, annotated_frame)

        gps_coord = [17.385044 + (frame_count * 0.0001), 78.486671 + (frame_count * 0.0001)]  # Simulated GPS
        gps_coordinates.append(gps_coord)
        for det in frame_detections:
            det["gps"] = gps_coord
        all_detections.extend(frame_detections)

        out.write(annotated_frame)
        if frame_skip > 1:
            for _ in range(frame_skip - 1):
                if frame_count + 1 <= total_frames:
                    out.write(annotated_frame)
                    frame_count += 1

        detected_counts.append(len(frame_detections))
        last_metrics = update_metrics(all_detections)
        detection_summary = {
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "frame": frame_count,
            "cracks": sum(1 for det in frame_detections if det["label"] == "crack"),
            "potholes": sum(1 for det in frame_detections if det["label"] == "pothole"),
            "gps": gps_coord,
            "processing_time_ms": (time.time() - start_time) * 1000 / processed_frames if processed_frames else 0
        }
        log_entries.append(json.dumps(detection_summary, indent=2))
        logging.info(json.dumps(detection_summary, indent=2))
        if len(log_entries) > 100:
            log_entries.pop(0)

        send_to_salesforce({
            "detections": frame_detections,
            "metrics": last_metrics,
            "timestamp": detection_summary["timestamp"],
            "frame_count": frame_count,
            "gps_coordinates": gps_coord
        })

    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
    cap.release()

    log_entries.append(f"Output video: {output_frames} frames, {output_fps} FPS, {output_duration:.2f} seconds")
    logging.info(f"Output video: {output_frames} frames, {output_fps} FPS, {output_duration:.2f} seconds")
    print(f"Output video: {output_frames} frames, {output_fps} FPS, {output_duration:.2f} seconds")
    print(f"Processing time: {time.time() - start_time:.2f} seconds")

    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("# Crack and Pothole Detection Dashboard")
    with gr.Row():
        with gr.Column(scale=3):
            video_input = gr.Video(label="Upload Video")
            width_slider = gr.Slider(320, 1280, value=320, label="Output Width", step=1)
            height_slider = gr.Slider(240, 720, 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=10, interactive=False)
    with gr.Row():
        video_output = gr.Video(label="Processed Video")
        issue_gallery = gr.Gallery(label="Detected Issues", columns=4, height="auto")
    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=8, 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()