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

# 1)  Use Ultralytics' built-in small model (22 MB) – no extra download
model = YOLO("yolov8s.pt")          # auto-downloads first run

REF_WIDTH_CM = 10.0                 # width of calibration card

def measure(image: Image.Image) -> Image.Image:
    frame = np.array(image)
    res   = model(frame, verbose=False)[0]
    boxes = res.boxes.xyxy.cpu().numpy()

    if len(boxes) < 2:              # need statue + card
        return image

    areas   = (boxes[:,2]-boxes[:,0]) * (boxes[:,3]-boxes[:,1])
    ref_idx = int(np.argmin(areas))      # smallest = card
    obj_idx = int(np.argmax(areas))      # largest  = statue

    x1r,y1r,x2r,y2r = boxes[ref_idx]
    x1o,y1o,x2o,y2o = boxes[obj_idx]

    px_per_cm = (x2r - x1r) / REF_WIDTH_CM
    w_cm  = round((x2o - x1o) / px_per_cm, 2)
    h_cm  = round((y2o - y1o) / px_per_cm, 2)

    canvas = frame.copy()
    cv2.rectangle(canvas, (int(x1o),int(y1o)), (int(x2o),int(y2o)), (0,255,0), 2)
    cv2.rectangle(canvas, (int(x1r),int(y1r)), (int(x2r),int(y2r)), (0,0,255), 1)
    cv2.putText(canvas,
                f"{w_cm} × {h_cm} cm",
                (int(x1o), int(max(y1o-10,0))),
                cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0,255,0), 2, cv2.LINE_AA)
    return Image.fromarray(canvas)

demo = gr.Interface(
    measure,
    gr.Image(type="pil", label="Upload studio photo (10 cm card bottom-left)"),
    gr.Image(type="pil", label="BBox + cm overlay"),
    title="SmartDimension – YOLOv8 demo",
    description="Detects gilded statues, finds 10 cm reference card, outputs real-world dimensions."
)

if __name__ == "__main__":
    demo.launch()