|
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 |
|
|
|
|
|
modelPath = download_model(BodyPixModelPaths.RESNET50_FLOAT_STRIDE_16) |
|
bodypix_model = load_model(modelPath) |
|
|
|
|
|
def predict(mask_threshold, image): |
|
|
|
image_array = tf.keras.preprocessing.image.img_to_array(image) |
|
result = bodypix_model.predict_single(image_array) |
|
|
|
|
|
mask = result.get_mask(threshold=mask_threshold) |
|
|
|
|
|
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() |
|
|
