import os
import tempfile
import cv2
import numpy as np
import gradio as gr
import supervision as sv
from ultralytics import YOLO

# 1. Cấu hình Model
MODEL_PATH = "best.pt" # Thay bằng model đã huấn luyện
model = YOLO(MODEL_PATH)

CLASS_NAMES_DICT = model.names
SELECTED_CLASS_NAMES = list(CLASS_NAMES_DICT.values())

# Khởi tạo Annotators
box_annotator = sv.BoxAnnotator(thickness=2)
label_annotator = sv.LabelAnnotator(text_thickness=1, text_scale=0.8)

def process_image(image):
    """Xử lý ảnh đơn lẻ."""
    if image is None:
        return None
    
    results = model(image, conf=0.25, verbose=False)[0]
    detections = sv.Detections.from_ultralytics(results)
    
    labels = [
        f"{CLASS_NAMES_DICT[class_id]} {conf:0.2f}"
        for conf, class_id in zip(detections.confidence, detections.class_id)
    ]
    
    annotated_image = image.copy()
    annotated_image = box_annotator.annotate(scene=annotated_image, detections=detections)
    annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections, labels=labels)
    
    return annotated_image

def process_video(video_path):
    """Xử lý video."""
    if video_path is None: return None
    
    output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    
    def callback(frame: np.ndarray, index: int) -> np.ndarray:
        results = model(frame, conf=0.25, verbose=False)[0]
        detections = sv.Detections.from_ultralytics(results)
        
        labels = [
            f"{CLASS_NAMES_DICT[class_id]} {conf:0.2f}"
            for conf, class_id in zip(detections.confidence, detections.class_id)
        ]
        
        annotated_frame = frame.copy()
        annotated_frame = box_annotator.annotate(scene=annotated_frame, detections=detections)
        annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels)
        
        return annotated_frame

    sv.process_video(source_path=video_path, target_path=output_path, callback=callback)
    return output_path

# --- Giao diện Gradio ---
with gr.Blocks(title="Road Damage", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚧 Road Damage Detection")
    gr.Markdown("Hệ thống nhận diện và đánh giá xuống cấp mặt đường.")

    with gr.Tabs():
        # Tab Xử lý Ảnh
        with gr.TabItem("🖼️ Xử lý Ảnh"):
            with gr.Row():
                img_input = gr.Image(label="Tải ảnh lên")
                img_output = gr.Image(label="Kết quả phân tích")
            img_btn = gr.Button("Phân tích Ảnh", variant="primary")
            img_btn.click(fn=process_image, inputs=img_input, outputs=img_output)

        # Tab Xử lý Video
        with gr.TabItem("🎥 Xử lý Video"):
            with gr.Row():
                vid_input = gr.Video(label="Tải video lên")
                vid_output = gr.Video(label="Video kết quả")
            vid_btn = gr.Button("Phân tích Video", variant="primary")
            vid_btn.click(fn=process_video, inputs=vid_input, outputs=vid_output)

    gr.Markdown("---")
    gr.Markdown("""
    ### Thành viên:
    - An Hoàng Anh - 223332813
    - Trần Hải Nam - 223332840
    - Đoàn Minh Thành - 223332848
    - Nguyễn Công Thành - 223332850
    """)

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