import numpy as np
import torch
import torch.nn.functional as F
from torchvision.transforms.functional import normalize
from huggingface_hub import hf_hub_download
import gradio as gr
from gradio_imageslider import ImageSlider
from briarmbg import BriaRMBG
import PIL
from PIL import Image
from typing import Tuple
import requests
import os

token_externo = os.getenv("token_server_externo")
token_api_id_img_adultos = os.getenv("token_leer")




net=BriaRMBG()
# model_path = "./model1.pth"
model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth')
if torch.cuda.is_available():
    net.load_state_dict(torch.load(model_path))
    net=net.cuda()
else:
    net.load_state_dict(torch.load(model_path,map_location="cpu"))
net.eval() 




def resize_image(image):
    image = image.convert('RGB')
    model_input_size = (1024, 1024)
    image = image.resize(model_input_size, Image.BILINEAR)
    return image



def process(image):

    respuesta = "Error"

    API_URL = "https://api-inference.huggingface.co/models/CowCowC/Adu_mod_id_img"
    headers = {"Authorization": "Bearer "+token_api_id_img_adultos+""}

    with open(image, "rb") as f:
        data = f.read()
    response = requests.post(API_URL, headers=headers, data=data, json={"options":{"wait_for_model":True}})
    #return response.json()

    respuesta_identifica_imagen = ""

    respuesta = response.json()
    for item in respuesta: 
        if item["label"] == "regular":
            if item["score"] > 0.8:
                respuesta_identifica_imagen = "imagen normal"
                print("Es una imagen normal")

                #URL
                # prepare input
                im = Image.open(image)
                orig_image = im
                #orig_image = Image.fromarray(image)
                #orig_image = Image.open(requests.get(image, stream=True).raw)
                #orig_image = np.array(im_url)
                w,h = orig_im_size = orig_image.size
                image = resize_image(orig_image)
                im_np = np.array(image)
                im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
                im_tensor = torch.unsqueeze(im_tensor,0)
                im_tensor = torch.divide(im_tensor,255.0)
                im_tensor = normalize(im_tensor,[0.5,0.5,0.5],[1.0,1.0,1.0])
                if torch.cuda.is_available():
                    im_tensor=im_tensor.cuda()
            
                #inference
                result=net(im_tensor)
                # post process
                result = torch.squeeze(F.interpolate(result[0][0], size=(h,w), mode='bilinear') ,0)
                ma = torch.max(result)
                mi = torch.min(result)
                result = (result-mi)/(ma-mi)    
                # image to pil
                im_array = (result*255).cpu().data.numpy().astype(np.uint8)
                pil_im = Image.fromarray(np.squeeze(im_array))
                # paste the mask on the original image
                new_im = Image.new("RGBA", pil_im.size, (0,0,0,0))
                new_im.paste(orig_image, mask=pil_im)
                # new_orig_image = orig_image.convert('RGBA')
                # return [new_orig_image, new_im]
                #return new_im

                new_im.save("img_transparente.png","PNG")
                url = 'https://imupmx.000webhostapp.com/'
                files = {'uploaded_file': open('img_transparente.png', 'rb')}
                r = requests.post(url, files=files)
                
                if r.status_code == 200:
                    respuesta = r.text
                else:
                    respuesta = "Error_servidor_externo"
            else:
                respuesta_identifica_imagen = "imagen adultos"
                print("Es una imagen para adultos")
                respuesta = "Error_imagen_no_valida"


    return respuesta
    



# output = ImageSlider(position=0.5,label='Image without background', type="pil", show_download_button=True)
# demo = gr.Interface(fn=process,inputs="image", outputs=output, examples=examples, title=title, description=description)
# demo = gr.Interface(fn=process,inputs="image", outputs="image")
demo = gr.Interface(fn=process,inputs=gr.Image(label="Input Image Component", type="filepath"), outputs="text", title="title", description="description")


#if __name__ == "__main__":
    #demo.launch(share=True)

demo.launch(share=True)