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

seg_tool = SegmentationTool()


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


demo = gr.Interface(fn=segment,
                    inputs=[
                        gr.Image(type='pil'),
                        gr.CheckboxGroup([("Door", 14), ("Window", 8)],
                                         value=[8, 14],
                                         label="Mask Items")],
                    outputs=['image'],
                    allow_flagging='never')
if __name__ == "__main__":
    demo.launch()