import gradio as gr
from roboflow import Roboflow
import tempfile
import os
from sahi.slicing import slice_image
from sahi import get_sliced_prediction
from sahi.postprocess import postprocess_predictions

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

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

    # Slice gambar menjadi potongan-potongan kecil
    slice_image_result = slice_image(
        image=temp_file_path,
        output_file_name="sliced_image",
        output_dir="/tmp/sliced/",
        slice_height=256,
        slice_width=256,
        overlap_height_ratio=0.1,
        overlap_width_ratio=0.1
    )

    # Jalankan prediksi pada setiap potongan gambar
    predictions = get_sliced_prediction(
        image_path=temp_file_path,
        model_path="path_to_your_model",
        model_type="yolov5",  # Adjust based on your model
        slice_height=256,
        slice_width=256,
        overlap_height_ratio=0.1,
        overlap_width_ratio=0.1,
        model_confidence_threshold=0.25
    )

    # Postprocess dan gabungkan hasil prediksi
    postprocessed_predictions = postprocess_predictions(
        predictions=predictions,
        postprocess_type='NMS',
        iou_threshold=0.5
    )

    # Annotate gambar dengan hasil prediksi
    annotated_image = model.annotate_image_with_predictions(temp_file_path, postprocessed_predictions)

    # Simpan gambar hasil annotasi
    output_image_path = "/tmp/prediction.jpg"
    annotated_image.save(output_image_path)

    # Menghitung jumlah objek per kelas
    class_count = {}
    for detection in postprocessed_predictions:
        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"

    # Hapus file sementara
    os.remove(temp_file_path)

    return output_image_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()