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.subheader("Este es una iteración para bucar mejorar el demo")

st.markdown(
    """
Este es un demo prueba para la clase de Platzi
"""
)

## 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)

st.markdown("Gracias por usar nuestro demo! Nos vemos pronto")