Create app.py

app.py

@@ -0,0 +1,84 @@
+import torch
+
+from torchvision import transforms
+from torch.utils.data import DataLoader, Dataset
+from transformers import MaskFormerForInstanceSegmentation, MaskFormerImageProcessor
+from PIL import Image
+from pathlib import Path
+
+# can upload from Huggingface Space "elliemci/maskformer_tumor_segmentation"
+model = MaskFormerForInstanceSegmentation.from_pretrained("elliemci/maskformer_tumor_segmentation")
+image_processor = MaskFormerImageProcessor.from_pretrained("elliemci/maskformer_tumor_segmentation")
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# Define a custom dataset class to handle images
+class ImageDataset(Dataset):
+    def __init__(self, image_paths, transform=None):
+        self.image_paths = image_paths
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        image = Image.open(self.image_paths[idx]).convert('RGB')
+
+        if self.transform:
+            image = self.transform(image)
+        return image
+
+def segment(image_files):
+    """Takes a list of UploadedFile objects and returns a list of segmented images."""
+
+    dataset = ImageDataset(image_files, transform=transforms.ToTensor())
+    dataloader = DataLoader(dataset, batch_size=len(image_files), shuffle=False)  # Batch size is the number of images
+
+    # process a batch
+    with torch.no_grad():
+        for batch in dataloader:  # Only one iteration since batch_size = len(image_files)
+            pixel_values = batch.to(device, dtype=torch.float32)
+            outputs = model(pixel_values=pixel_values)
+
+            # Post-processing
+            original_images = outputs.get("org_images", batch)
+            target_sizes = [(image.shape[-2], image.shape[-1]) for image in original_images]
+            predicted_masks = image_processor.post_process_semantic_segmentation(outputs, target_sizes=target_sizes)
+
+            return predicted_masks  # Return the list of segmented images
+
+# components for Gradion interface
+def update_gallery(images):
+    print(f"Type in update_gallery: {type(images[0])}")
+    gallery_data = []
+
+    if images:
+        segmented_images = segment(images)  # Process all images in one batch
+
+        for i, image in enumerate(images):
+            segmented_image_pil = transforms.ToPILImage()(segmented_images[i].to(device, dtype=torch.float32))
+            gallery_data.extend([(image, "Original Image"), (segmented_image_pil, "Segmented Image")])
+
+    return gallery_data
+
+# Gradio UI for MEI segmentation
+import gradio as gr
+
+with gr.Blocks() as demo:
+  gr.Markdown("<h1 style='text-align: center;'>MRI Brain Tumor Segmentation App</h1>")
+
+  with gr.Column():
+    with gr.Column():
+      image_files = gr.Files(label="Upload MRI files",
+                                     file_count="multiple",
+                                     type="filepath")
+      with gr.Row():
+        gallery = gr.Gallery(label="Brain Images and Tumor Segmentation")
+
+        image_files.change(
+             fn=update_gallery,
+             inputs=[image_files],
+             outputs=[gallery])
+
+demo.launch()