Updated Code

For MICCAI 2024

Author: philsaurabh

data/.DS_Store

@@ -0,0 +1,206 @@
+# encoding: utf-8
+"""
+Read images and corresponding labels.
+"""
+
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+import pandas as pd
+import numpy as np
+from PIL import Image
+import os
+
+from PIL import ImageFile
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+#############################################################
+##### Dataset with memory bank and contrastive samples. #####
+#############################################################
+
+
+class ISIC_InstanceSample(Dataset):
+
+    def __init__(self, root_dir, csv_file, CCD_mode, transform=None, p=10, k=4096,
+                 mode='exact', is_sample=True, percent=1.0):
+        super(ISIC_InstanceSample, self).__init__()
+
+        self.p = p
+        self.k = k
+        self.mode = mode
+        self.CCD_mode = CCD_mode
+        self.is_sample = is_sample
+
+        file = pd.read_csv(csv_file)
+
+        self.root_dir = root_dir
+        self.images = file['id_code'].values  # image name
+        self.labels = file['diagnosis'].values.astype(int)  # scalar label
+        n_classes = len(np.unique(self.labels))
+        # one hot. [num_images, num_classes]
+        self.labels = np.eye(n_classes)[self.labels.reshape(-1)]
+        self.transform = transform
+
+        print('Total # images:{}, labels:{}'.format(
+            len(self.images), len(self.labels)))
+
+        num_samples = len(self.images)
+        label = np.argmax(self.labels, axis=1)
+
+        self.cls_positive = [[] for i in range(n_classes)]
+        for i in range(num_samples):
+            self.cls_positive[label[i]].append(i)
+
+        self.cls_negative = [[] for i in range(n_classes)]
+        for i in range(n_classes):
+            for j in range(n_classes):
+                if j == i:
+                    continue
+                self.cls_negative[i].extend(self.cls_positive[j])
+
+        self.cls_positive = [np.asarray(self.cls_positive[i])
+                             for i in range(n_classes)]
+        self.cls_negative = [np.asarray(self.cls_negative[i])
+                             for i in range(n_classes)]
+
+        self.class_index = self.cls_positive
+
+        if 0 < percent < 1:
+            n = int(len(self.cls_negative[0]) * percent)
+            self.cls_negative = [np.random.permutation(self.cls_negative[i])[0:n]
+                                 for i in range(n_classes)]
+
+        self.cls_positive = np.asarray(self.cls_positive, dtype=object)
+        self.cls_negative = np.asarray(self.cls_negative, dtype=object)
+
+    def __getitem__(self, index):
+        image_name = os.path.join(self.root_dir, self.images[index]+'.png')
+        img = Image.open(image_name).convert('RGB')
+        target = np.argmax(self.labels, axis=1)[index]
+        label = self.labels[index]
+
+        if self.transform is not None:
+            img = self.transform(img)
+
+        if not self.is_sample:
+            return img, target, index
+        else:
+            # sample contrastive examples
+            if self.mode == 'exact':
+                pos_idx = index
+            elif self.mode == 'relax':
+                pos_idx = np.random.choice(self.cls_positive[target], 1)[0]
+            elif self.mode == 'multi_pos':
+                pos_idx = np.random.choice(
+                    self.cls_positive[target], self.p, replace=False)
+            else:
+                raise NotImplementedError(self.mode)
+
+            if self.CCD_mode == "sup":
+                replace = True if self.k > len(
+                    self.cls_negative[target]) else False
+                neg_idx = np.random.choice(
+                    self.cls_negative[target], self.k, replace=replace)
+            elif self.CCD_mode == "unsup":
+                pos_others = np.setdiff1d(self.cls_positive[target], ([index]))
+                all_negative = np.hstack(
+                    (pos_others, self.cls_negative[target]))
+                neg_idx = np.random.choice(all_negative, self.k, replace=True)
+
+            if self.mode == 'exact' or self.mode == 'relax':
+                sample_idx = np.hstack((np.asarray([pos_idx]), neg_idx))
+            elif self.mode == 'multi_pos':
+                sample_idx = np.hstack((pos_idx, neg_idx))
+
+            return img, label, index, sample_idx
+
+    def __len__(self):
+        return len(self.images)
+
+
+def load_dataset(args, p=10, k=4096, mode='exact', is_sample=True, percent=1.0):
+    csv_file_train = args.csv_file_path + \
+        args.dataset + '/' + args.split + '_train.csv'
+    csv_file_test = args.csv_file_path + args.dataset + '/' + args.split + '_test.csv'
+
+    train_transform = TransformTwice(transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.RandomAffine(degrees=10, translate=(0.02, 0.02)),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ]))
+
+    test_transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ])
+
+    train_set = ISIC_InstanceSample(root_dir=args.root_path,
+                                    csv_file=csv_file_train,
+                                    CCD_mode=args.CCD_mode,
+                                    transform=train_transform,
+                                    p=p,
+                                    k=k,
+                                    mode=mode,
+                                    is_sample=is_sample,
+                                    percent=percent)
+
+    test_set = ISIC_Dataset(root_dir=args.root_path,
+                            csv_file=csv_file_test, transform=test_transform)
+
+    return train_set, test_set
+
+
+# Dataset without memory bank
+class ISIC_Dataset(Dataset):
+    def __init__(self, root_dir, csv_file, transform=None):
+        """
+        Args:
+            data_dir: path to image directory.
+            csv_file: path to the file containing images
+                with corresponding labels.
+            transform: optional transform to be applied on a sample.
+        """
+        super(ISIC_Dataset, self).__init__()
+        file = pd.read_csv(csv_file)
+
+        self.root_dir = root_dir
+        self.images = file['id_code'].values  # image name
+        self.labels = file['diagnosis'].values.astype(int)
+        self.n_classes = len(np.unique(self.labels))
+        self.labels = np.eye(self.n_classes)[
+            self.labels.reshape(-1)]  # one_hot labels
+        self.transform = transform
+
+        print('Total # images:{}, labels:{}'.format(
+            len(self.images), len(self.labels)))
+
+    def __getitem__(self, index):
+        """
+        Args:
+            index: the index of item
+        Returns:
+            image and its labels
+        """
+        image_name = os.path.join(self.root_dir, self.images[index]+'.png')
+        image = Image.open(image_name).convert('RGB')
+        label = self.labels[index]
+        if self.transform is not None:
+            image = self.transform(image)
+
+        return image, label
+
+    def __len__(self):
+        return len(self.images)
+
+
+class TransformTwice:
+    def __init__(self, transform):
+        self.transform = transform
+
+    def __call__(self, inp):
+        out1 = self.transform(inp)
+        out2 = self.transform(inp)
+        return out1, out2

