import gradio as gr
from PIL import Image
from fastai.vision.all import *
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.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="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)