import streamlit as st
from PIL import Image
import numpy as np
import cv2
import tensorflow as tf
from utils.model import Unet_3

st.header("Segmentación de pulmón con Rayos X")


st.markdown(
    """
Este es un demo para la clase de Platzi, utilizando un módelo de segmentación desarrollado como parte de mi trabajo de grado.

Más Información del modelo se puede encontrar en el siguiente artículo:

https://www.sciencedirect.com/science/article/pii/S2666827021000694?via%3Dihub

"""
)

## Creamos la función del modelo con tensorflow
model=Unet_3()

model.load_weights("utils/Unet_3_M+J+N_1.h5")

## 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 = ["Lung_1.png", "Lung_2.png", "Lung_3.png","Lung_4.png"]
path="assets/Examples/"
## Creamos tres columnas; en cada una estará una imagen ejemplo
col1, col2, col3, col4 = st.columns(4)
with col1:
    ## Se carga la imagen y se muestra en la interfaz
    ex = Image.open(path+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 = path+ejemplos[0]

with col2:
    ex1 = Image.open(path+ejemplos[1])
    st.image(ex1, width=200)
    if st.button("Corre este ejemplo 2"):
        archivo_imagen = path+ejemplos[1]

with col3:
    ex2 = Image.open(path+ejemplos[2])
    st.image(ex2, width=200)
    if st.button("Corre este ejemplo 3"):
        archivo_imagen = path+ejemplos[2]

with col4:
    ex2 = Image.open(path+ejemplos[3])
    st.image(ex2, width=200)
    if st.button("Corre este ejemplo 4"):
        archivo_imagen = path+ejemplos[3]

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

    #Variables estadisticas de las imagenes con las que se entrenó
    mean_seg=0.5397013
    std_seg=0.21162419

    ## Procesamos la imagen para ingresarla al modelo
    img_cv=cv2.resize(img, (224,224))
    
    img_cv = convertir_one_channel(img_cv)
    #Normalizacion
    maxi=np.max(img_cv)
    mini=np.min(img_cv)
    img_N=(img_cv-mini)/(maxi-mini)
    #Estandarizacion
    img_P=(img_N-mean_seg)/std_seg
    
    ## Ingresamos el array de NumPy al modelo
    mask=model.predict(np.expand_dims(np.expand_dims(img_P,axis=-1),axis=0))
    #mask=np.resize(mask,(224,224))
    #predicted=np.array(mask>0.5).astype(int)
    #img_seg_1=img_cv*mask
    predicted = mask[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, (0, 255, 0), 3)
    #output=img_seg_1
    ## 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)