test test test

segmentation.py

@@ -20,37 +20,36 @@ def get_mask(image:Image) -> Image:
             mask  = corresponding maks of the image
     """
 
-        # model_path = "model_checkpoint.h5"
-    # CLASSES = ['sofa']
+    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'
+    n_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)
-    # LR=0.0001
+    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 = 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)
+    model.load_weights(model_path) 
+    #model = keras.models.load_model('model_final.h5', compile=False)
     print('loaded model')
-    test_img = np.array(image)#cv2.imread(path, cv2.IMREAD_COLOR)       
+    test_img = np.array(image) 
     test_img = cv2.resize(test_img, (640, 640))
-    return test_img
     test_img = cv2.cvtColor(test_img, cv2.COLOR_RGB2BGR)
     test_img = np.expand_dims(test_img, axis=0)