fix resize

app.py

@@ -17,13 +17,13 @@ def style_sofa(input_img: np.ndarray, style_img: np.ndarray):
     """
     
     # preprocess input images to be (640,640) squares to fit requirements of the segmentation model
-    resized_img = resize_sofa(input_img)
+    resized_img,box = resize_sofa(input_img)
     resized_style = resize_style(style_img)
     # generate mask for image
     mask = get_mask(resized_img)
     styled_sofa = create_styledSofa(resized_img,resized_style)
-    print('Created styled image')
     new_sofa = replace_sofa(resized_img,mask,styled_sofa)
+    new_sofa = Image.fromarray(new_sofa).crop(box)
     return new_sofa
 
 image = gr.inputs.Image()

segmentation.py

@@ -51,18 +51,14 @@ def replace_sofa(image,mask,styled_sofa):
     # print(mask.shape)
     # mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
     # print(mask.shape)
-    print('started function')
     image = np.array(image)
     #image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     styled_sofa = cv2.cvtColor(styled_sofa, cv2.COLOR_BGR2RGB)
 
-    print('reading masks')
     _, mask = cv2.threshold(mask, 10, 255, cv2.THRESH_BINARY)
     mask_inv = cv2.bitwise_not(mask)
-    print('creating parts')
     image_bg = cv2.bitwise_and(image,image,mask = mask_inv)
     sofa_fg = cv2.bitwise_and(styled_sofa,styled_sofa,mask = mask)
-    print('combining parts')
     new_image = cv2.add(image_bg,sofa_fg)
     return new_image
 

sofaApp.py

@@ -1,50 +0,0 @@
-import numpy as np
-import gradio as gr
-from segmentation import get_mask,replace_sofa
-from styleTransfer import resize_sofa,resize_style,create_styledSofa
-from PIL import Image
-
-def style_sofa(input_img: np.ndarray, style_img: np.ndarray):
-    """
-    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
-    Return:
-        new_sofa  = image containing the styled sofa
-    """
-    
-    # preprocess input images to be (640,640) squares to fit requirements of the segmentation model
-    resized_img = resize_sofa(input_img)
-    resized_style = resize_style(style_img)
-    # generate mask for image
-    mask = get_mask(resized_img)
-    styled_sofa = create_styledSofa(resized_img,resized_style)
-    new_sofa = replace_sofa(resized_img,mask,styled_sofa)
-    return new_sofa
-
-image = gr.inputs.Image()
-style = gr.inputs.Image()
-
-demo = gr.Interface(
-    style_sofa,
-    [image,style],
-    'image',
-    examples=[
-        ['sofa_example1.jpg','style_example1.jpg'],
-        ['sofa_example1.jpg','style_example2.jpg'],
-        ['sofa_example1.jpg','style_example3.jpg'],
-        ['sofa_example1.jpg','style_example4.jpg'],
-        ['sofa_example1.jpg','style_example5.jpg'],
-    ],
-    title="Style your sofa",
-    description="🛋 Customize your sofa to your wildest dreams! 🛋",
-)
-
-if __name__ == "__main__":
-    demo.launch(share=True)
-
-
-#https://github.com/dhawan98/Post-Processing-of-Image-Segmentation-using-CRF

styleTransfer.py

@@ -8,17 +8,19 @@ def resize_sofa(img):
     img = Image.fromarray(img)
     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 = ((height-width)//2*newsize[0]/width, 0, newsize[0] - (height-width)//2*newsize[0]/width, newsize[1])
     else:
         img1 = Image.new(img.mode, (width, width), (255, 255, 255))
         img1.paste(img, (0, (width-height)//2))
+        box = (0, (height-width)//2*newsize[1]/height, newsize[0], newsize[1]-(height-width)//2*newsize[1]/height)
 
-    newsize = (640, 640) # parameters from test script
     im1 = img1.resize(newsize)
-    return im1
+    return im1,box
 
 def resize_style(img):
     #img = Image.open(path)#"../style5.jpg")
@@ -68,6 +70,10 @@ def create_styledSofa(sofa,style):
 
     return styled_sofa
 
-# image = Image.open('input/sofa.jpg')  
-# image = np.array(image)[:,:600]
-# image = resize_sofa(image)  
+# image = Image.open('sofa_office.jpg') 
+# image.show() 
+# image = np.array(image)
+# image,box = resize_sofa(image)  
+# image = image.crop(box)
+# print(box)
+# image.show()

test.py

@@ -148,7 +148,6 @@ style_tf = test_transform(style_size, crop)
 
 for content_path in content_paths:
     for style_path in style_paths:
-        print(content_path)
        
       
         content_tf1 = content_transform()