app.py

# -*- coding: utf-8 -*-
"""
Aplikasi Gradio untuk Analisis Komparatif Deteksi Helm Keselamatan
==================================================================
Deskripsi:
Aplikasi ini memungkinkan pengguna untuk mengunggah video dari lingkungan konstruksi
dan membandingkan kinerja dua model AI (YOLOv5m dan YOLOv8m) dalam mendeteksi
helm keselamatan secara real-time.
"""

# --- 1. IMPORT LIBRARY ---
import gradio as gr
from ultralyticsplus import YOLO, render_result
import cv2
import numpy as np
import os
import tempfile
import time
import logging
import subprocess

# --- 2. SETUP & KONFIGURASI AWAL ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Set YOLO_CONFIG_DIR untuk mengatasi masalah izin
os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"

# --- 3. PEMUATAN MODEL AI ---
try:
    logger.info("Memulai pemuatan model AI...")
    model_spesialis = YOLO('keremberke/yolov8m-hard-hat-detection')
    model_spesialis.overrides['conf'] = 0.25
    model_spesialis.overrides['iou'] = 0.45
    model_spesialis.overrides['agnostic_nms'] = False
    model_spesialis.overrides['max_det'] = 1000
    logger.info("✅ Model Spesialis (YOLOv8m-HH) berhasil dimuat.")

    model_generalis = YOLO('keremberke/yolov5m-construction-safety')
    model_generalis.overrides['conf'] = 0.25
    model_generalis.overrides['iou'] = 0.45
    model_generalis.overrides['agnostic_nms'] = False
    model_generalis.overrides['max_det'] = 1000
    logger.info("✅ Model Generalis (YOLOv5m-CS) berhasil dimuat.")

    models = {
        "YOLOv8m (Spesialis Helm)": model_spesialis,
        "YOLOv5m (Generalis Konstruksi)": model_generalis
    }
    logger.info("Semua model siap digunakan.")
except Exception as e:
    logger.error(f"Gagal memuat model AI: {e}")
    raise RuntimeError(f"Tidak dapat memuat model AI. Error: {e}")

# --- 4. FUNGSI KONVERSI VIDEO ---
def convert_video_to_mp4(input_path):
    """Konversi video input ke MP4 jika format tidak didukung."""
    temp_mp4_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    ffmpeg_command = [
        "ffmpeg",
        "-y",
        "-i", input_path,
        "-c:v", "libx264",
        "-preset", "veryfast",
        "-pix_fmt", "yuv420p",
        "-t", "30",
        temp_mp4_path
    ]
    result = subprocess.run(ffmpeg_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if result.returncode == 0 and os.path.exists(temp_mp4_path) and os.path.getsize(temp_mp4_path) > 0:
        logger.info("Video dikonversi ke MP4: %s", temp_mp4_path)
        if input_path != temp_mp4_path:
            os.remove(input_path)
        return temp_mp4_path
    else:
        logger.error("Konversi video gagal: %s", result.stderr)
        return None

# --- 5. FUNGSI UTAMA PEMROSESAN VIDEO ---
def process_video_and_analyze(video_path, selected_model_name, progress=gr.Progress(track_tqdm=True)):
    """Fungsi utama untuk memproses video, deteksi objek, dan analisis."""
    if video_path is None:
        return None, "Status: Silakan unggah video terlebih dahulu.", ""

    try:
        logger.info(f"Memulai pemrosesan video: {video_path} menggunakan model: {selected_model_name}")
        start_time = time.time()

        # Pilih model
        model = models[selected_model_name]

        # Buka video dan coba konversi jika gagal
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            logger.warning("Format video tidak didukung, mencoba konversi ke MP4...")
            converted_path = convert_video_to_mp4(video_path)
            if converted_path:
                video_path = converted_path
                cap = cv2.VideoCapture(video_path)
            if not cap.isOpened():
                logger.error("Gagal membuka video setelah konversi")
                return None, "Error: Video tidak dapat dibuka meskipun dikonversi.", ""

        # Dapatkan properti video
        fps = cap.get(cv2.CAP_PROP_FPS) or 15
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        max_duration = 30
        max_frames = int(fps * max_duration)
        if total_frames > max_frames:
            logger.warning("Video lebih dari 30 detik, dipotong ke %d frame", max_frames)
            total_frames = max_frames

        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        target_width = 640
        target_height = int(height * (target_width / width)) if width > 0 else 480

        # Konfigurasi video output
        temp_output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        out = cv2.VideoWriter(temp_output_path, fourcc, fps, (target_width, target_height))
        if not out.isOpened():
            logger.error("Gagal membuka VideoWriter")
            cap.release()
            return None, "Error: Gagal membuat file video sementara.", ""

        # Variabel untuk analisis
        detection_count = 0
        helm_detected_count = 0
        frame_count = 0

        progress(0, desc="Memulai Pemrosesan...")
        while cap.isOpened() and frame_count < max_frames:
            ret, frame = cap.read()
            if not ret:
                break
            frame_count += 1

            # Update progress
            progress(frame_count / total_frames, desc=f"Memproses Frame {frame_count}/{total_frames}")

            # Resize frame
            frame_resized = cv2.resize(frame, (target_width, target_height))

            # Proses frame (skip setiap frame ke-3 untuk performa)
            if frame_count % 3 != 0:
                out.write(frame_resized)
                continue

            # Prediksi
            results = model.predict(frame_resized)
            annotated_frame = render_result(model=model, image=frame_resized, result=results[0])
            annotated_frame = np.array(annotated_frame)  # Konversi PIL ke NumPy
            annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)  # Konversi RGB ke BGR

            # Hitung deteksi
            detection_count += len(results[0].boxes)
            for box in results[0].boxes:
                class_id = int(box.cls)
                class_name = model.names[class_id].lower()
                if class_name in ['hardhat', 'helmet']:
                    helm_detected_count += 1

            out.write(annotated_frame)

        end_time = time.time()
        processing_time = end_time - start_time

        cap.release()
        out.release()

        # Verifikasi file output
        if not os.path.exists(temp_output_path) or os.path.getsize(temp_output_path) == 0:
            logger.error("File video sementara tidak valid atau kosong")
            return None, "Error: File video sementara tidak valid.", ""

        logger.info(f"Video berhasil diproses dalam {processing_time:.2f} detik.")

        # Analisis
        analysis_text = f"""
        ### Analisis Kinerja Model: {selected_model_name}
        - **Waktu Proses Total:** {processing_time:.2f} detik
        - **Total Frame Diproses:** {frame_count}
        - **Jumlah Deteksi Keseluruhan:** {detection_count} objek
        - **Jumlah Deteksi Helm:** {helm_detected_count} objek
        **Catatan:**
        - **Model Spesialis (YOLOv8m):** Fokus pada helm, akurasi tinggi untuk 'Hardhat'/'Helmet'.
        - **Model Generalis (YOLOv5m):** Deteksi berbagai objek, akurasi helm mungkin lebih rendah.
        """

        return temp_output_path, f"Status: Video berhasil diproses! ({processing_time:.2f} detik)", analysis_text

    except Exception as e:
        logger.error(f"Terjadi error saat memproses video: {e}", exc_info=True)
        return None, f"Error: Terjadi kesalahan - {e}", ""
    finally:
        if 'cap' in locals() and cap.isOpened():
            cap.release()
        if 'out' in locals() and out.isOpened():
            out.release()

