app.py

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"
    "<a href='https://tianchi.aliyun.com/specials/promotion/alibaba-3d-future' \
    target='_blank'>\
    1. The data that was used to train the segmentation model. \
    </a>  \n"
    "<a href='https://github.com/qubvel/segmentation_models' \
    target='_blank'> \
    2. Github repository used to train a segmentation model with transfer. \
    learning.\
    </a>  \n"
    "<a href='https://github.com/diyiiyiii/StyTR-2' \
    target='_blank'> \
    3. The github repository that is used for the style transformer. \
    </a>  \n"
    "<a href='https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2' \
    target='_blank'> \
    4. A tensorflow model for fast arbitrary image style transfer. \
    </a>  \n"
    "<a href='https://github.com/PaddlePaddle/PaddleHub/tree/release/v2.2/modules/image/Image_gan/style_transfer/stylepro_artistic' \
    target='_blank'> \
    5. A paddleHub model for parameter free style transfer. \
    </a>  \n",
)

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