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" # "1. Tutorial to implement Fast Neural Style Transfer using the pretrained model from TensorFlow Hub \n" # "2. The idea to build a neural style transfer application was inspired from this Hugging Face Space " ) if __name__ == "__main__": demo.launch() #https://github.com/dhawan98/Post-Processing-of-Image-Segmentation-using-CRF