fix error

segmentation.py

@@ -7,34 +7,6 @@ import matplotlib.pyplot as plt
 from PIL import Image
 import segmentation_models as sm
 
-# model_path = "model_checkpoint.h5"
-# CLASSES = ['sofa']
-BACKBONE = 'resnet50'
-
-# define network parameters
-# in_classes = 1 if len(CLASSES) == 1 else (len(CLASSES) + 1)  # case for binary and multiclass segmentation
-# actvation = 'sigmoid' if n_classes == 1 else 'softmax'
-preprocess_input = sm.get_preprocessing(BACKBONE)
-sm.set_framework('tf.keras')
-# 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)
 
 def get_mask(image:Image) -> Image:
     """
@@ -46,6 +18,36 @@ def get_mask(image:Image) -> Image:
     Return:
             mask  = corresponding maks of the image
     """
+
+        # model_path = "model_checkpoint.h5"
+    # CLASSES = ['sofa']
+    BACKBONE = 'resnet50'
+
+    # define network parameters
+    # in_classes = 1 if len(CLASSES) == 1 else (len(CLASSES) + 1)  # case for binary and multiclass segmentation
+    # actvation = 'sigmoid' if n_classes == 1 else 'softmax'
+    preprocess_input = sm.get_preprocessing(BACKBONE)
+    sm.set_framework('tf.keras')
+    # 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)
+
     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)

styleTransfer.py

@@ -34,19 +34,20 @@ def content_transform():
     transform = transforms.Compose(transform_list)
     return transform
 
-# Advanced options
-content_size=640
-style_size=640
 
-vgg = StyTR.vgg
-vgg.load_state_dict(torch.load(vgg_path))
-vgg = nn.Sequential(*list(vgg.children())[:44])
+def StyleTransformer(content_img: Image, style_img: Image):
 
-decoder = StyTR.decoder
-Trans = transformer.Transformer()
-embedding = StyTR.PatchEmbed()
+    # Advanced options
+    content_size=640
+    style_size=640
 
-def StyleTransformer(content_img: Image, style_img: Image):
+    vgg = StyTR.vgg
+    vgg.load_state_dict(torch.load(vgg_path))
+    vgg = nn.Sequential(*list(vgg.children())[:44])
+
+    decoder = StyTR.decoder
+    Trans = transformer.Transformer()
+    embedding = StyTR.PatchEmbed()
 
     decoder.eval()
     Trans.eval()
@@ -105,6 +106,11 @@ def styleProjection(content_image,style_image):
     }])
 
     return Image.fromarray(np.uint8(result[0]['data'])[:,:,::-1]).convert('RGB')
+
+
+def create_styledSofa(content_image,style_image):
+    output = StyleTransformer(content_image,style_image)
+    return output