import streamlit as st
import torch
import cv2
import tempfile
import os
from PIL import Image
import numpy as np

# Muat model YOLOv5 yang sudah dilatih
@st.cache_resource
def load_model():
    model = torch.hub.load('ultralytics/yolov11', 'custom', path='best.pt', force_reload=True)
    return model

model = load_model()

# Fungsi untuk deteksi pada gambar
def detect_image(image):
    results = model(image)
    results.render()  # Tambahkan bounding box ke gambar
    annotated_image = results.imgs[0]  # Ambil gambar yang sudah dianotasi
    return annotated_image

# Fungsi untuk deteksi pada video
def detect_video(video_path, output_path):
    cap = cv2.VideoCapture(video_path)
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        results = model(frame)
        results.render()
        annotated_frame = np.squeeze(results.render())
        out.write(annotated_frame)

    cap.release()
    out.release()

# Streamlit Interface
st.title("YOLO Object Detection")

# Opsi untuk memilih gambar atau video
option = st.radio("Pilih jenis input:", ("Gambar", "Video"))

if option == "Gambar":
    uploaded_image = st.file_uploader("Unggah gambar", type=["jpg", "jpeg", "png"])
    if uploaded_image is not None:
        image = Image.open(uploaded_image)
        st.image(image, caption="Gambar asli", use_column_width=True)
        
        # Deteksi objek
        annotated_image = detect_image(np.array(image))
        st.image(annotated_image, caption="Hasil deteksi", use_column_width=True)

elif option == "Video":
    uploaded_video = st.file_uploader("Unggah video", type=["mp4", "avi", "mov"])
    if uploaded_video is not None:
        # Simpan video sementara
        temp_video_path = tempfile.NamedTemporaryFile(delete=False).name
        with open(temp_video_path, "wb") as f:
            f.write(uploaded_video.read())
        
        # Proses video
        output_video_path = "output_video.mp4"
        st.text("Sedang memproses video...")
        detect_video(temp_video_path, output_video_path)
        st.text("Proses selesai!")

        # Tampilkan hasil video
        st.video(output_video_path)