miyuki17/npy-test-0608 · issue01-目前代码不支持上传npz格式

Absolutely! Here's the enhanced code that gracefully handles .npz files, along with explanations and best practices for medical image segmentation tasks:

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)

Key Improvements:

NPZ Handling: Robustly loads .npz files using np.load and extracts the image array.
File Type Check: Uses file.name.endswith('.npz') for explicit checking of file type.
Dimension Handling: The example assumes the NPZ array is named arr_0. If your NPZ files have different structures, adapt the loading logic. If the channels are not in the last dimension, add a .transpose((1,2,0)) call as shown in the commented line.
Clear Descriptions: Enhances the Gradio interface with more informative titles and descriptions.
Error Handling: (Recommended) Add a try-except block around the image loading logic to provide user-friendly error messages in case of invalid file types or formats.

Example NPZ structure:

np.savez('image_data.npz', arr_0=your_image_array)

Important Considerations:

Medical Data: Medical images often have specific formats (DICOM, NIfTI) and require preprocessing. Ensure your model is trained on data that matches the format and preprocessing you'll use in deployment.
Data Normalization: If your model expects normalized pixel values (e.g., 0-1 range), add normalization steps after loading the image.
Visualizations: For medical image segmentation, consider using gr.Image for both input and output to visually display the original image and the segmentation results.

Let me know if you'd like help with data preprocessing, visualization, or error handling.