Update app.py

app.py

@@ -190,6 +190,45 @@ class readDataset:
             np.array(train_sar_aug), np.array(train_optic_aug), np.array(train_masks_aug),
             np.array(test_sar), np.array(test_optic), np.array(test_masks)
         )
+
+@tf.keras.saving.register_keras_serializable()
+def dice_score(y_true, y_pred, threshold=0.5, smooth=1.0):
+    #determine binary or multiclass segmentation
+    is_multiclass = y_true.shape[-1] > 1
+    
+    if not is_multiclass:
+        # Binary segmentation
+        y_true_flat = tf.cast(tf.reshape(y_true, [-1]), dtype=tf.float32)
+        y_pred_flat = tf.cast(tf.reshape(y_pred >= threshold, [-1]), dtype=tf.float32)
+        intersection = tf.reduce_sum(y_true_flat * y_pred_flat)
+        score = (2. * intersection + smooth) / (tf.reduce_sum(y_true_flat) + tf.reduce_sum(y_pred_flat) + smooth) 
+        return score
+    else:
+        # Multiclass segmentation
+        num_classes = y_true.shape[-1]
+        score_per_class = []
+        
+        for i in range(num_classes):
+            y_true_flat = tf.cast(tf.reshape(y_true, [-1]), dtype=tf.float32)
+            y_pred_flat = tf.cast(tf.reshape(y_pred >= threshold, [-1]), dtype=tf.float32)
+            intersection = tf.reduce_sum(y_true_flat * y_pred_flat)
+            score = (2. * intersection + smooth) / (tf.reduce_sum(y_true_flat) + tf.reduce_sum(y_pred_flat) + smooth)
+            score_per_class.append(score)
+        
+        return tf.reduce_mean(score_per_class)
+    
+@tf.keras.saving.register_keras_serializable() 
+def dice_loss(y_true, y_pred):
+    dice = dice_score(y_true, y_pred)
+    loss = 1. - dice
+    return tf.cast(loss, dtype=tf.float32)
+
+@tf.keras.saving.register_keras_serializable()
+def cce_dice_loss(y_true, y_pred):
+    cce = tf.keras.losses.CategoricalCrossentropy()(y_true, y_pred)
+    dice = dice_loss(y_true, y_pred)
+    return tf.cast(cce, dtype=tf.float32) + dice
+    
 # Streamlit App Title
 st.title("Satellite Mining Segmentation: SAR + Optic Image Inference")
 
@@ -224,7 +263,8 @@ if st.button("Run Inference"):
         optic_images = dataset.normalizeImages(optic_images, 'i')
 
         # Load model
-        model = load_model(model_path)
+        model = tf.keras.models.load_model(model_path,
+                                  custom_objects={"cce_dice_loss": cce_dice_loss, "dice_score": dice_score})
 
         # Predict
         pred_masks = model.predict([optic_images, sar_images], verbose=0)