# --- 6. PEMBUATAN ANTARMUKA PENGGUNA (GRADIO UI) ---
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky")) as iface:
    gr.Markdown(
        """
        # 🛡️ Analisis Komparatif Deteksi Helm Keselamatan
        Aplikasi ini membandingkan kinerja dua model AI dalam mendeteksi helm keselamatan pada video konstruksi.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### **Langkah 1: Konfigurasi & Unggah**")
            selected_model_ui = gr.Radio(
                choices=list(models.keys()),
                label="Pilih Model AI",
                value="YOLOv8m (Spesialis Helm)"
            )
            video_input = gr.Video(
                label="Unggah Video"
            )
            detect_btn = gr.Button("🚀 Mulai Deteksi", variant="primary")
            
            gr.Markdown("### **Cara Penggunaan**")
            gr.Markdown(
                """
                1. Pilih model AI.
                2. Unggah video atau gunakan contoh.
                3. Klik "Mulai Deteksi".
                4. Lihat hasil dan analisis.
                """
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### **Langkah 2: Hasil Deteksi**")
            video_output = gr.Video(label="Video Hasil Deteksi")
            status_text = gr.Textbox(label="Status Proses", interactive=False)
            gr.Markdown("### **Ringkasan Analisis**")
            analysis_output = gr.Markdown(label="Analisis Kinerja")

    # Contoh video
    gr.Examples(
        examples=[
            ["video.mp4", "YOLOv8m (Spesialis Helm)"]
        ],
        inputs=[video_input, selected_model_ui],
        outputs=[video_output, status_text, analysis_output],
        fn=process_video_and_analyze,
        cache_examples=False,
        label="Contoh Video"
    )

    # Informasi dan kredit
    with gr.Accordion("ℹ️ Informasi & Kredit", open=False):
        gr.Markdown(
            """
            ### **Identitas Prototipe**
            - **Pengembang:** Faisal Fahmi Yuliawan
            - **Tujuan:** Membandingkan deteksi helm dengan model AI.
            ### **Kredit**
            - **Model AI:** keremberke/yolov8m-hard-hat-detection, keremberke/yolov5m-construction-safety
            - **Teknologi:** Ultralytics YOLO, Gradio, OpenCV, FFmpeg
            """
        )

    detect_btn.click(
        fn=process_video_and_analyze,
        inputs=[video_input, selected_model_ui],
        outputs=[video_output, status_text, analysis_output]
    )

# --- 7. LUNCURKAN APLIKASI ---
if __name__ == "__main__":
    iface.launch(debug=True)