import cv2
import numpy as np
from PIL import Image
import gradio as gr
from transformers import SamModel, SamProcessor

# Download model automatically (no .pth needed)
processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
model = SamModel.from_pretrained("facebook/sam-vit-base")

def auto_dimension(image, height, width):
    image_np = np.array(image)

    # Prepare for SAM
    inputs = processor(image, input_points=None, return_tensors="pt")

    # Run SAM
    outputs = model(**inputs)
    masks = processor.post_process_masks(
        outputs.pred_masks,
        inputs["original_sizes"],
        inputs["reshaped_input_sizes"]
    )
    mask = masks[0][0].numpy().astype(np.uint8) * 255

    # Find largest contour
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cnt = max(contours, key=cv2.contourArea)
    x, y, w, h = cv2.boundingRect(cnt)

    # Scale using real height
    scale = height / h
    computed_width = w * scale

    output = image_np.copy()

    # Draw dimension lines
    def draw_dim(pt1, pt2, text):
        cv2.line(output, pt1, pt2, (0, 255, 0), 3)
        cv2.putText(output, text, (pt1[0]+10, pt1[1]-10),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,0), 2)

    draw_dim((x-40, y), (x-40, y+h), f"H: {height} cm")
    draw_dim((x, y+h+40), (x+w, y+h+40), f"W: {width} cm")

    return Image.fromarray(output)

iface = gr.Interface(
    fn=auto_dimension,
    inputs=[
        gr.Image(type="pil"),
        gr.Number(label="Height (cm)"),
        gr.Number(label="Width (cm)")
    ],
    outputs=gr.Image(),
    title="Automatic Dimension Overlay Tool",
    description="Upload an image → enter real dimensions → AI adds dimension graphics automatically."
)

iface.launch()