File size: 1,063 Bytes
0ff11f5 23e74e7 ca7c022 23e74e7 7fa8d07 be29020 d769eaf 7c85c75 d769eaf 7c85c75 ca7c022 0ff11f5 d769eaf 96ea0c4 0ff11f5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
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()
|