finishing touches

.flake8

@@ -0,0 +1,12 @@
+[flake8]
+exclude =
+    .git,
+    *.egg-info,
+    __pycache__,
+    .tox,
+    .pytest_cache,
+    build,
+    dist,
+    tests
+max-line-length = 88
+ignore = D202,W503,E203 # conflicts with black

README.md

@@ -1,8 +1,8 @@
 ---
 title: SofaStyler
 emoji: 🛋
-colorFrom: pink
-colorTo: yellow
+colorFrom: blue
+colorTo: green
 sdk: gradio
 sdk_version: 2.9.4
 app_file: app.py

app.py

@@ -1,138 +1,186 @@
-from cv2 import transpose
-import numpy as np
+from random import randint
+from typing import Tuple
+
 import gradio as gr
-from Segmentation.segmentation import get_mask,replace_sofa
-from StyleTransfer.styleTransfer import create_styledSofa
+import numpy as np
 from PIL import Image
-from random import randint
 
+from Segmentation.segmentation import get_mask, replace_sofa
+from StyleTransfer.styleTransfer import create_styledSofa
 
-#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):
+def resize_sofa(img: Image.Image) -> Tuple[Image.Image, tuple]:
     """
-    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.
+    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
+    width, height = img.size
+    idx = np.argmin([width, height])
+    newsize = (640, 640)  # parameters from test script
 
-    if idx==0:
+    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])
+        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)
+        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
+    return im1, box
+
 
-def resize_style(img):
+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.
+    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])
+    idx = np.argmin([width, height])
 
     # Makes the image square by cropping
-    if idx==0:
-        top= (height-width)//2
-        bottom= height-(height-width)//2
+    if idx == 0:
+        top = (height - width) // 2
+        bottom = height - (height - width) // 2
         left = 0
-        right= width
+        right = width
     else:
-        left = (width-height)//2
-        right = width - (width-height)//2
+        left = (width - height) // 2
+        right = width - (width - height) // 2
         top = 0
         bottom = height
-    newsize = (640, 640) # parameters from test script
+    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))
+    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.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):
+
+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.
+    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 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)
+    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')
+    # resized_style.save('resized_style.jpg')
     # generate mask for image
-    print('generating mask...')
+    print("generating mask...")
     mask = get_mask(resized_img)
-    #mask.save('mask.jpg')
+    # 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')
+    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)
+    print("Replacing sofa...")
+    new_sofa = replace_sofa(resized_img, mask, styled_sofa)
     new_sofa = new_sofa.crop(box)
-    print('Finishing job', id)
+    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']],
+    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 and a desired pattern\
-                 or choose one of the examples below. ⬇",
+                \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://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>"
+    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()
-
-
+    demo.launch(cache_examples=True)

gradio_cached_examples/log.csv

@@ -0,0 +1,4 @@
+'output'
+'output/0.png'
+'output/1.png'
+'output/2.png'