data/loader.py

@@ -0,0 +1,32 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import densenet121, DenseNet121_Weights
+
+
+class DenseNet121(nn.Module):
+    def __init__(self, hidden_units, out_size, drop_rate=0):
+        super(DenseNet121, self).__init__()
+        self.densenet121 = densenet121(
+            weights=DenseNet121_Weights.IMAGENET1K_V1)
+        num_ftrs = self.densenet121.classifier.in_features
+
+        self.densenet121.fc_layer = nn.Linear(num_ftrs, hidden_units)
+        self.densenet121.classifier = nn.Linear(hidden_units, out_size)
+
+        self.drop_rate = drop_rate
+        self.drop_layer = nn.Dropout(p=drop_rate)
+
+    def forward(self, x):
+        features = self.densenet121.features(x)
+        fmaps_b4 = F.relu(features, inplace=True)
+
+        out = F.adaptive_avg_pool2d(
+            fmaps_b4, (1, 1)).view(fmaps_b4.size(0), -1)
+
+        if self.drop_rate > 0:
+            out = self.drop_layer(out)
+
+        feature4 = self.densenet121.fc_layer(out)
+        logit_b4 = self.densenet121.classifier(feature4)
+        return feature4, logit_b4

logs/aptos/.DS_Store

@@ -0,0 +1,144 @@
+from utils.utils import get_labels_frequency, set_logger
+from utils.trainer import fit
+from models.model import DenseNet121
+from data.loader import load_dataset
+from torch.utils.data import DataLoader
+import torch.backends.cudnn as cudnn
+import torch
+import numpy as np
+import random
+import logging
+import sys
+import os
+import argparse
+import warnings
+warnings.simplefilter('ignore')
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root_path', type=str,
+                        default='../Datasets/APTOS/APTOS_images/train_images')
+    parser.add_argument('--csv_file_path', type=str, default='../CSVs/')
+    parser.add_argument("--logdir", type=str, required=False,
+                        default="./logs/aptos/", help="Log directory path")
+    parser.add_argument('--dataset', type=str, default='aptos')
+    parser.add_argument('--split', type=str, default='split1')
+
+    parser.add_argument('--n_distill', type=int, default=20,
+                        help='start to use the kld loss')
+
+    parser.add_argument('--mode', default='exact', type=str,
+                        choices=['exact', 'relax', 'multi_pos'])
+    parser.add_argument('--nce_p', default=1, type=int,
+                        help='number of positive samples for NCE')
+    parser.add_argument('--nce_k', default=4096, type=int,
+                        help='number of negative samples for NCE')
+    parser.add_argument('--nce_t', default=0.07, type=float,
+                        help='temperature parameter for softmax')
+    parser.add_argument('--nce_m', default=0.5, type=float,
+                        help='momentum for non-parametric updates')
+    parser.add_argument('--CCD_mode', type=str,
+                        default="sup", choices=['sup', 'unsup'])
+    parser.add_argument('--rel_weight', type=float, default=25,
+                        help='whether use the relation loss')
+    parser.add_argument('--ccd_weight', type=float,
+                        default=0.1, help='whether use the CCD loss')
+
+    parser.add_argument('--anchor_type', type=str,
+                        default="center", choices=['center', 'class'])
+    parser.add_argument('--class_anchor', default=30, type=int,
+                        help='number of anchors in each class')
+
+    parser.add_argument('--feat_dim', type=int, default=128,
+                        help='reduced feature dimension')
+    parser.add_argument('--s_dim', type=int, default=128,
+                        help='feature dim of the student model')
+    parser.add_argument('--t_dim', type=int, default=128,
+                        help='feature dim of the EMA teacher')
+    parser.add_argument('--n_data', type=int, default=3662,
+                        help='total number of training samples.')
+    parser.add_argument('--t_decay', type=float,
+                        default=0.99, help='ema_decay')
+
+    parser.add_argument('--epochs', type=int,  default=80,
+                        help='maximum epoch number to train')
+    parser.add_argument('--batch_size', type=int,
+                        default=64, help='batch_size per gpu')
+    parser.add_argument('--drop_rate', type=int,
+                        default=0, help='dropout rate')
+    parser.add_argument('--lr', type=float,  default=1e-4,
+                        help='learning rate')
+    parser.add_argument('--seed', type=int,  default=2024, help='random seed')
+
+    parser.add_argument('--optimizer', type=str,  default='adam', help='optim')
+    parser.add_argument('--scheduler', type=str,
+                        default='OneCycleLR', help='sch_str')
+    parser.add_argument('--device', type=str, default='cuda:0', help='device')
+
+    parser.add_argument('--consistency', type=float,
+                        default=1, help='consistency')
+    parser.add_argument('--consistency_rampup', type=float,
+                        default=30, help='consistency_rampup')
+
+    args = parser.parse_args()
+    return args
+
+# Function to set the seed for all random number generators to ensure reproducibility
+
+
+def set_seed(seed):
+    cudnn.benchmark = False
+    cudnn.deterministic = True
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+
+
+if __name__ == "__main__":
+    # Get arguments
+    args = get_args()
+
+    # Set seed
+    set_seed(args.seed)
+
+    # Set Logger
+    if not os.path.exists(args.logdir):
+        os.makedirs(args.logdir)
+    logger = set_logger(args)
+    logger.info(args)
+
+    # Loading Data
+    train_ds, test_ds = load_dataset(args, p=args.nce_p, mode=args.mode)
+    n_classes = test_ds.n_classes
+    class_index = train_ds.class_index
+    print(n_classes)
+
+    def worker_init_fn(worker_id):
+        random.seed(args.seed+worker_id)
+    train_dl = DataLoader(train_ds, batch_size=args.batch_size,
+                          shuffle=True, num_workers=12, pin_memory=True,
+                          worker_init_fn=worker_init_fn)
+
+    test_dl = DataLoader(test_ds, batch_size=args.batch_size,
+                         shuffle=False, num_workers=12, pin_memory=True,
+                         worker_init_fn=worker_init_fn)
+    freq = get_labels_frequency(args.csv_file_path + args.dataset +
+                                '/' + args.split + '_train.csv', 'diagnosis', 'id_code')
+    freq = freq.values
+    weights = freq.sum() / freq
+    print(weights)
+
+    # Loading Models
+    student = DenseNet121(hidden_units=args.feat_dim,
+                          out_size=n_classes, drop_rate=args.drop_rate)
+    teacher = DenseNet121(hidden_units=args.feat_dim,
+                          out_size=n_classes, drop_rate=args.drop_rate)
+
+    for param in teacher.parameters():
+        param.detach_()
+
+    # Fit the model
+    fit(student, teacher, train_dl, test_dl, weights,
+        class_index, logger, args, device=args.device)