import streamlit as st
from ultralytics import YOLO
from PIL import Image, ImageDraw
import requests
from io import BytesIO
import tempfile
import os

# Load model dari Hugging Face
@st.cache_resource
def load_model():
    model_path = "best.pt"
    if not os.path.exists(model_path):
        hf_url = "https://huggingface.co/rakharaditya/cerfixscan/resolve/main/best.pt"
        response = requests.get(hf_url)
        with open(model_path, "wb") as f:
            f.write(response.content)
    return YOLO(model_path)

model = load_model()

# UI
st.title("📷 CerfixScan: Serviks Detection")
st.markdown("Ambil gambar serviks via **kamera** atau **upload file**.")

# Input dari kamera atau upload manual
img_input = st.camera_input("Ambil gambar dari kamera (klik tombol kamera di bawah)", key="camera")
uploaded_file = st.file_uploader("Atau upload gambar dari galeri", type=["jpg", "jpeg", "png"])

image = None
if img_input:
    image = Image.open(img_input).convert("RGB")
elif uploaded_file:
    image = Image.open(uploaded_file).convert("RGB")

if image:
    st.image(image, caption="Input Gambar", use_column_width=True)

    threshold = 0.92

    # Simpan image sementara
    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp:
        image.save(tmp.name)
        results = model.predict(source=tmp.name, conf=0.5)
        os.unlink(tmp.name)

    result = results[0]
    draw = ImageDraw.Draw(image)
    download_allowed = True

    for box in result.boxes:
        conf = float(box.conf)
        cls = int(box.cls)
        label = result.names[cls]
        xyxy = box.xyxy[0].tolist()
        x1, y1, x2, y2 = map(int, xyxy)

        color = "green" if conf >= threshold else "red"
        if conf < threshold:
            download_allowed = False

        draw.rectangle([x1, y1, x2, y2], outline=color, width=4)
        draw.text((x1, y1 - 10), f"{label} {conf:.2f}", fill=color)

    st.image(image, caption="Hasil Deteksi", use_column_width=True)

    if download_allowed:
        st.success("Confidence ≥ 0.92 ✅ Serviks terdeteksi")
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        st.download_button("📥 Download Gambar", buffered.getvalue(), file_name="deteksi_serviks.jpg", mime="image/jpeg")
    else:
        st.warning("Confidence < 0.92 ❌ Gambar tidak valid untuk deteksi serviks")