import gradio as gr
from test import inference_img
from models import *
import numpy as np 
from PIL import Image

device='cpu'
model = StyleMatte()
model = model.to(device)
checkpoint = f"stylematte.pth"
state_dict = torch.load(checkpoint, map_location=f'{device}')

model.load_state_dict(state_dict)
model.eval()

def predict(inp):
    print("***********Inference****************")
    mask = inference_img(model, inp) 
    inp_np = np.array(inp)
    fg = np.uint8((mask[:,:,None]*inp_np))
    alpha_channel = (mask*255).astype(np.uint8)
    print(fg.max(), alpha_channel.max(), fg.shape, alpha_channel.shape)
    print("***********Inference finish****************")
    # print("***********MASK****************", inp_np.max(), mask.max())
    fg = np.dstack((fg, alpha_channel))
    fg_pil = Image.fromarray(fg, 'RGBA')
    
    return [mask, fg_pil]

print("MODEL LOADED")
print("************************************")

iface = gr.Interface(fn=predict, 
             inputs=gr.Image(type="numpy"),
             outputs=[gr.Image(type="numpy"),gr.Image(type="pil", image_mode='RGBA')],
             examples=["./logo.jpeg"])
print("****************Interface created******************")

iface.launch()