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)