Add files for the gradio app

__init__.py

@@ -0,0 +1 @@
+from . import data, models, tools, visualization

app.py

@@ -0,0 +1,49 @@
+import os
+import gradio as gr
+from tools.predict import single_prediction
+
+KAGGLE_NOTEBOOK = "[![Static Badge](https://img.shields.io/badge/Open_Notebook_in_Kaggle-gray?logo=kaggle&logoColor=white&labelColor=20BEFF)](https://www.kaggle.com/code/mmenendezg/mobilevit-fluorescent-neuronal-cells/notebook)"
+GITHUB_REPOSITORY = "[![Static Badge](https://img.shields.io/badge/Git_Repository-gray?logo=github&logoColor=white&labelColor=181717)](https://github.com/mmenendezg/mobilevit-fluorescent-cells)"
+HF_SPACE = "[![Open in Spaces](https://huggingface.co/datasets/huggingface/badges/resolve/main/open-in-hf-spaces-md-dark.svg)](https://huggingface.co/spaces/mmenendezg/mobilevit-fluorescent-neuronal-cells)"
+
+# Gradio interface
+demo = gr.Blocks()
+with demo:
+    gr.Markdown(
+        f"""
+    # Fluorescent Neuronal Cells Segmentation
+
+    This model extracts a segmentation mask of the neuronal cells on an image.
+
+    {KAGGLE_NOTEBOOK}
+
+    {GITHUB_REPOSITORY}
+
+    {HF_SPACE}
+    """
+    )
+    with gr.Tab("Image Segmentation"):
+        with gr.Row():
+            with gr.Column():
+                uploaded_image = gr.Image(
+                    label="Neuronal Cells Image",
+                    sources=["upload", "clipboard"],
+                    type="pil",
+                    height=550,
+                )
+            with gr.Column():
+                mask_image = gr.Image(label="Segmented Neurons", height=550)
+        with gr.Row():
+            classify_btn = gr.Button("Segment the image", variant="primary")
+            clear_btn = gr.ClearButton(components=[uploaded_image, mask_image])
+        classify_btn.click(
+            fn=single_prediction, inputs=uploaded_image, outputs=[mask_image]
+        )
+        gr.Examples(
+            examples=[
+                os.path.join(os.path.dirname(__file__), "examples/example_1.png"),
+                os.path.join(os.path.dirname(__file__), "examples/example_2.png"),
+            ],
+            inputs=uploaded_image,
+        )
+demo.launch(show_error=True)

config/fluorescent_mobilevit_hps.yaml

@@ -0,0 +1,3 @@
+learning_rate: 0.0005480015685663855
+weight_decay: 1.544480236681167e-05
+batch_size: 2

data/__init__.py

@@ -0,0 +1 @@
+from . import data_preprocessing

data/data_preprocessing.py

