Update app.py

app.py

@@ -88,11 +88,12 @@ focal_loss = sm.losses.CategoricalFocalLoss()
 total_loss = dice_loss + (1 * focal_loss)
 
 # ... (previous code)
-satellite_model = tf.saved_model.load("satellite_segmentation_v3_full")
+# ... (previous code)
+satellite_model = tf.keras.models.load_model("satellite_segmentation_v3_full", custom_objects={'dice_loss_plus_1focal_loss': total_loss, 'jaccard_coef': jaccard_coef, 'DiceLoss': sm.losses.DiceLoss, 'CategoricalFocalLoss': sm.losses.CategoricalFocalLoss})
 
 def process_input_image(image_source):
     image = np.expand_dims(image_source, 0)
-    prediction = satellite_model(image) # Change this line
+    prediction = satellite_model.predict(image)
     predicted_image = np.argmax(prediction, axis=3)
     predicted_image = predicted_image[0, :, :]
     rgb_image = label_to_rgb(predicted_image)
@@ -101,6 +102,9 @@ def process_input_image(image_source):
 # ... (rest of the code)
 
 
+# ... (rest of the code)
+
+
 my_app = gr.Blocks()
 
 # Define the custom legend HTML