import torch
import cv2
import os
import numpy as np
import gradio as gr

# Muat model pre-trained YOLOv5
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

# Fungsi untuk memproses video dan menghitung jumlah manusia
def process_video(video_path):
    # Direktori output
    output_dir = "output_videos"
    os.makedirs(output_dir, exist_ok=True)
    output_path = os.path.join(output_dir, "person_counter_output.mp4")

    # Buka video input
    cap = cv2.VideoCapture(video_path)

    # Dapatkan spesifikasi video
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))

    # Buat VideoWriter untuk menyimpan video output
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))

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

        # Inferensi dengan YOLOv5
        results = model(frame)
        detections = results.pred[0]
        names = model.names

        # Filter hanya label 'person'
        person_detections = [d for d in detections if names[int(d[-1])] == "person"]
        person_count = len(person_detections)

        # Render frame dan buat salinan eksplisit
        annotated_frame = results.render()[0]
        annotated_frame = np.copy(annotated_frame)

        # Tambahkan teks ke frame
        cv2.putText(annotated_frame, f"Person Count: {person_count}", (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        # Tulis frame yang telah dianotasi ke video output
        out.write(annotated_frame)

    # Tutup video input dan output
    cap.release()
    out.release()

    # Mengembalikan video yang telah diproses (tidak menjumlahkan seluruh frame)
    return output_path

# Fungsi Gradio untuk antarmuka
def gradio_interface(video_file):
    output_path = process_video(video_file)
    return output_path

# Antarmuka Gradio
iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.File(type="filepath"),  # Input berupa file video
    outputs=gr.File(label="Processed Video"),
    title="Person Counter using YOLOv5",
    description="Upload a video file to detect and count the number of people in each frame using YOLOv5."
)

# Menjalankan aplikasi
if __name__ == "__main__":
    iface.launch()