@@ -0,0 +1,155 @@
+import os
+import numpy as np
+from PIL import Image
+import cv2
+import torch
+from torch.utils.data import DataLoader
+import albumentations as A
+import pytorch_lightning as pl
+from transformers import AutoImageProcessor
+from datasets import Dataset, DatasetDict
+
+# Checkpoint of the model used in the projec
+MODEL_CHECKPOINT = "apple/deeplabv3-mobilevit-xx-small"
+# Size of the image used to train the model
+IMG_SIZE = [256, 256]
+
+
+class FluorescentNeuronalDataModule(pl.LightningDataModule):
+    def __init__(self, batch_size, data_dir, dataset_size=1.0):
+        super().__init__()
+        self.data_dir = data_dir
+        self.batch_size = batch_size
+        self.image_processor = AutoImageProcessor.from_pretrained(
+            MODEL_CHECKPOINT, do_reduce_labels=False
+        )
+        self.image_resizer = A.Compose(
+            [
+                A.Resize(
+                    width=IMG_SIZE[0],
+                    height=IMG_SIZE[1],
+                    interpolation=cv2.INTER_NEAREST,
+                )
+            ]
+        )
+        self.image_augmentator = A.Compose(
+            [
+                A.HorizontalFlip(p=0.6),
+                A.VerticalFlip(p=0.6),
+                A.RandomBrightnessContrast(p=0.6),
+                A.RandomGamma(p=0.6),
+                A.HueSaturationValue(p=0.6),
+            ]
+        )
+
+        # Percentage of the dataset
+        self.dataset_size = dataset_size
+
+    def _create_dataset(self):
+        images_path = os.path.join(self.data_dir, "all_images", "images")
+        masks_path = os.path.join(self.data_dir, "all_masks", "masks")
+        list_images = os.listdir(images_path)
+
+        # Determine the size of the dataset
+        if self.dataset_size < 1.0:
+            n_images = int(len(list_images) * self.dataset_size)
+            list_images = list_images[:n_images]
+
+        images = []
+        masks = []
+        for image_filename in list_images:
+            image_path = os.path.join(images_path, image_filename)
+            mask_path = os.path.join(masks_path, image_filename)
+
+            image = np.array(Image.open(image_path).convert("RGB"), dtype=np.uint8)
+            mask = np.array(Image.open(mask_path).convert("L"), dtype=np.uint8)
+            mask = (mask / 255).astype(np.uint8)
+
+            images.append(image)
+            masks.append(mask)
+
+        dataset = Dataset.from_dict({"image": images, "mask": masks})
+
+        # Split the dataset into train, val, and test sets
+        dataset = dataset.train_test_split(test_size=0.1)
+        train_val = dataset["train"]
+        test_ds = dataset["test"]
+        del dataset
+
+        train_val = train_val.train_test_split(test_size=0.2)
+        train_ds = train_val["train"]
+        valid_ds = train_val["test"]
+        del train_val
+
+        dataset = DatasetDict(
+            {"train": train_ds, "validation": valid_ds, "test": test_ds}
+        )
+        del train_ds, valid_ds, test_ds
+        return dataset
+
+    def _transform_train_data(self, batch):
+        # Preprocess the images
+        images, masks = [], []
+        for i, m in zip(batch["image"], batch["mask"]):
+            img = np.asarray(i, dtype=np.uint8)
+            mask = np.asarray(m, dtype=np.uint8)
+            # First resize the images and masks
+            resized_outputs = self.image_resizer(image=img, mask=mask)
+            images.append(resized_outputs["image"])
+            masks.append(resized_outputs["mask"])
+
+            # Then augment the images
+            augmented_outputs = self.image_augmentator(
+                image=resized_outputs["image"], mask=resized_outputs["mask"]
+            )
+            images.append(augmented_outputs["image"])
+            masks.append(augmented_outputs["mask"])
+
+        inputs = self.image_processor(
+            images=images,
+            return_tensors="pt",
+        )
+        inputs["labels"] = torch.tensor(masks, dtype=torch.long)
+        return inputs
+
+    def _transform_data(self, batch):
+        # Preprocess the images
+        images, masks = [], []
+        for i, m in zip(batch["image"], batch["mask"]):
+            img = np.asarray(i, dtype=np.uint8)
+            mask = np.asarray(m, dtype=np.uint8)
+            # Resize the images and masks
+            resized_outputs = self.image_resizer(image=img, mask=mask)
+            images.append(resized_outputs["image"])
+            masks.append(resized_outputs["mask"])
+
+        inputs = self.image_processor(
+            images=images,
+            return_tensors="pt",
+        )
+        inputs["labels"] = inputs["labels"] = torch.tensor(masks, dtype=torch.long)
+        return inputs
+
+    def setup(self, stage=None):
+        dataset = self._create_dataset()
+        train_ds = dataset["train"]
+        valid_ds = dataset["validation"]
+        test_ds = dataset["test"]
+
+        if stage is None or stage == "fit":
+            self.train_ds = train_ds.with_transform(self._transform_train_data)
+            self.valid_ds = valid_ds.with_transform(self._transform_data)
+        if stage is None or stage == "test" or stage == "predict":
+            self.test_ds = test_ds.with_transform(self._transform_data)
+
+    def train_dataloader(self):
+        return DataLoader(self.train_ds, batch_size=self.batch_size, shuffle=True)
+
+    def val_dataloader(self):
+        return DataLoader(self.valid_ds, batch_size=self.batch_size)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_ds, batch_size=self.batch_size)
+
+    def predict_dataloader(self):
+        return DataLoader(self.test_ds, batch_size=self.batch_size)

models/__init__.py

