import gradio as gr

from pathlib import Path
import tensorflow as tf
from tf_bodypix.api import download_model, load_model, BodyPixModelPaths
import numpy as np
from PIL import Image

# load model
modelPath = download_model(BodyPixModelPaths.RESNET50_FLOAT_STRIDE_16)
bodypix_model = load_model(modelPath)


def predict(mask_threshold, image):    
    # get prediction result
    image_array = tf.keras.preprocessing.image.img_to_array(image)
    result = bodypix_model.predict_single(image_array)
    
    # simple mask
    mask = result.get_mask(threshold=mask_threshold)
    
    # colored mask (separate colour for each body part)
    colored_mask = result.get_colored_part_mask(mask)

    colored_mask_image = Image.fromarray(colored_mask.astype('uint8'), 'RGB')

    pred_img = np.array(image) * 0.5 + colored_mask * 0.5
    pred_img = pred_img.astype(np.uint8)
    pred_img

    return colored_mask_image, pred_img;

iface = gr.Interface(fn=predict, inputs=[gr.Number(label='Mask Threshold', value=0.5),"image"], outputs=["image","image"])
iface.launch()