import gradio as gr
from roboflow import Roboflow
import tempfile
import os
import cv2
import supervision as sv
import numpy as np

# Inisialisasi Roboflow
rf = Roboflow(api_key="Otg64Ra6wNOgDyjuhMYU")
project = rf.workspace("alat-pelindung-diri").project("nescafe-4base")
model = project.version(16).model

# Fungsi untuk deteksi objek dengan supervision InferenceSlicer
def detect_objects(image):
    # Simpan gambar sementara
    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
        image.save(temp_file, format="JPEG")
        temp_file_path = temp_file.name

    # Membaca gambar dengan OpenCV
    img = cv2.imread(temp_file_path)

    # Callback function untuk model prediksi
    def callback(image_slice: np.ndarray) -> sv.Detections:
        # Lakukan inferensi pada setiap potongan gambar
        predictions = model.predict(image_slice, confidence=50, overlap=30).json()
        return sv.Detections.from_inference(predictions)

    # Menggunakan InferenceSlicer
    slicer = sv.InferenceSlicer(callback=callback)

    # Proses gambar dengan slicer
    detections = slicer(img)

    # Filter deteksi yang tumpang tindih (gunakan NMM atau tanpa filter)
    filtered_detections = detections.filter(strategy=sv.OverlapFilter.NON_MAX_MERGE, iou_threshold=0.5)

    # Annotasi gambar dengan deteksi
    annotated_image = sv.BoxAnnotator().annotate(scene=img.copy(), detections=filtered_detections)

    # Simpan gambar dengan prediksi
    output_path = "/tmp/prediction.jpg"
    cv2.imwrite(output_path, annotated_image)

    # Hapus file sementara
    os.remove(temp_file_path)

    # Menghitung jumlah objek per kelas
    class_count = {}
    for detection in filtered_detections:
        class_name = detection.class_name
        if class_name in class_count:
            class_count[class_name] += 1
        else:
            class_count[class_name] = 1

    # Hasil perhitungan objek
    result_text = "Jumlah objek per kelas:\n"
    for class_name, count in class_count.items():
        result_text += f"{class_name}: {count} objek\n"

    return output_path, result_text

# Membuat antarmuka Gradio
iface = gr.Interface(
    fn=detect_objects,                         # Fungsi yang dipanggil saat gambar diupload
    inputs=gr.Image(type="pil"),               # Input berupa gambar
    outputs=[gr.Image(), gr.Textbox()],        # Output gambar dan teks
    live=True                                    # Menampilkan hasil secara langsung
)

# Menjalankan antarmuka
iface.launch()