from random import randint from typing import Tuple import gradio as gr import numpy as np from PIL import Image from Segmentation.segmentation import get_mask, replace_sofa from StyleTransfer.styleTransfer import create_styledSofa def resize_sofa(img: Image.Image) -> Tuple[Image.Image, tuple]: """ This function adds padding to make the original image square and 640by640. It also returns the original 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: Image.Image) -> Image.Image: """ 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_image: np.ndarray, Style_image: np.ndarray, Choice_of_algorithm: str ) -> Image.Image: """ 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 choice = Style transfer algorithm to use Return: new_sofa = image containing the styled sofa """ id = randint(0, 10) print("Starting job ", id) # preprocess input images to fit requirements of the segmentation model input_img, style_img = Image.fromarray(Input_image), Image.fromarray(Style_image) 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, Choice_of_algorithm) # 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(), gr.inputs.Radio( ["Style Transformer", "Style FAST", "Style Projection"], default="Style FAST", ), ], outputs="image", examples=[ [ "figures/sofa_example1.jpg", "figures/style_example1.jpg", "Style Transformer", ], [ "figures/sofa_example3.jpg", "figures/style_example10.jpg", "Style FAST", ], [ "figures/sofa_example2.jpg", "figures/style_example6.jpg", "Style Projection", ], ], title="🛋 Style your sofa 🛋 ", description="Customize your sofa to your wildest dreams 💭!\ \nProvide a picture of your sofa, a desired pattern\ and (optionally) choose one of the algorithms.\ \nOr just pick one of the examples below. ⬇", theme="huggingface", article="**References**\n\n" "\ 1. The data that was used to train the segmentation model. \ \n" " \ 2. Github repository used to train a segmentation model with transfer. \ learning.\ \n" " \ 3. The github repository that is used for the style transformer. \ \n" " \ 4. A tensorflow model for fast arbitrary image style transfer. \ \n" " \ 5. A paddleHub model for parameter free style transfer. \ \n", ) if __name__ == "__main__": demo.launch(cache_examples=True)