import gradio as gr
import numpy as np 
from PIL import Image
from deepfloorplan_inference import DeepFloorPlanModel

# Load model once at startup
model = DeepFloorPlanModel(model_dir='pretrained')

def predict_floorplan(image):
    # image: PIL Image from Gradio
    result = model.predict(image)
    # Convert numpy array to PIL Image for Gradio output
    return Image.fromarray(result.astype(np.uint8))

iface = gr.Interface(
    fn=predict_floorplan,
    inputs=gr.Image(type="pil", label="Upload Floorplan Image"),
    outputs=gr.Image(type="pil", label="Predicted Segmentation"),
    title="Deep Floor Plan Segmentation",
    description="Upload a floorplan image to get the predicted segmentation using the Deep Floor Plan model.",
    allow_flagging="never"
)

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