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 identificación de dientes en un archivo de rayos X!
"""
)

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