from cv2 import transpose
import numpy as np
import gradio as gr
from segmentation import get_mask,replace_sofa
from styleTransfer import create_styledSofa
from PIL import Image

def resize_sofa(img):
    """
    This function adds padding to make the orignal image square and 640by640.
    It also returns the orignal ratio of the image, such that it can be reverted later.
    Parameters:
        img = original image
    Return:
        im1 = squared image
        box = parameters to later crop the image to it original ratio
    """
    width, height = img.size 
    idx = np.argmin([width,height])
    newsize = (640, 640) # parameters from test script

    if idx==0:
        img1 = Image.new(img.mode, (height, height), (255, 255, 255))
        img1.paste(img, ((height-width)//2, 0))
        box = ( newsize[0]*(1-width/height)//2,
                0, 
                newsize[0]-newsize[0]*(1-width/height)//2,
                newsize[1])
    else:
        img1 = Image.new(img.mode, (width, width), (255, 255, 255))
        img1.paste(img, (0, (width-height)//2)) 
        box = (0, 
                newsize[1]*(1-height/width)//2, 
                newsize[0], 
                newsize[1]-newsize[1]*(1-height/width)//2)
    im1 = img1.resize(newsize)
    return im1,box

def resize_style(img):
    """
    This function generates a zoomed out version of the style image and resizes it to a 640by640 square.
    Parameters:
        img = image containing the style/pattern
    Return:
        dst = a zoomed-out and resized version of the pattern
    """
    width, height = img.size
    idx = np.argmin([width,height])

    # Makes the image square by cropping
    if idx==0:
        top= (height-width)//2
        bottom= height-(height-width)//2
        left = 0
        right= width
    else:
        left = (width-height)//2
        right = width - (width-height)//2
        top = 0
        bottom = height
    newsize = (640, 640) # parameters from test script
    im1 = img.crop((left, top, right, bottom))

    # Constructs a zoomed-out version
    copies = 8
    resize = (newsize[0]//copies,newsize[1]//copies)
    dst = Image.new('RGB', (resize[0]*copies,resize[1]*copies))
    im2 = im1.resize((resize))
    for row in range(copies):
        im2 = im2.transpose(Image.FLIP_LEFT_RIGHT)
        for column in range(copies):
            im2 = im2.transpose(Image.FLIP_TOP_BOTTOM)
            dst.paste(im2, (resize[0]*row, resize[1]*column))
    dst = dst.resize((newsize))
    return dst

def style_sofa(input_img: np.ndarray, style_img: np.ndarray):
    """
    Styles (all) the sofas in the image to the given style.
    This function uses a transformer to combine the image with the desired style according 
    to a generated mask of the sofas in the image.
    Input:
        input_img = image containing a sofa
        style_img = image containing a style
    Return:
        new_sofa  = image containing the styled sofa
    """
    
    # preprocess input images to be (640,640) squares to fit requirements of the segmentation model
    input_img,style_img = Image.fromarray(input_img),Image.fromarray(style_img)
    resized_img,box = resize_sofa(input_img)
    resized_style = resize_style(style_img)
    resized_style.save('resized_style.jpg')
    # generate mask for image
    mask = get_mask(resized_img)
    mask.save('mask.jpg')
    styled_sofa = create_styledSofa(resized_img,resized_style)
    styled_sofa.save('styled_sofa.jpg')
    # postprocess the final image
    new_sofa = replace_sofa(resized_img,mask,styled_sofa)
    new_sofa = new_sofa.crop(box)
    return new_sofa

image = gr.inputs.Image()
style = gr.inputs.Image()

# Examples
example1 = ['sofa_example1.jpg','style_example1.jpg']
example2 = ['sofa_example1.jpg','style_example2.jpg']
example3 = ['sofa_example1.jpg','style_example3.jpg']
example4 = ['sofa_example1.jpg','style_example4.jpg']
example5 = ['sofa_example1.jpg','style_example5.jpg']

demo = gr.Interface(
    style_sofa,
    inputs = [image,style],
    outputs = 'image',
    examples= [example1, example2, example3 ,example4 ,example5],
    title="🛋 Style your sofa 🛋 ",
    description="Customize your sofa to your wildest dreams 💭!\
                \nProvide a picture of your sofa and a desired pattern\
                 or choose one of the examples below. ⬇",
    # article="**References**\n\n"
    # "<a href='https://www.tensorflow.org/hub/tutorials/tf2_arbitrary_image_stylization' target='_blank'>1. Tutorial to implement Fast Neural Style Transfer using the pretrained model from TensorFlow Hub</a>  \n"
    # "<a href='https://huggingface.co/spaces/luca-martial/neural-style-transfer' target='_blank'>2. The idea to build a neural style transfer application was inspired from this Hugging Face Space </a>"
)

if __name__ == "__main__":
    demo.launch()


#https://github.com/dhawan98/Post-Processing-of-Image-Segmentation-using-CRF