first commit

README.md

@@ -8,5 +8,3 @@ sdk_version: 3.36.1
 app_file: app.py
 pinned: false
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,230 @@
+import os
+import numpy as np
+import gradio as gr
+from glob import glob
+from functools import partial
+from dataclasses import dataclass
+
+import torch
+import torchvision
+import torch.nn as nn
+import lightning.pytorch as pl
+import torchvision.transforms as TF
+
+from torchmetrics import MeanMetric
+from torchmetrics.classification import MultilabelF1Score
+
+
+@dataclass
+class DatasetConfig:
+    IMAGE_SIZE: tuple = (384, 384)  # (W, H)
+    CHANNELS: int = 3
+    NUM_CLASSES: int = 10
+    MEAN: tuple = (0.485, 0.456, 0.406)
+    STD: tuple = (0.229, 0.224, 0.225)
+
+
+@dataclass
+class TrainingConfig:
+    METRIC_THRESH: float = 0.4
+    MODEL_NAME: str = "efficientnet_v2_s"
+    FREEZE_BACKBONE: bool = False
+
+
+def get_model(model_name: str, num_classes: int, freeze_backbone: bool = True):
+    """A helper function to load and prepare any classification model
+    available in Torchvision for transfer learning or fine-tuning."""
+
+    model = getattr(torchvision.models, model_name)(weights="DEFAULT")
+
+    if freeze_backbone:
+        # Set all layer to be non-trainable
+        for param in model.parameters():
+            param.requires_grad = False
+
+    model_childrens = [name for name, _ in model.named_children()]
+
+    try:
+        final_layer_in_features = getattr(model, f"{model_childrens[-1]}")[-1].in_features
+    except Exception as e:
+        final_layer_in_features = getattr(model, f"{model_childrens[-1]}").in_features
+
+    new_output_layer = nn.Linear(in_features=final_layer_in_features, out_features=num_classes)
+
+    try:
+        getattr(model, f"{model_childrens[-1]}")[-1] = new_output_layer
+    except:
+        setattr(model, model_childrens[-1], new_output_layer)
+
+    return model
+
+
+class ProteinModel(pl.LightningModule):
+    def __init__(
+        self,
+        model_name: str,
+        num_classes: int = 10,
+        freeze_backbone: bool = False,
+        init_lr: float = 0.001,
+        optimizer_name: str = "Adam",
+        weight_decay: float = 1e-4,
+        use_scheduler: bool = False,
+        f1_metric_threshold: float = 0.4,
+    ):
+        super().__init__()
+
+        # Save the arguments as hyperparameters.
+        self.save_hyperparameters()
+
+        # Loading model using the function defined above.
+        self.model = get_model(
+            model_name=self.hparams.model_name,
+            num_classes=self.hparams.num_classes,
+            freeze_backbone=self.hparams.freeze_backbone,
+        )
+
+        # Intialize loss class.
+        self.loss_fn = nn.BCEWithLogitsLoss()
+
+        # Initializing the required metric objects.
+        self.mean_train_loss = MeanMetric()
+        self.mean_train_f1 = MultilabelF1Score(num_labels=self.hparams.num_classes, average="macro", threshold=self.hparams.f1_metric_threshold)
+        self.mean_valid_loss = MeanMetric()
+        self.mean_valid_f1 = MultilabelF1Score(num_labels=self.hparams.num_classes, average="macro", threshold=self.hparams.f1_metric_threshold)
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, *args, **kwargs):
+        data, target = batch
+        logits = self(data)
+        loss = self.loss_fn(logits, target)
+
+        self.mean_train_loss(loss, weight=data.shape[0])
+        self.mean_train_f1(logits, target)
+
+        self.log("train/batch_loss", self.mean_train_loss, prog_bar=True)
+        self.log("train/batch_f1", self.mean_train_f1, prog_bar=True)
+        return loss
+
+    def on_train_epoch_end(self):
+        # Computing and logging the training mean loss & mean f1.
+        self.log("train/loss", self.mean_train_loss, prog_bar=True)
+        self.log("train/f1", self.mean_train_f1, prog_bar=True)
+        self.log("step", self.current_epoch)
+
+    def validation_step(self, batch, *args, **kwargs):
+        data, target = batch  # Unpacking validation dataloader tuple
+        logits = self(data)
+        loss = self.loss_fn(logits, target)
+
+        self.mean_valid_loss.update(loss, weight=data.shape[0])
+        self.mean_valid_f1.update(logits, target)
+
+    def on_validation_epoch_end(self):
+        # Computing and logging the validation mean loss & mean f1.
+        self.log("valid/loss", self.mean_valid_loss, prog_bar=True)
+        self.log("valid/f1", self.mean_valid_f1, prog_bar=True)
+        self.log("step", self.current_epoch)
+
+    def configure_optimizers(self):
+        optimizer = getattr(torch.optim, self.hparams.optimizer_name)(
+            filter(lambda p: p.requires_grad, self.model.parameters()),
+            lr=self.hparams.init_lr,
+            weight_decay=self.hparams.weight_decay,
+        )
+
+        if self.hparams.use_scheduler:
+            lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+                optimizer,
+                milestones=[
+                    self.trainer.max_epochs // 2,
+                ],
+                gamma=0.1,
+            )
+
+            # The lr_scheduler_config is a dictionary that contains the scheduler
+            # and its associated configuration.
+            lr_scheduler_config = {
+                "scheduler": lr_scheduler,
+                "interval": "epoch",
+                "name": "multi_step_lr",
+            }
+            return {"optimizer": optimizer, "lr_scheduler": lr_scheduler_config}
+
+        else:
+            return optimizer
+
+
+@torch.inference_mode()
+def predict(input_image, threshold=0.4, model=None, preprocess_fn=None, device="cpu", idx2labels=None):
+    input_tensor = preprocess_fn(input_image)
+    input_tensor = input_tensor.unsqueeze(0).to(device)
+
+    # Generate predictions
+    output = model(input_tensor).cpu()
+
+    probabilities = torch.sigmoid(output)[0].numpy().tolist()
+
+    output_probs = dict()
+    predicted_classes = []
+
+    for idx, prob in enumerate(probabilities):
+        output_probs[idx2labels[idx]] = prob
+        if prob >= threshold:
+            predicted_classes.append(idx2labels[idx])
+
+    predicted_classes = "\n".join(predicted_classes)
+    return predicted_classes, output_probs
+
+
+if __name__ == "__main__":
+    labels = {
+        0: "Mitochondria",
+        1: "Nuclear bodies",
+        2: "Nucleoli",
+        3: "Golgi apparatus",
+        4: "Nucleoplasm",
+        5: "Nucleoli fibrillar center",
+        6: "Cytosol",
+        7: "Plasma membrane",
+        8: "Centrosome",
+        9: "Nuclear speckles",
+    }
+
+    DEVICE = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+    CKPT_PATH = os.path.join(os.getcwd(), r"ckpt_024-vloss_0.1816_vf1_0.7855.ckpt")
+    model = ProteinModel.load_from_checkpoint(CKPT_PATH)
+    model.to(DEVICE)
+    model.eval()
+    _ = model(torch.randn(1, DatasetConfig.CHANNELS, *DatasetConfig.IMAGE_SIZE[::-1], device=DEVICE))
+
+    preprocess = TF.Compose(
+        [
+            TF.Resize(size=DatasetConfig.IMAGE_SIZE[::-1]),
+            TF.ToTensor(),
+            TF.Normalize(DatasetConfig.MEAN, DatasetConfig.STD, inplace=True),
+        ]
+    )
+
+    images_dir = glob(os.path.join(os.getcwd(), "samples") + os.sep + "*.png")
+    examples = [[i, TrainingConfig.METRIC_THRESH] for i in np.random.choice(images_dir, size=8, replace=False)]
+    print(examples)
+
+    iface = gr.Interface(
+        fn=partial(predict, model=model, preprocess_fn=preprocess, device=DEVICE, idx2labels=labels),
+        inputs=[
+            gr.Image(type="pil", label="Image"),
+            gr.Slider(0.0, 1.0, value=0.4, label="Threshold", info="Select the cut-off threshold for a node to be considered as a valid output."),
+        ],
+        outputs=[
+            gr.Textbox(label="Labels Present"),
+            gr.Label(label="Probabilities", show_label=False),
+        ],
+        examples=examples,
+        cache_examples=False,
+        allow_flagging="never",
+        title="Medical Multi-Label Image Classification",
+    )
+
+    iface.launch()

ckpt_024-vloss_0.1816_vf1_0.7855.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eeba4764adb310bf3a35ba2479326fdbf38acaed3242a9f020ff2d7eba47b2ca
+size 243578302

requirements.txt

@@ -0,0 +1,5 @@
+--find-links https://download.pytorch.org/whl/torch_stable.html
+torch==2.0.0+cpu
+torchvision==0.15.0
+torchmetrics==1.0.0
+lightning==2.0.4