app.py

"""
Vehicle Detection, Tracking & Counting Application

Main entry point untuk Streamlit app.
Menggunakan RT-DETR untuk deteksi, ByteTrack untuk tracking,
dan virtual line / polygon region untuk counting kendaraan.
"""

import streamlit as st
import cv2
import numpy as np
import tempfile
import time
from pathlib import Path

# ── Compatibility helper ──────────────────────────────────────────────────────
# `use_container_width` on st.image() was renamed from `use_column_width`.
# Streamlit 1.28.x still uses `use_column_width` for images.
import streamlit as _st_ver
_st_version = tuple(int(x) for x in _st_ver.__version__.split(".")[:3])

def _image_full_width(placeholder, frame, **kwargs):
    """Display an image spanning the full column width, version-safe."""
    if _st_version >= (1, 32, 0):
        placeholder.image(frame, use_container_width=True, **kwargs)
    else:
        placeholder.image(frame, use_column_width=True, **kwargs)

from core.detector import VehicleDetector
from core.tracker import ByteTracker
from core.counter import VirtualLineCounter, PolygonRegionCounter
from core.exporter import export_counts_to_csv, create_summary_dataframe
from ui.sidebar import render_sidebar
from utils import (
    draw_tracking,
    draw_counting_line,
    draw_polygon_region,
    draw_stats_overlay,
    resize_frame,
    calculate_fps,
    format_time
)


# konfigurasi page
st.set_page_config(
    page_title="Vehicle Detection & Counting - RT-DETR",
    page_icon="🚗",
    layout="wide",
    initial_sidebar_state="expanded"
)


def main():
    st.title("Vehicle Detection, Tracking & Counting")
    st.markdown(
        "Deteksi dan hitung kendaraan secara otomatis menggunakan "
        "**RT-DETR** + **ByteTrack**. Upload video dan lihat hasilnya."
    )

    # render sidebar, dapetin config
    config = render_sidebar()

    # auto-detect model .pt di folder models/
    models_dir = Path("models")
    available_models = sorted(models_dir.glob("*.pt"))

    if not available_models:
        st.warning(
            "Tidak ada model `.pt` ditemukan di folder `models/`. "
            "Letakkan file model (misal `rtdetr-l.pt`) di folder `models/`."
        )
        st.info(
            "Kalau belum training, jalankan notebook di `notebooks/kaggle_training.ipynb` "
            "di Kaggle terlebih dahulu."
        )
        return

    if len(available_models) == 1:
        model_path = available_models[0]
    else:
        model_names = [m.name for m in available_models]
        selected = st.selectbox("Pilih Model", model_names, index=0)
        model_path = models_dir / selected

    # inisialisasi model (cache supaya tidak load ulang terus)
    # reload jika model yang dipilih berubah
    if (
        "detector" not in st.session_state
        or st.session_state.get("loaded_model_path") != str(model_path)
    ):
        with st.spinner(f"Loading model `{model_path.name}`..."):
            st.session_state.detector = VehicleDetector(
                model_path=str(model_path),
                confidence=config["confidence"]
            )
            st.session_state.loaded_model_path = str(model_path)
    else:
        # update confidence kalau berubah
        st.session_state.detector.set_confidence(config["confidence"])

    detector = st.session_state.detector

    # tampilkan info model
    model_info = detector.get_model_info()
    with st.expander("Info Model", expanded=False):
        col1, col2 = st.columns(2)
        with col1:
            st.markdown(f"**Device:** {model_info['device']}")
            st.markdown(f"**Confidence:** {model_info['confidence_threshold']}")
        with col2:
            st.markdown(f"**Jumlah Kelas:** {model_info['num_classes']}")
            st.markdown(f"**Kelas:** {', '.join(model_info['classes'])}")

    st.markdown("---")

    # proses video kalau sudah di-upload
    if config["uploaded_video"] is not None:
        process_video(config, detector)
    else:
        st.info("Upload video di sidebar untuk memulai deteksi.")

        # tampilkan guide singkat
        st.markdown("### Cara Penggunaan")
        st.markdown("""
        1. Upload file video (.mp4 / .avi) di sidebar
        2. Atur confidence threshold sesuai kebutuhan
        3. Pilih metode counting (Virtual Line atau Polygon Region)
        4. Klik **Mulai Proses** dan tunggu sampai selesai
        5. Download hasil video dan CSV
        """)


