import numpy as np
import gradio as gr
from model import SegmentationTool

seg_tool = SegmentationTool()


def segment(input_img):
    mask_image, transparent_mask_image, image, items, room = seg_tool.get_mask(image=input_img)
    return mask_image


demo = gr.Interface(segment, gr.Image(shape=(400, 600)), "image")
if __name__ == "__main__":
    demo.launch()