import gradio as gr
from ultralytics import YOLO, SAM
from PIL import Image, ImageDraw
import numpy as np

yolo_model = YOLO("best.onnx", task="segment")
sam_model = SAM("sam_b.pt")

def predict(image):
    img_array = np.array(image)

    yolo_results = yolo_model.predict(img_array, conf=0.25, verbose=False)[0]

    detections = []
    boxes = []

    if yolo_results.boxes is not None:
        for box in yolo_results.boxes:
            bbox = box.xyxy[0].tolist()
            boxes.append(bbox)
            cls = yolo_results.names[int(box.cls)]
            conf = round(float(box.conf), 3)
            detections.append(f"{cls}: {conf}")

    sam_masks_count = 0
    if boxes:
        sam_results = sam_model(img_array, bboxes=boxes)
        if sam_results and sam_results[0].masks is not None:
            sam_masks_count = len(sam_results[0].masks)

    result_text = f"Aniqlangan qismlar: {len(detections)}\nSAM masks: {sam_masks_count}\n\n"
    result_text += "\n".join(detections)

    annotated = yolo_results.plot()
    annotated_image = Image.fromarray(annotated[..., ::-1])

    return annotated_image, result_text

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Mashina rasmini yuklang"),
    outputs=[
        gr.Image(type="pil", label="Segmentatsiya natijasi"),
        gr.Textbox(label="Aniqlangan qismlar")
    ],
    title="🚗 Auto Damage Segmentation",
    description="Mashina qismlarini aniqlash — YOLO + SAM Pipeline",
    examples=[]
)

demo.launch()