documentation

app.py

@@ -3,7 +3,6 @@ import gradio as gr
 from segmentation import get_mask,replace_sofa
 from styleTransfer import resize_sofa,resize_style,create_styledSofa
 from PIL import Image
-import cv2
 
 def style_sofa(input_img: np.ndarray, style_img: np.ndarray):
     """
@@ -21,11 +20,9 @@ def style_sofa(input_img: np.ndarray, style_img: np.ndarray):
     resized_img,box = resize_sofa(input_img)
     resized_style = resize_style(style_img)
     # generate mask for image
-    print('starting mask')
     mask = get_mask(resized_img)
-    print('got mask')
     styled_sofa = create_styledSofa(resized_img,resized_style)
-    print('resizing')
+    # postprocess the final image
     new_sofa = replace_sofa(resized_img,mask,styled_sofa)
     new_sofa = Image.fromarray(new_sofa).crop(box)
     return new_sofa
@@ -44,8 +41,10 @@ demo = gr.Interface(
         ['sofa_example1.jpg','style_example4.jpg'],
         ['sofa_example1.jpg','style_example5.jpg'],
     ],
-    title="Style your sofa",
-    description="🛋 Customize your sofa to your wildest dreams! 🛋",
+    title="🛋 Style your sofa 🛋 ",
+    description="Customize your sofa to your wildest dreams!\
+                \nProvide a picture of your sofa and a desired pattern\
+                \n or choose one of the examples below",
 )
 
 if __name__ == "__main__":

segmentation.py

@@ -8,6 +8,15 @@ from PIL import Image
 import segmentation_models as sm
 
 def get_mask(image):
+    """
+    This function generates a mask of the image that highlights all the sofas in the image.
+    This uses a pre-trained Unet model with a resnet50 backbone.
+    Remark: The model was trained on 640by640 images and it is therefore best that the image has the same size.
+    Parameters:
+            image = original image
+    Return:
+            mask  = corresponding maks of the image
+    """
     model_path = "model_checkpoint.h5"
     CLASSES = ['sofa']
     BACKBONE = 'resnet50'
@@ -31,14 +40,11 @@ def get_mask(image):
     # total_loss = sm.losses.binary_focal_dice_loss # or sm.losses.categorical_focal_dice_loss 
     metrics = [sm.metrics.IOUScore(threshold=0.5), sm.metrics.FScore(threshold=0.5)]
     # compile keras model with defined optimozer, loss and metrics
-    print('compiling')
     model.compile(optim, total_loss, metrics)
 
-    print('loading model weights')
     #load model
     model.load_weights(model_path) 
 
-    print('processing image')
     test_img = np.array(image)#cv2.imread(path, cv2.IMREAD_COLOR)       
     test_img = cv2.resize(test_img, (640, 640))
     test_img = cv2.cvtColor(test_img, cv2.COLOR_RGB2BGR)
@@ -51,9 +57,17 @@ def get_mask(image):
     return np.array(mask) 
 
 def replace_sofa(image,mask,styled_sofa):
-    # print(mask.shape)
-    # mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
-    # print(mask.shape)
+    """
+    This function replaces the original sofa in the image by the new styled sofa according
+    to the mask.
+    Remark: All images should have the same size.
+    Input:
+        image       = Original image
+        mask        = Generated masks highlighting the sofas in the image
+        styled_sofa = Styled image
+    Return:
+        new_image   = Image containing the styled sofa
+    """
     image = np.array(image)
     #image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     styled_sofa = cv2.cvtColor(styled_sofa, cv2.COLOR_BGR2RGB)

styleTransfer.py

@@ -1,4 +1,3 @@
-import torch
 from PIL import Image
 import numpy as np
 import os