Spaces:
Build error
Build error
| from cv2 import transpose | |
| import numpy as np | |
| import gradio as gr | |
| from Segmentation.segmentation import get_mask,replace_sofa | |
| from StyleTransfer.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 new_sofa | |
| demo = gr.Interface( | |
| style_sofa, | |
| inputs = [gr.inputs.Image(),gr.inputs.Image()], | |
| outputs = 'image', | |
| examples= [['figures/sofa_example1.jpg','figures/style_example1.jpg'], | |
| ['figures/sofa_example1.jpg','figures/style_example2.jpg'], | |
| ['figures/sofa_example1.jpg','figures/style_example3.jpg'], | |
| ['figures/sofa_example1.jpg','figures/style_example4.jpg'], | |
| ['figures/sofa_example1.jpg','figures/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" | |
| # "<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() | |