'''
This Python script is a web application that performs human body part segmentation 
using a pre-trained deep learning model called DeepLabv3+. 
The application is built using the Streamlit library and uses the Hugging Face Hub 
to download the pre-trained model.
'''

# import libraries
import numpy as np
import tensorflow as tf
import streamlit as st
from PIL import Image
from huggingface_hub import from_pretrained_keras
import cv2

# The model used is the DeepLabv3+ model with a ResNet50 backbone.
model = from_pretrained_keras("keras-io/deeplabv3p-resnet50")

# A colormap is defined to map the predicted segmentation masks to colors for better visualization
colormap = np.array([[0,0,0], [31,119,180], [44,160,44], [44, 127, 125], [52, 225, 143],
                    [217, 222, 163], [254, 128, 37], [130, 162, 128], [121, 7, 166], [136, 183, 248],
                    [85, 1, 76], [22, 23, 62], [159, 50, 15], [101, 93, 152], [252, 229, 92],
                    [167, 173, 17], [218, 252, 252], [238, 126, 197], [116, 157, 140], [214, 220, 252]], dtype=np.uint8)

# size of the input image is defined as 512x512 pixels                    
img_size = 512
                    
def read_image(image):
    '''
    read_image: reads in the input image and preprocesses it 
                by resizing it to the defined size and normalizing it to values between -1 and 1
    '''
    image = tf.convert_to_tensor(image)
    image.set_shape([None, None, 3])
    image = tf.image.resize(images=image, size=[img_size, img_size])
    image = image / 255
    return image


def infer(model, image_tensor):
    '''
    infer: performs inference using the pre-trained model and returns the predicted segmentation mask.
    '''
    predictions = model.predict(np.expand_dims((image_tensor), axis=0))
    predictions = np.squeeze(predictions)
    predictions = np.argmax(predictions, axis=2)
    return predictions


def decode_segmentation_masks(mask, colormap, n_classes):
    '''
     decode_segmentation_masks: maps the predicted segmentation mask to the defined colormap 
                                to produce a colored mask.
     '''
    r = np.zeros_like(mask).astype(np.uint8)
    g = np.zeros_like(mask).astype(np.uint8)
    b = np.zeros_like(mask).astype(np.uint8)
    for l in range(0, n_classes):
        idx = mask == l
        r[idx] = colormap[l, 0]
        g[idx] = colormap[l, 1]
        b[idx] = colormap[l, 2]
    rgb = np.stack([r, g, b], axis=2)
    return rgb


def get_overlay(image, colored_mask):
    '''
    get_overlay: overlays the colored mask on the original image for visualization
    '''
    image = tf.keras.preprocessing.image.array_to_img(image)
    image = np.array(image).astype(np.uint8)
    overlay = cv2.addWeighted(image, 0.35, colored_mask, 0.65, 0)
    return overlay


def segmentation(input_image):
    '''
    segmentation: 
        returns,
           - prediction_colormap: function is used to convert the prediction mask into a colored mask, 
             where each class is assigned a unique color from a predefined color map.

           - overlay: used to create an overlay image by blending the original input image with the colored mask
    '''
    image_tensor = read_image(input_image)
    prediction_mask = infer(image_tensor=image_tensor, model=model)
    prediction_colormap = decode_segmentation_masks(prediction_mask, colormap, 20)
    overlay = get_overlay(image_tensor, prediction_colormap)
    return (overlay, prediction_colormap)


## Streamlit interface

st.header("Segmentación de partes del cuerpo humano")

st.subheader("Demo de Spaces usando Streamlit y segmentacion de imagenes [Space original](https://huggingface.co/spaces/PKaushik/Human-Part-Segmentation)")

st.markdown("Sube una imagen o selecciona un ejemplo para segmentar las distintas partes del cuerpo humano")

file_imagen = st.file_uploader("Sube aquí tu imagen", type=["png", "jpg", "jpeg"])

examples = ["example_image_1.jpg", "example_image_2.jpg", "example_image_3.jpg"]

col1, col2, col3 = st.columns(3)
with col1:
    ex1 = Image.open(examples[0])
    st.image(ex1, width=200)
    if st.button("Corre ejemplo 1"):
        file_imagen = examples[0]

with col2:
    ex2 = Image.open(examples[1])
    st.image(ex2, width=200)
    if st.button("Corre ejemplo 2"):
        file_imagen = examples[1]

with col3:
    ex3 = Image.open(examples[2])
    st.image(ex3, width=200)
    if st.button("Corre ejemplo 3"):
        file_imagen = examples[2]

if file_imagen is not None:
    img = Image.open(file_imagen)

    output = segmentation(img)
    if output is not None:
        st.subheader("Original: ")
        st.image(img, width=850)
        col1, col2 = st.columns(2)

        with col1:
            st.subheader("Segmentación: ")
            st.image(output[0], width=425)

        with col2:
            st.subheader("Mask: ")
            st.image(output[1], width=425)
else:
    st.write("Ocurrió un error al segmentar la imagen, intenta con otra imagen")