@@ -0,0 +1 @@
+from . import mobilevit

models/mobilevit.py

@@ -0,0 +1,120 @@
+import numpy as np
+import torch
+from torch import nn
+import pytorch_lightning as pl
+from transformers import MobileViTForSemanticSegmentation
+import evaluate
+
+MODEL_CHECKPOINT = "mmenendezg/mobilevit-fluorescent-neuronal-cells"
+CLASSES = {0: "Background", 1: "Neuron"}
+
+
+class MobileVIT(pl.LightningModule):
+    def __init__(self, learning_rate=None, weight_decay=None):
+        super().__init__()
+        self.id2label = CLASSES
+        self.label2id = {v: k for k, v in self.id2label.items()}
+        self.num_classes = len(self.id2label.keys())
+        self.model = MobileViTForSemanticSegmentation.from_pretrained(
+            MODEL_CHECKPOINT,
+            num_labels=self.num_classes,
+            id2label=self.id2label,
+            label2id=self.label2id,
+            ignore_mismatched_sizes=True,
+        )
+        self.metric = evaluate.load("mean_iou")
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+
+    def forward(self, pixel_values, labels):
+        return self.model(pixel_values=pixel_values, labels=labels)
+
+    def common_step(self, batch, batch_idx):
+        pixel_values = batch["pixel_values"]
+        labels = batch["labels"]
+
+        outputs = self.model(pixel_values=pixel_values, labels=labels)
+
+        loss = outputs.loss
+        logits = outputs.logits
+        return loss, logits
+
+    def compute_metric(self, logits, labels):
+        logits_tensor = nn.functional.interpolate(
+            logits,
+            size=labels.shape[-2:],
+            mode="bilinear",
+            align_corners=False,
+        ).argmax(dim=1)
+        pred_labels = logits_tensor.detach().cpu().numpy()
+        metrics = self.metric.compute(
+            predictions=pred_labels,
+            references=labels,
+            num_labels=self.num_classes,
+            ignore_index=255,
+            reduce_labels=False,
+        )
+
+        return metrics
+
+    def training_step(self, batch, batch_idx):
+        labels = batch["labels"]
+
+        # Calculate and log the loss
+        loss, logits = self.common_step(batch, batch_idx)
+        self.log("train_loss", loss)
+
+        # Calculate and log the metrics
+        metrics = self.compute_metric(logits, labels)
+        metrics = {key: np.float32(value) for key, value in metrics.items()}
+
+        self.log("train_mean_iou", metrics["mean_iou"])
+        self.log("train_mean_accuracy", metrics["mean_accuracy"])
+        self.log("train_overall_accuracy", metrics["overall_accuracy"])
+
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        labels = batch["labels"]
+
+        # Calculate and log the loss
+        loss, logits = self.common_step(batch, batch_idx)
+        self.log("val_loss", loss)
+
+        # Calculate and log the metrics
+        metrics = self.compute_metric(logits, labels)
+        metrics = {key: np.float32(value) for key, value in metrics.items()}
+        self.log("val_mean_iou", metrics["mean_iou"])
+        self.log("val_mean_accuracy", metrics["mean_accuracy"])
+        self.log("val_overall_accuracy", metrics["overall_accuracy"])
+
+        return loss
+
+    def test_step(self, batch, batch_idx):
+        labels = batch["labels"]
+
+        # Calculate and log the loss
+        loss, logits = self.common_step(batch, batch_idx)
+        self.log("test_loss", loss)
+
+        # Calculate and log the metrics
+        metrics = self.compute_metric(logits, labels)
+        metrics = {key: np.float32(value) for key, value in metrics.items()}
+        # for k, v in metrics.items():
+        #     self.log(f"val_{k}", v.item())
+        self.log("test_mean_iou", metrics["mean_iou"])
+        self.log("test_mean_accuracy", metrics["mean_accuracy"])
+        self.log("test_overall_accuracy", metrics["overall_accuracy"])
+
+        return loss
+
+    def configure_optimizers(self):
+        param_dicts = [
+            {
+                "params": [p for n, p in self.named_parameters()],
+                "lr": self.learning_rate,
+            }
+        ]
+        return torch.optim.AdamW(
+            param_dicts, lr=self.learning_rate, weight_decay=self.weight_decay
+        )

