Update app.py

app.py

@@ -1,32 +1,44 @@
 import gradio as gr
 from PIL import Image
 from fastai.vision.all import *
-
-def predict_image(image):
-#    learner = load_learner('model-epoch=119.ckpt')
-    learner = load_learner('export_2.pkl')
-    img = PILImage.create(image)
+import numpy as np
+
+# Model Loading (Choose ONE)
+# learner = load_learner('export_2.pkl')  # If using exported pickle model
+learner = load_learner('model-epoch=119.ckpt')  # If using checkpoint
+
+def predict_image(file):
+    if file.name.endswith('.npz'):
+        # Load NPZ data
+        with np.load(file.name) as data:
+            # Assuming your NPZ file contains an array named 'arr_0'
+            img_array = data['arr_0']  
+            # Ensure correct dimensions (e.g., channels last)
+            # img_array = img_array.transpose((1, 2, 0))  # If needed
+            img = Image.fromarray(img_array)
+    else:
+        # Load regular image files (JPG, PNG, etc.)
+        img = PILImage.create(file.name)
     pred = learner.predict(img)
     return pred[0]
 
 
 def create_interface():
-    
-    image_input = gr.Image()
-
+    image_input = gr.File(
+        label="Upload Image or NPZ",
+        file_types=[".jpg", ".jpeg", ".png", ".npz"]
+    )
     output = gr.Text()
 
     iface = gr.Interface(
         fn=predict_image,
         inputs=image_input,
         outputs=output,
-        title="npy seg",
-        description="Upload an image to seg."
+        title="Medical Image Segmentation",  # More informative title
+        description="Upload an image or NPZ file to segment."  # Clear description
     )
     return iface
 
 if __name__ == "__main__":
     iface = create_interface()
     iface.launch(share=True)
-
-