changed model loading

app.py

@@ -5,6 +5,10 @@ from segmentation import get_mask,replace_sofa
 from styleTransfer import create_styledSofa
 from PIL import Image
 
+#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):
     """
     This function adds padding to make the orignal image square and 640by640.
@@ -130,4 +134,3 @@ if __name__ == "__main__":
     demo.launch()
 
 
-#https://github.com/dhawan98/Post-Processing-of-Image-Segmentation-using-CRF

segmentation.py

@@ -7,6 +7,34 @@ 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
+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)
+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) 
+
 def get_mask(image:Image) -> Image:
     """
     This function generates a mask of the image that highlights all the sofas in the image.
@@ -17,34 +45,6 @@ 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
-    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)
-    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) 
-
     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)