tools/__init__.py

@@ -0,0 +1 @@
+from . import hyperparameters_tuning, train_model

tools/hyperparameters_tuning.py

@@ -0,0 +1,104 @@
+import os
+from pathlib import Path
+import yaml
+import torch
+import optuna
+import pytorch_lightning as pl
+import click
+from pytorch_lightning.callbacks.early_stopping import EarlyStopping
+
+from models.mobilevit import MobileVIT
+from data.data_preprocessing import FluorescentNeuronalDataModule
+
+
+MODEL_CHECKPOINT = "apple/deeplabv3-mobilevit-xx-small"
+
+# Define the accelerator
+if torch.backends.mps.is_available():
+    DEVICE = torch.device("mps:0")
+    ACCELERATOR = "mps"
+elif torch.cuda.is_available():
+    DEVICE = torch.device("cuda")
+    ACCELERATOR = "gpu"
+else:
+    DEVICE = torch.device("cpu")
+    ACCELERATOR = "cpu"
+
+RAW_DATA_PATH = "./data/raw/"
+DEFAULT_CONFIG_FILE = "./config/fluorescent_mobilevit_hps.yaml"
+
+CLASSES = {0: "Background", 1: "Neuron"}
+
+IMG_SIZE = [256, 256]
+
+
+@click.command()
+@click.option(
+    "--data_dir",
+    type=click.Path(exists=True, file_okay=True, path_type=Path),
+    default=RAW_DATA_PATH,
+)
+@click.option(
+    "--config_file",
+    type=click.Path(exists=True, file_okay=True, path_type=Path),
+    default=DEFAULT_CONFIG_FILE,
+)
+@click.option("--dataset_size", type=click.FLOAT, default=0.25)
+@click.option("--force-tune/--no-force-tune", default=False)
+def get_best_params(data_dir, config_file, dataset_size, force_tune) -> dict:
+    def objective(trial: optuna.Trial, dataset_size=dataset_size) -> float:
+        # Suggest values of the hyperparameters for the trials
+        learning_rate = trial.suggest_float("learning_rate", 1e-6, 1e-3, log=True)
+        weight_decay = trial.suggest_float("weight_decay", 1e-6, 1e-3, log=True)
+        batch_size = trial.suggest_int("batch_size", 2, 4, log=True)
+
+        # Define the callbacks of the model
+        early_stopping_cb = EarlyStopping(monitor="val_loss", patience=2)
+
+        # Create the model
+        model = MobileVIT(learning_rate=learning_rate, weight_decay=weight_decay)
+
+        # Instantiate the data module
+        data_module = FluorescentNeuronalDataModule(
+            batch_size=batch_size, dataset_size=dataset_size, data_dir=data_dir
+        )
+        data_module.setup()
+
+        # Train the model
+        trainer = pl.Trainer(
+            devices=1,
+            accelerator=ACCELERATOR,
+            precision="16-mixed",
+            max_epochs=5,
+            log_every_n_steps=5,
+            callbacks=[early_stopping_cb],
+        )
+        trainer.fit(
+            model,
+            train_dataloaders=data_module.train_dataloader(),
+            val_dataloaders=data_module.val_dataloader(),
+        )
+        return trainer.callback_metrics["val_loss"].item()
+
+    if os.path.exists(config_file) and force_tune:
+        os.remove(config_file)
+        pruner = optuna.pruners.MedianPruner()
+        study = optuna.create_study(direction="maximize", pruner=pruner)
+
+        study.optimize(objective, n_trials=25)
+        best_params = study.best_params
+        with open(config_file, "w") as file:
+            yaml.dump(best_params, file)
+    elif os.path.exists(config_file):
+        with open(config_file, "r") as file:
+            best_params = yaml.safe_load(file)
+    else:
+        pruner = optuna.pruners.MedianPruner()
+        study = optuna.create_study(direction="minimize", pruner=pruner)
+
+        study.optimize(objective, n_trials=25)
+        best_params = study.best_params
+        with open(config_file, "w") as file:
+            yaml.dump(best_params, file)
+
+    click.echo(f"The best parameters are:\n{best_params}")

tools/predict.py

