espejelomar's picture
Update app.py
5ceaa50
import streamlit as st
from PIL import Image
import numpy as np
import cv2
from huggingface_hub import from_pretrained_keras
st.header("Segmentaci贸n de dientes con rayos X")
st.markdown(
"""
Hola estudiantes de Platzi 馃殌. Este modelo usa UNet para segmentar im谩genes
de dientes en rayos X. Se utila un modelo de Keras importado con la funci贸n
`huggingface_hub.from_pretrained_keras`. Recuerda que el Hub de Hugging Face est谩 integrado
con muchas librer铆as como Keras, scikit-learn, fastai y otras.
El modelo fue creado por [SerdarHelli](https://huggingface.co/SerdarHelli/Segmentation-of-Teeth-in-Panoramic-X-ray-Image-Using-U-Net).
"""
)
## Seleccionamos y cargamos el modelo
model_id = "SerdarHelli/Segmentation-of-Teeth-in-Panoramic-X-ray-Image-Using-U-Net"
model = from_pretrained_keras(model_id)
## Permitimos a la usuaria cargar una imagen
archivo_imagen = st.file_uploader("Sube aqu铆 tu imagen.", type=["png", "jpg", "jpeg"])
## Si una imagen tiene m谩s de un canal entonces se convierte a escala de grises (1 canal)
def convertir_one_channel(img):
if len(img.shape) > 2:
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
return img
else:
return img
def convertir_rgb(img):
if len(img.shape) == 2:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
return img
else:
return img
## Manipularemos la interfaz para que podamos usar im谩genes ejemplo
## Si el usuario da click en un ejemplo entonces el modelo correr谩 con 茅l
ejemplos = ["dientes_1.png", "dientes_2.png", "dientes_3.png"]
## Creamos tres columnas; en cada una estar谩 una imagen ejemplo
col1, col2, col3 = st.columns(3)
with col1:
## Se carga la imagen y se muestra en la interfaz
ex = Image.open(ejemplos[0])
st.image(ex, width=200)
## Si oprime el bot贸n entonces usaremos ese ejemplo en el modelo
if st.button("Corre este ejemplo 1"):
archivo_imagen = ejemplos[0]
with col2:
ex1 = Image.open(ejemplos[1])
st.image(ex1, width=200)
if st.button("Corre este ejemplo 2"):
archivo_imagen = ejemplos[1]
with col3:
ex2 = Image.open(ejemplos[2])
st.image(ex2, width=200)
if st.button("Corre este ejemplo 3"):
archivo_imagen = ejemplos[2]
## Si tenemos una imagen para ingresar en el modelo entonces
## la procesamos e ingresamos al modelo
if archivo_imagen is not None:
## Cargamos la imagen con PIL, la mostramos y la convertimos a un array de NumPy
img = Image.open(archivo_imagen)
st.image(img, width=850)
img = np.asarray(img)
## Procesamos la imagen para ingresarla al modelo
img_cv = convertir_one_channel(img)
img_cv = cv2.resize(img_cv, (512, 512), interpolation=cv2.INTER_LANCZOS4)
img_cv = np.float32(img_cv / 255)
img_cv = np.reshape(img_cv, (1, 512, 512, 1))
## Ingresamos el array de NumPy al modelo
predicted = model.predict(img_cv)
predicted = predicted[0]
## Regresamos la imagen a su forma original y agregamos las m谩scaras de la segmentaci贸n
predicted = cv2.resize(
predicted, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_LANCZOS4
)
mask = np.uint8(predicted * 255) #
_, mask = cv2.threshold(
mask, thresh=0, maxval=255, type=cv2.THRESH_BINARY + cv2.THRESH_OTSU
)
kernel = np.ones((5, 5), dtype=np.float32)
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=1)
cnts, hieararch = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
output = cv2.drawContours(convertir_one_channel(img), cnts, -1, (255, 0, 0), 3)
## Si obtuvimos exitosamente un resultadod entonces lo mostramos en la interfaz
if output is not None:
st.subheader("Segmentaci贸n:")
st.write(output.shape)
st.image(output, width=850)