Restructure code and fix error

.gitattributes

@@ -25,5 +25,3 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-
-*.pth filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,8 +0,0 @@
-__pycache__/box_ops.cpython-37.pyc
-__pycache__/function.cpython-37.pyc
-__pycache__/misc.cpython-37.pyc
-__pycache__/segmentation.cpython-37.pyc
-__pycache__/styleTransfer.cpython-37.pyc
-__pycache__/StyTR.cpython-37.pyc
-__pycache__/transformer.cpython-37.pyc
-__pycache__/ViT_helper.cpython-37.pyc

README.md

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

segmentation.py → Segmentation/segmentation.py

@@ -8,6 +8,30 @@ from PIL import Image
 import segmentation_models as sm
 sm.set_framework('tf.keras')
 
+model_path = "Segmentation/model_checkpoint.h5"
+CLASSES = ['sofa']
+BACKBONE = 'resnet50'
+
+# define network parameters
+n_classes = 1 if len(CLASSES) == 1 else (len(CLASSES) + 1)  # case for binary and multiclass segmentation
+activation = 'sigmoid' if n_classes == 1 else 'softmax'
+preprocess_input = sm.get_preprocessing(BACKBONE)
+LR=0.0001
+
+#create model architecture
+model = sm.Unet(BACKBONE, classes=n_classes, activation=activation)
+# define optomizer
+optim = keras.optimizers.Adam(LR)
+# Segmentation models losses can be combined together by '+' and scaled by integer or float factor
+dice_loss = sm.losses.DiceLoss()
+focal_loss = sm.losses.BinaryFocalLoss() if n_classes == 1 else sm.losses.CategoricalFocalLoss()
+total_loss = dice_loss + (1 * focal_loss)
+# actulally total_loss can be imported directly from library, above example just show you how to manipulate with losses
+# 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
+model.compile(optim, total_loss, metrics)
+model.load_weights(model_path) 
 
 def get_mask(image:Image) -> Image:
     """
@@ -20,35 +44,10 @@ def get_mask(image:Image) -> Image:
             mask  = corresponding maks of the image
     """
 
-    model_path = "model_checkpoint.h5"
-    CLASSES = ['sofa']
-    BACKBONE = 'resnet50'
-
-    # define network parameters
-    n_classes = 1 if len(CLASSES) == 1 else (len(CLASSES) + 1)  # case for binary and multiclass segmentation
-    activation = 'sigmoid' if n_classes == 1 else 'softmax'
-    preprocess_input = sm.get_preprocessing(BACKBONE)
-    LR=0.0001
-
-    #create model architecture
-    model = sm.Unet(BACKBONE, classes=n_classes, activation=activation)
-    # define optomizer
-    optim = keras.optimizers.Adam(LR)
-    # Segmentation models losses can be combined together by '+' and scaled by integer or float factor
-    dice_loss = sm.losses.DiceLoss()
-    focal_loss = sm.losses.BinaryFocalLoss() if n_classes == 1 else sm.losses.CategoricalFocalLoss()
-    total_loss = dice_loss + (1 * focal_loss)
-    # actulally total_loss can be imported directly from library, above example just show you how to manipulate with losses
-    # 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
-    model.compile(optim, total_loss, metrics)
-
     # #load model
-    model.load_weights(model_path) 
+
     #model = keras.models.load_model('model_final.h5', compile=False)
     print('loaded model')
-    return image
     test_img = np.array(image) 
     test_img = cv2.resize(test_img, (640, 640))
     test_img = cv2.cvtColor(test_img, cv2.COLOR_RGB2BGR)

StyTR.py → StyleTransfer/StyTR.py

@@ -1,15 +1,12 @@
 import torch
 import torch.nn.functional as F
 from torch import nn
