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 from random import randint #https://colab.research.google.com/drive/11CtQpSeRBGAuw4TtE_rL470tRo-1X-p2#scrollTo=edGukUHXyymr #https://colab.research.google.com/drive/1xq33YKf0LVKCkbbUZIoNPzgpR_4Kd0qL#scrollTo=sPuM8Xypjs-c #https://github.com/dhawan98/Post-Processing-of-Image-Segmentation-using-CRF 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 """ id = randint(0, 10) print('Starting job ', id) # 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 print('generating mask...') mask = get_mask(resized_img) #mask.save('mask.jpg') # Created a styled sofa # print('Styling sofa...') # styled_sofa = create_styledSofa(resized_img,resized_style) # #styled_sofa.save('styled_sofa.jpg') # # postprocess the final image # print('Replacing sofa...') # new_sofa = replace_sofa(resized_img,mask,styled_sofa) # new_sofa = new_sofa.crop(box) print('Finishing job', id) return mask demo = gr.Interface( style_sofa, inputs = [gr.inputs.Image(),gr.inputs.Image()], outputs = '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 💭!\ \nProvide a picture of your sofa and a desired pattern\ or choose one of the examples below. ⬇", theme="huggingface", # 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()