alexrods
add model output in app.py
4e0c08c
raw
history blame
3.64 kB
import numpy as np
import tensorflow as tf
import streamlit as st
from PIL import Image
from huggingface_hub import from_pretrained_keras
import cv2
model = from_pretrained_keras("keras-io/deeplabv3p-resnet50")
colormap = np.array([[0,0,0], [31,119,180], [44,160,44], [44, 127, 125], [52, 225, 143],
[217, 222, 163], [254, 128, 37], [130, 162, 128], [121, 7, 166], [136, 183, 248],
[85, 1, 76], [22, 23, 62], [159, 50, 15], [101, 93, 152], [252, 229, 92],
[167, 173, 17], [218, 252, 252], [238, 126, 197], [116, 157, 140], [214, 220, 252]], dtype=np.uint8)
img_size = 512
def read_image(image):
image = tf.convert_to_tensor(image)
image.set_shape([None, None, 3])
image = tf.image.resize(images=image, size=[img_size, img_size])
image = image / 127.5 - 1
return image
def infer(model, image_tensor):
predictions = model.predict(np.expand_dims((image_tensor), axis=0))
predictions = np.squeeze(predictions)
predictions = np.argmax(predictions, axis=2)
return predictions
def decode_segmentation_masks(mask, colormap, n_classes):
r = np.zeros_like(mask).astype(np.uint8)
g = np.zeros_like(mask).astype(np.uint8)
b = np.zeros_like(mask).astype(np.uint8)
for l in range(0, n_classes):
idx = mask == l
r[idx] = colormap[l, 0]
g[idx] = colormap[l, 1]
b[idx] = colormap[l, 2]
rgb = np.stack([r, g, b], axis=2)
return rgb
def get_overlay(image, colored_mask):
image = tf.keras.preprocessing.image.array_to_img(image)
image = np.array(image).astype(np.uint8)
overlay = cv2.addWeighted(image, 0.35, colored_mask, 0.65, 0)
return overlay
def segmentation(input_image):
image_tensor = read_image(input_image)
prediction_mask = infer(image_tensor=image_tensor, model=model)
prediction_colormap = decode_segmentation_masks(prediction_mask, colormap, 20)
overlay = get_overlay(image_tensor, prediction_colormap)
return (overlay, prediction_colormap)
# i = gr.inputs.Image()
# o = [gr.outputs.Image('pil'), gr.outputs.Image('pil')]
st.header("Segmentacion de partes del cuerpo humano")
st.markdown("Sube una imagen o selecciona un ejemplo para segmentar las distintas partes del cuerpo humano")
file_imagen = st.file_uploader("Sube aqui tu imagen", type=["png", "jpg", "jpeg"])
examples = ["example_image_1.jpg", "example_image_2.jpg", "example_image_3.jpg"]
col1, col2, col3 = st.columns(3)
with col1:
ex1 = Image.open(examples[0])
st.image(ex1, width=200)
if st.button("Corre ejemplo 1"):
file_imagen = examples[0]
with col2:
ex2 = Image.open(examples[1])
st.image(ex2, width=200)
if st.button("Corre ejemplo 2"):
file_imagen = examples[1]
with col3:
ex3 = Image.open(examples[2])
st.image(ex3, width=200)
if st.button("Corre ejemplo 3"):
file_imagen = examples[2]
if file_imagen is not None:
img = Image.open(file_imagen)
output = segmentation(img)
if output is not None:
st.subheader("Segmentacion: ")
st.write(output.shape)
st.image(output[0], width=850)
st.subheader("Mask: ")
st.write(output.shape)
st.image(output[1], width=850)
# article = "<div style='text-align: center;'><a href='https://keras.io/examples/vision/deeplabv3_plus/' target='_blank'>Keras example by Praveen Kaushik</a></div>"
# gr.Interface(segmentation, i, o, examples=examples, allow_flagging=False, analytics_enabled=False,
# title=title, description=description, article=article).launch(enable_queue=True)