-import numpy as np
-import box_ops
-from misc import (NestedTensor, nested_tensor_from_tensor_list,
+from StyleTransfer.misc import (NestedTensor, nested_tensor_from_tensor_list,
                        accuracy, get_world_size, interpolate,
                        is_dist_avail_and_initialized)
-from function import normal,normal_style
-from function import calc_mean_std
-import scipy.stats as stats
-from ViT_helper import DropPath, to_2tuple, trunc_normal_
+from StyleTransfer.function import normal,normal_style
+from StyleTransfer.function import calc_mean_std
+from StyleTransfer.ViT_helper import DropPath, to_2tuple, trunc_normal_
 
 class PatchEmbed(nn.Module):
     """ Image to Patch Embedding

styleTransfer.py → StyleTransfer/styleTransfer.py

@@ -3,13 +3,11 @@ import numpy as np
 import torch
 import torch.nn as nn
 from torchvision import transforms
-import transformer as transformer
-import StyTR as StyTR
+import StyleTransfer.transformer as transformer
+import StyleTransfer.StyTR as StyTR
 from collections import OrderedDict
 import tensorflow_hub as tfhub
 import tensorflow as tf
-import os
-import cv2 
 import paddlehub as phub
 
 
@@ -30,10 +28,10 @@ def content_transform():
     return transform
 
 def StyleTransformer(content_img: Image, style_img: Image):
-    vgg_path        = 'vgg_normalised.pth'
-    decoder_path    = 'decoder_iter_160000.pth'
-    Trans_path      = 'transformer_iter_160000.pth'
-    embedding_path  = 'embedding_iter_160000.pth'
+    vgg_path        = 'StyleTransfer/models/vgg_normalised.pth'
+    decoder_path    = 'StyleTransfer/models/decoder_iter_160000.pth'
+    Trans_path      = 'StyleTransfer/models/transformer_iter_160000.pth'
+    embedding_path  = 'StyleTransfer/models/embedding_iter_160000.pth'
     # Advanced options
     content_size=640
     style_size=640

transformer.py → StyleTransfer/transformer.py

@@ -4,7 +4,7 @@ from typing import Optional, List
 import torch
 import torch.nn.functional as F
 from torch import nn, Tensor
-from function import normal,normal_style
+from StyleTransfer.function import normal,normal_style
 import numpy as np
 import os
 device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")

app.py

@@ -1,11 +1,12 @@
 from cv2 import transpose
 import numpy as np
 import gradio as gr
-from segmentation import get_mask,replace_sofa
-#from styleTransfer import create_styledSofa
+from Segmentation.segmentation import get_mask,replace_sofa
+from StyleTransfer.styleTransfer import create_styledSofa
 from PIL import Image
 from random import randint
 
+
 #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
@@ -102,25 +103,25 @@ def style_sofa(input_img: np.ndarray, style_img: np.ndarray):
     mask = get_mask(resized_img)
     #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')
-    # # postprocess the final image
-    # print('Replacing sofa...')
-    # new_sofa = replace_sofa(resized_img,mask,styled_sofa)
-    # new_sofa = new_sofa.crop(box)
+    print('Styling sofa...')
+    styled_sofa = create_styledSofa(resized_img,resized_style)
+    #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 mask
+    return new_sofa
 
 demo = gr.Interface(
     style_sofa,
     inputs = [gr.inputs.Image(),gr.inputs.Image()],
     outputs = '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']],
+    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']],
     title="🛋 Style your sofa 🛋 ",
     description="Customize your sofa to your wildest dreams 💭!\
                 \nProvide a picture of your sofa and a desired pattern\

box_ops.py

@@ -1,88 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
-"""
-Utilities for bounding box manipulation and GIoU.
-"""
-import torch
-from torchvision.ops.boxes import box_area
-
-
-def box_cxcywh_to_xyxy(x):
-    x_c, y_c, w, h = x.unbind(-1)
-    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
-         (x_c + 0.5 * w), (y_c + 0.5 * h)]
-    return torch.stack(b, dim=-1)
-
-
-def box_xyxy_to_cxcywh(x):
-    x0, y0, x1, y1 = x.unbind(-1)
-    b = [(x0 + x1) / 2, (y0 + y1) / 2,
-         (x1 - x0), (y1 - y0)]
-    return torch.stack(b, dim=-1)
-
-
-# modified from torchvision to also return the union
-def box_iou(boxes1, boxes2):
-    area1 = box_area(boxes1)
-    area2 = box_area(boxes2)
-
-    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
-    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]
-
-    wh = (rb - lt).clamp(min=0)  # [N,M,2]
-    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]
-
-    union = area1[:, None] + area2 - inter
-
-    iou = inter / union
-    return iou, union
-
-
-def generalized_box_iou(boxes1, boxes2):
-    """
-    Generalized IoU from https://giou.stanford.edu/
-
-    The boxes should be in [x0, y0, x1, y1] format
-
-    Returns a [N, M] pairwise matrix, where N = len(boxes1)
-    and M = len(boxes2)
-    """
-    # degenerate boxes gives inf / nan results
-    # so do an early check
-    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
-    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
-    iou, union = box_iou(boxes1, boxes2)
-
-    lt = torch.min(boxes1[:, None, :2], boxes2[:, :2])
-    rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])
-
-    wh = (rb - lt).clamp(min=0)  # [N,M,2]
-    area = wh[:, :, 0] * wh[:, :, 1]
-
-    return iou - (area - union) / area
-
-
-def masks_to_boxes(masks):
-    """Compute the bounding boxes around the provided masks
-
-    The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions.
-
-    Returns a [N, 4] tensors, with the boxes in xyxy format
-    """
-    if masks.numel() == 0:
-        return torch.zeros((0, 4), device=masks.device)
-
-    h, w = masks.shape[-2:]
-
-    y = torch.arange(0, h, dtype=torch.float)
-    x = torch.arange(0, w, dtype=torch.float)
-    y, x = torch.meshgrid(y, x)
-
-    x_mask = (masks * x.unsqueeze(0))
-    x_max = x_mask.flatten(1).max(-1)[0]
-    x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]
-
-    y_mask = (masks * y.unsqueeze(0))
-    y_max = y_mask.flatten(1).max(-1)[0]
-    y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]
-
-    return torch.stack([x_min, y_min, x_max, y_max], 1)

model_final.h5

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:9a456f38c83897d9d8b5c8dd989ff7ee2fe13bb123a70a00b6e987d4efac1c6e
-size 130858696