import gradio as gr
from transformers import AutoImageProcessor, AutoModelForObjectDetection
import torch
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
import io

# โหลดโมเดลและตัวประมวลผล
processor = AutoImageProcessor.from_pretrained("0llheaven/Conditional-detr-finetuned")
model = AutoModelForObjectDetection.from_pretrained("0llheaven/Conditional-detr-finetuned")

def detect_objects(image):
    # แปลงรูปภาพเป็น RGB หากเป็น grayscale
    if image.mode != "RGB":
        image = image.convert("RGB")
    
    # เตรียม input สำหรับโมเดล
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)

    # กรองการทำนายที่มีความแม่นยำมากกว่า 0.5
    target_sizes = torch.tensor([image.size[::-1]])
    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)

    # วาดกรอบรอบวัตถุที่ตรวจพบในภาพ
    draw = ImageDraw.Draw(image)
    for result in results:
        scores = result["scores"]
        labels = result["labels"]
        boxes = result["boxes"]

        for score, label, box in zip(scores, labels, boxes):
            box = [round(i, 2) for i in box.tolist()]
            label_name = "Pneumonia" if label.item() == 0 else "Other"
            draw.rectangle(box, outline="red", width=3)
            draw.text((box[0], box[1]), f"{label_name}: {round(score.item(), 3)}", fill="red")
    
    # แปลงภาพเป็นรูปแบบที่สามารถแสดงผลได้ใน Gradio
    output_image = io.BytesIO()
    image.save(output_image, format='PNG')
    output_image.seek(0)
    
    return output_image

# สร้างอินเตอร์เฟซด้วย Gradio
interface = gr.Interface(
    fn=detect_objects, 
    inputs=gr.inputs.Image(type="pil"), 
    outputs=gr.outputs.Image(type="auto"),
    title="Object Detection with Transformers",
    description="Upload an image to detect objects using a fine-tuned Conditional-DETR model."
)

# เปิดใช้งานอินเตอร์เฟซ
interface.launch()