import gradio as gr
import os
import random
import numpy as np
import torch
from torch import nn
from torchvision import transforms
from transformers import SegformerForSemanticSegmentation

# examples
os.system("wget -O 073.png https://huggingface.co/spaces/alkzar90/rock-glacier-segmentation/resolve/main/example_images/buenos_resultados/073.png")
os.system("wget -O 356.png https://huggingface.co/spaces/alkzar90/rock-glacier-segmentation/resolve/main/example_images/buenos_resultados/356.png")
os.system("wget -O 599.png https://huggingface.co/spaces/alkzar90/rock-glacier-segmentation/resolve/main/example_images/buenos_resultados/599.png")
os.system("wget -O 630.png https://huggingface.co/spaces/alkzar90/rock-glacier-segmentation/resolve/main/example_images/buenos_resultados/630.png")
os.system("wget -O 673.png https://huggingface.co/spaces/alkzar90/rock-glacier-segmentation/resolve/main/example_images/buenos_resultados/673.png")


os.system("wget -O 019.png https://huggingface.co/spaces/alkzar90/rock-glacier-segmentation/resolve/main/example_images/malos_resultados/019.png")
os.system("wget -O 261.png https://huggingface.co/spaces/alkzar90/rock-glacier-segmentation/resolve/main/example_images/malos_resultados/261.png")
os.system("wget -O 524.png https://huggingface.co/spaces/alkzar90/rock-glacier-segmentation/resolve/main/example_images/malos_resultados/524.png")
os.system("wget -O 716.png https://huggingface.co/spaces/alkzar90/rock-glacier-segmentation/resolve/main/example_images/malos_resultados/716.png")
os.system("wget -O 898.png https://huggingface.co/spaces/alkzar90/rock-glacier-segmentation/resolve/main/example_images/malos_resultados/898.png")

# model-setting
MODEL_PATH="./best_model_mixto/"

device = torch.device("cpu")

preprocessor = transforms.Compose([
                        transforms.Resize(128),
                        transforms.ToTensor()
                 ])
model = SegformerForSemanticSegmentation.from_pretrained(MODEL_PATH)
model.eval()

# inference-functions
def upscale_logits(logit_outputs, size):
   """Escala los logits a (4W)x(4H) para recobrar dimensiones originales del input"""
   return nn.functional.interpolate(
       logit_outputs,
       size=size,
       mode="bilinear",
       align_corners=False
   )


def visualize_instance_seg_mask(mask):
    """Agrega colores RGB a cada una de las clases en la mask"""
    image = np.zeros((mask.shape[0], mask.shape[1], 3))
    labels = np.unique(mask)
    label2color = {label: (random.randint(0, 1), 
                           random.randint(0, 255), 	
                           random.randint(0, 255)) for label in labels}
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            image[i, j, :] = label2color[mask[i, j]]
    image = image / 255
    return image


def query_image(img):
    """Función para generar predicciones a la escala origina"""
    inputs = preprocessor(img).unsqueeze(0)
    with torch.no_grad():
        preds = model(inputs)["logits"]
    preds_upscale = upscale_logits(preds, preds.shape[2])
    predict_label = torch.argmax(preds_upscale, dim=1).to(device)
    result = predict_label[0,:,:].detach().cpu().numpy()
    return visualize_instance_seg_mask(result)

# demo
demo = gr.Interface(
    query_image,
    #inputs=[gr.Image(type="pil").style(full_width=True, height=256)],
    inputs=[gr.ImageMask(type="pil").style(full_width=True, height=256)],
    outputs=[gr.Image().style(full_width=True, height=256)],
    title="Skyguard: segmentador de glaciares de roca  🛰️ +️ 🛡️ ️",
    description="Modelo de segmentación de imágenes para detectar glaciares de roca.<br> Se entrenó un modelo [nvidia/SegFormer](https://huggingface.co/nvidia/mit-b0) con _fine-tuning_ en el [rock-glacier-dataset](https://huggingface.co/datasets/alkzar90/rock-glacier-dataset)",
    examples=[["073.png"], ["356.png"], ["599.png"], ["630.png"], ["673.png"], 
              ["019.png"], ["261.png"], ["524.png"], ["716.png"], ["898.png"]],
    cache_examples=False
)

demo.launch()