Diana1234/Pulmonary_Nodules_Classifier

1import os2import gradio as gr3from fastai.learner import load_learner4from fastMONAI.vision_core import MedImage5from monai.transforms import LoadImage6import torch7import numpy as np  # Ensure numpy is imported8import matplotlib.pyplot as plt  # For visualizing images9 10# Load the trained model and variables11learn = load_learner('learner.pkl', cpu=True)12 13def predict_from_nii(file):14    try:15        # Validate the file type16        if not file.name.endswith(".nii.gz"):17            return "Invalid file type. Please upload a .nii.gz file.", None18 19        # Step 1: Load the image using MONAI's LoadImage20        image_loader = LoadImage(image_only=True)21        image_data = image_loader(file.name)22 23        # Debug: Print the type of image_data24        print(f"Type of image_data: {type(image_data)}")25 26        # Ensure image_data is a numpy array27        if not isinstance(image_data, np.ndarray):28            image_data = np.array(image_data, dtype=np.float32)29        else:30            image_data = image_data.astype(np.float32)31 32        # Step 2: Convert the loaded image data into a tensor33        image_tensor = torch.from_numpy(image_data)34 35        # Adjust the shape to (channels, x, y, z)36        if len(image_tensor.shape) == 3:37            image_tensor = image_tensor.unsqueeze(0)38        elif len(image_tensor.shape) == 2:39            image_tensor = image_tensor.unsqueeze(0).unsqueeze(0)40 41        # Wrap the tensor into MedImage42        med_image = MedImage(image_tensor)43 44        # Step 3: Perform inference45        pred_class, pred_idx, probs = learn.predict(med_image)46 47        # Convert the image to a matplotlib figure for display48        plt.imshow(np.squeeze(image_data), cmap="gray")49        plt.axis("off")50        plt.savefig("temp_image.png")  # Save the image temporarily51        plt.close()52 53        # Format the output54        probs_list = probs.tolist()55        prediction_result = (56            f"Predicted Class: {pred_class}\n"57            f"Class Index: {pred_idx}\n"58            f"Probabilities: {probs_list}"59        )60 61        # Return both prediction text and image path62        return prediction_result, "temp_image.png"63 64    except Exception as e:65        return f"Error processing file: {str(e)}", None66 67# Example folder path68example_folder = "./examples"69 70# Get all .nii.gz files in the example folder71example_files = [os.path.join(example_folder, f) for f in os.listdir(example_folder) if f.endswith('.nii.gz')]72 73interface = gr.Interface(74    fn=predict_from_nii,75    inputs=gr.File(label="Upload a .nii.gz file", file_types=["file"]),76    outputs=[77        gr.Textbox(label="Prediction Results"),  # Text output78        gr.Image(label="CT Scan Image")  # Image output79    ],80    title="Pulmonary Nodules (Lung CT Scans) Classifier",81    description=(82        "Model is trained on the NoduleMNIST dataset, part of the MedMNIST collection, "83        "focused on pulmonary nodules in lung CT scans. Upload a .nii.gz file to classify "84        "the image using a trained model and view the corresponding image."85    ),86    examples=[[file] for file in example_files],  # Dynamically adding examples87)88 89 90 91# Launch the Gradio interface92interface.launch()93