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

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
    resized_img,box = resize_sofa(input_img)
    resized_style = resize_style(style_img)
    # generate mask for image
    print('starting mask')
    mask = get_mask(resized_img)
    print('got mask')
    styled_sofa = create_styledSofa(resized_img,resized_style)
    print('resizing')
    new_sofa = replace_sofa(resized_img,mask,styled_sofa)
    new_sofa = Image.fromarray(new_sofa).crop(box)
    return new_sofa

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

demo = gr.Interface(
    style_sofa,
    [image,style],
    'image',
    examples=[
        ['sofa_example1.jpg','style_example1.jpg'],
        ['sofa_example1.jpg','style_example2.jpg'],
        ['sofa_example1.jpg','style_example3.jpg'],
        ['sofa_example1.jpg','style_example4.jpg'],
        ['sofa_example1.jpg','style_example5.jpg'],
    ],
    title="Style your sofa",
    description="🛋 Customize your sofa to your wildest dreams! 🛋",
)

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


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