Spaces:
Build error
Build error
File size: 6,287 Bytes
09012f9 3dacec1 09012f9 3dacec1 09012f9 4e64649 a37eb28 09012f9 e4fb230 09012f9 e4fb230 09012f9 e4fb230 09012f9 e4fb230 09012f9 e4fb230 09012f9 e4fb230 09012f9 e4fb230 09012f9 e4fb230 09012f9 e4fb230 09012f9 e4fb230 09012f9 e4fb230 09012f9 e4fb230 09012f9 e4fb230 09012f9 e4fb230 09012f9 e4fb230 5000f96 e4fb230 5000f96 09012f9 e4fb230 09012f9 3dacec1 09012f9 3dacec1 09012f9 3dacec1 d1be09c 09012f9 3dacec1 09012f9 3dacec1 09012f9 f2e8e46 09012f9 baaaa83 09012f9 4e64649 09012f9 4e64649 09012f9 4e64649 3dacec1 09012f9 3dacec1 09012f9 79c6687 a402a5e 09012f9 c63f571 c06d8b9 09012f9 3dacec1 09012f9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 |
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",
enable_queue=True,
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)
|