Fix model loading: use TFSMLayer for Keras 3 compatibility

TFSMLayer handles legacy SavedModel format that load_model no longer supports in Keras 3.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

app.py

@@ -13,12 +13,18 @@ from itertools import cycle, islice
 
 
 
-#Function that predicts on only 1 sample 
-def predict_sample(image):
-  prediction = model.predict(image[tf.newaxis, ...])
-  prediction[prediction > 0.5 ] = 1
-  prediction[prediction !=1] = 0
-  result = prediction[0]*255
+#Function that predicts on only 1 sample
+def predict_sample(img):
+  input_tensor = tf.cast(img[tf.newaxis, ...], tf.float32)
+  output = model(input_tensor)
+  # TFSMLayer returns a dict; get the first output
+  if isinstance(output, dict):
+    prediction = list(output.values())[0].numpy()
+  else:
+    prediction = output.numpy()
+  prediction[prediction > 0.5] = 1
+  prediction[prediction != 1] = 0
+  result = prediction[0] * 255
   return result
 
 
@@ -45,7 +51,7 @@ def create_input_image(data, visualize=False):
   return input
 
 model_path = snapshot_download("tareknaous/unet-visual-clustering")
-model = tf.keras.models.load_model(model_path)
+model = tf.keras.layers.TFSMLayer(model_path, call_endpoint='serving_default')
 
 
 def get_instances(prediction, data, max_filter_size=1):