changed model loading

segmentation.py

@@ -7,33 +7,34 @@ import matplotlib.pyplot as plt
 from PIL import Image
 import segmentation_models as sm
 
-model_path = "model_checkpoint.h5"
-CLASSES = ['sofa']
+# 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'
+# 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
+# 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) 
+# #load model
+# model.load_weights(model_path) 
+model = keras.models.load_model('model_final.h5', compile=False)
 
 def get_mask(image:Image) -> Image:
     """