Update app.py

app.py

@@ -37,11 +37,12 @@ def FRFsegment(input_img):
     # Model prediction
     est_label_dict = model(img)
 
-    # Extract the actual tensor from the dictionary
-    if isinstance(est_label_dict, dict):
-        est_label = est_label_dict['serving_default']  # Use the appropriate key
-    else:
-        est_label = est_label_dict
+    # Print available keys to understand what we are dealing with
+    print("Available keys in the output dictionary:", est_label_dict.keys())
+
+    # Extract the actual tensor from the dictionary using a dynamic key lookup
+    key = list(est_label_dict.keys())[0]  # Use the first available key
+    est_label = est_label_dict[key]
 
     # Check the shape of the predicted output
     if len(est_label.shape) == 4 and est_label.shape[-1] > 1:
@@ -86,4 +87,3 @@ FRFSegapp = gr.Interface(
 
 FRFSegapp.launch()
 
-