def process_video(config, detector):
    """
    Proses video: deteksi, tracking, counting frame by frame.
    
    Args:
        config: dict dari render_sidebar()
        detector: VehicleDetector instance
    """
    uploaded_video = config["uploaded_video"]

    # simpan video sementara supaya bisa dibaca OpenCV
    tfile = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    tfile.write(uploaded_video.read())
    tfile.flush()
    video_path = tfile.name

    # buka video
    cap = cv2.VideoCapture(video_path)

    if not cap.isOpened():
        st.error("Gagal membuka video. Pastikan format video valid.")
        return

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps_video = cap.get(cv2.CAP_PROP_FPS)
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    duration = total_frames / fps_video if fps_video > 0 else 0

    # tampilkan info video
    st.subheader("Info Video")
    col1, col2, col3, col4 = st.columns(4)
    col1.metric("Resolusi", f"{frame_width}x{frame_height}")
    col2.metric("FPS", f"{fps_video:.1f}")
    col3.metric("Total Frame", str(total_frames))
    col4.metric("Durasi", format_time(duration))

    st.markdown("---")

    # tombol mulai
    start_button = st.button("Mulai Proses", type="primary")

    if not start_button:
        cap.release()
        return

    # inisialisasi tracker dan counter
    tracker = ByteTracker(
        track_thresh=config["confidence"],
        match_thresh=0.3,
        track_buffer=30
    )
    tracker.reset()

    if config["counting_mode"] == "Virtual Line":
        counter = VirtualLineCounter(
            line_position_ratio=config["line_position"],
            frame_height=frame_height
        )
    else:
        counter = PolygonRegionCounter(
            frame_width=frame_width,
            frame_height=frame_height
        )

    # setup output video writer
    output_path = "outputs/result_video.mp4"
    Path("outputs").mkdir(exist_ok=True)

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out_writer = cv2.VideoWriter(
        output_path, fourcc, fps_video, (frame_width, frame_height)
    )

    # UI elements untuk progress
    progress_bar = st.progress(0)
    status_text = st.empty()

    # area untuk menampilkan frame dan stats
    col_video, col_stats = st.columns([3, 1])

    with col_video:
        frame_display = st.empty()

    with col_stats:
        stats_display = st.empty()
        count_display = st.empty()

    # mulai processing
    frame_count = 0
    start_time = time.time()
    frame_logs = []

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        frame_count += 1

        # deteksi
        detections = detector.detect(frame)

        # tracking
        tracked_objects = tracker.update(detections)

        # counting
        counter.update(tracked_objects)
        counts = counter.get_counts()

        # hitung FPS inference
        current_fps = calculate_fps(start_time, frame_count)

        # annotate frame
        annotated = draw_tracking(frame, tracked_objects)

        if config["counting_mode"] == "Virtual Line":
            count_text = f"Total: {counts['total']}"
            annotated = draw_counting_line(
                annotated, config["line_position"], count_text
            )
        else:
            annotated = draw_polygon_region(
                annotated, counter.get_polygon_points()
            )

        # stats overlay
        stats = {
            "FPS": f"{current_fps:.1f}",
            "Frame": f"{frame_count}/{total_frames}",
            "Total": str(counts["total"])
        }
        annotated = draw_stats_overlay(annotated, stats)

        # tulis ke output video
        out_writer.write(annotated)

        # update display (tidak setiap frame, biar tidak terlalu lambat)
        if frame_count % 3 == 0 or frame_count == total_frames:
            # convert BGR ke RGB untuk Streamlit
            display_frame = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
            display_frame = resize_frame(display_frame, max_width=800)

            _image_full_width(frame_display, display_frame, channels="RGB")

            # update stats
            elapsed = time.time() - start_time
            per_class = counts.get("per_class", {})

            # build per-class stats dynamically
            class_lines = ""
            for cls_name, cls_count in sorted(per_class.items()):
                class_lines += f"            - {cls_name}: {cls_count}\n"

            stats_md = f"""
            **Processing Stats**
            
            - FPS: {current_fps:.1f}
            - Frame: {frame_count}/{total_frames}
            - Waktu: {format_time(elapsed)}
            
            **Counting Results**
            
            - Total: **{counts['total']}**
{class_lines}            """
            stats_display.markdown(stats_md)

        # update progress bar
        progress = frame_count / total_frames
        progress_bar.progress(progress)
        status_text.text(f"Processing frame {frame_count}/{total_frames}...")

        # log per frame
        frame_logs.append({
            "frame_number": frame_count,
            "num_detections": len(detections),
            "num_tracked": len(tracked_objects),
            "cumulative_count": counts["total"],
            "fps": round(current_fps, 2)
        })

    # selesai
    cap.release()
    out_writer.release()

    total_time = time.time() - start_time
    status_text.text(f"Selesai. Total waktu: {format_time(total_time)}")
    progress_bar.progress(1.0)

    st.success(f"Processing selesai. {counts['total']} kendaraan terdeteksi.")

    st.markdown("---")

    # tampilkan hasil akhir
    st.subheader("Hasil Akhir")

    # tabel counting
    summary_df = create_summary_dataframe(counts)
    st.dataframe(summary_df, use_container_width=True)

    # metrics - tampilkan per kelas secara dinamis
    per_class = counts.get("per_class", {})
    class_names = sorted(per_class.keys())

    if class_names:
        cols = st.columns(len(class_names))
        for i, cls_name in enumerate(class_names):
            cols[i].metric(cls_name, per_class[cls_name])

    st.markdown("---")

    # download buttons
    st.subheader("Download Hasil")

    col_dl1, col_dl2 = st.columns(2)

    # download video
    with col_dl1:
        if Path(output_path).exists():
            with open(output_path, "rb") as f:
                st.download_button(
                    label="Download Video Hasil",
                    data=f,
                    file_name="vehicle_counting_result.mp4",
                    mime="video/mp4"
                )

    # download CSV
    with col_dl2:
        csv_path = export_counts_to_csv(counts)
        if Path(csv_path).exists():
            with open(csv_path, "rb") as f:
                st.download_button(
                    label="Download CSV Counting",
                    data=f,
                    file_name="counting_results.csv",
                    mime="text/csv"
                )


if __name__ == "__main__":
    main()