@@ -0,0 +1,47 @@
+import numpy as np
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor
+
+from models.mobilevit import MobileVIT
+
+# Checkpoint of the model used in the projec
+MODEL_CHECKPOINT = "mmenendezg/mobilevit-fluorescent-neuronal-cells"
+
+# Define the accelerator
+if torch.backends.mps.is_available():
+    DEVICE = torch.device("mps:0")
+    ACCELERATOR = "mps"
+elif torch.cuda.is_available():
+    DEVICE = torch.device("cuda")
+    ACCELERATOR = "gpu"
+else:
+    DEVICE = torch.device("cpu")
+    ACCELERATOR = "cpu"
+
+
+def single_prediction(image):
+    # Instantiate the model from the checkpoint and using the hparams file
+    mobilevit_model = MobileVIT()
+    mobilevit_model.to(DEVICE)
+    # Instantiate the image_processor
+    image_processor = AutoImageProcessor.from_pretrained(
+        MODEL_CHECKPOINT, do_reduce_labels=False
+    )
+    # Load the image
+    image = image.convert("RGB")
+    # Convert the image to numpy array
+    np_image = np.asarray(image, dtype=np.uint8)
+    # Preprocess the image and move the image to the GPU Device
+    processed_image = image_processor(images=np_image, return_tensors="pt")
+    processed_image.to(DEVICE)
+    # Make the prediction and resize the predicted mask
+    logits = mobilevit_model.model(pixel_values=processed_image["pixel_values"])
+    post_processed_image = image_processor.post_process_semantic_segmentation(
+        outputs=logits, target_sizes=[(np_image.shape[0], np_image.shape[1])]
+    )
+    # Process the mask
+    mask = post_processed_image[0].data.cpu().numpy().astype(np.uint8) * 255
+    mask = Image.fromarray(mask)
+
+    return mask

tools/train_model.py

@@ -0,0 +1,70 @@
+import yaml
+from pathlib import Path
+import click
+import torch
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
+from pytorch_lightning.loggers import TensorBoardLogger
+
+from models.mobilevit import MobileVIT
+from data.data_preprocessing import FluorescentNeuronalDataModule
+
+CONFIG_FILE = "config/fluorescent_mobilevit_hps.yaml"
+DATA_DIR = "data/raw/"
+LOGS_DIR = "reports/logs/FluorescentMobileVIT"
+MODEL_DIR = "models/FluorescentMobileVIT"
+
+# Define the accelerator
+if torch.backends.mps.is_available():
+    DEVICE = torch.device("mps:0")
+    ACCELERATOR = "mps"
+elif torch.cuda.is_available():
+    DEVICE = torch.device("cuda")
+    ACCELERATOR = "gpu"
+else:
+    DEVICE = torch.device("cpu")
+    ACCELERATOR = "cpu"
+
+
+@click.command()
+@click.option(
+    "--data_dir",
+    type=click.Path(exists=True, file_okay=True, path_type=Path),
+    default=DATA_DIR,
+)
+@click.option(
+    "--config_file",
+    type=click.Path(exists=True, file_okay=True, path_type=Path),
+    default=CONFIG_FILE,
+)
+def train_model(data_dir, config_file):
+    # Load the best parameters
+    with open(config_file, "r") as file:
+        best_params = yaml.safe_load(file)
+    # Instantiate the model
+    model = MobileVIT(
+        learning_rate=best_params["learning_rate"],
+        weight_decay=best_params["weight_decay"],
+    )
+    # Define the callbacks of the model
+    model_checkpoint_cb = ModelCheckpoint(
+        save_top_k=1, dirpath=MODEL_DIR, monitor="val_loss"
+    )
+    logger = TensorBoardLogger(save_dir=LOGS_DIR)
+
+    # Create the trainer with its parameters
+    trainer = pl.Trainer(
+        logger=logger,
+        devices=1,
+        accelerator=ACCELERATOR,
+        precision=16,
+        max_epochs=100,
+        log_every_n_steps=20,
+        callbacks=[model_checkpoint_cb],
+    )
+    data_module = FluorescentNeuronalDataModule(
+        data_dir=data_dir, batch_size=best_params["batch_size"]
+    )
+    trainer.fit(model=model, datamodule=data_module)
+    trainer.test(model=model, datamodule=data_module)
+    click.echo("\n\n==========The Training has Finished!==========")