import cv2
import paddlehub as hub
import gradio as gr


module = hub.Module(name="lseg")


def segment(image, labels):
    long_size = max(image.shape[:2])
    if long_size > 512:
        f = 512 / long_size
        image = cv2.resize(image, (0, 0), fx=f, fy=f)
    results = module.segment(
        image=image[..., ::-1],
        labels=[item for item in labels.split('\n') if item != '']
    )
    return [
        results['color'][..., ::-1],
        results['mix'][..., ::-1],
        *[v[..., ::-1] for v in results['classes'].values()]
    ]


gr.Interface(
    fn=segment,
    inputs=['image', gr.Textbox(
        placeholder='other\ncat', lines=5, max_lines=50)],
    outputs=[gr.Gallery().style(grid=[2, 3], height="auto")]
).launch()