Initial commit

dl_supervised_pipeline.py

@@ -0,0 +1,158 @@
+# Code modified from pytorch-image-classification
+# obtained from https://colab.research.google.com/github/bentrevett/pytorch-image-classification/blob/master/5_resnet.ipynb#scrollTo=4QmwmcXuPuLo
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import torch.optim as optim
+import torch.optim.lr_scheduler as lr_scheduler
+
+import torch.utils.data as data
+
+import numpy as np
+import random
+import tqdm
+import os
+from pathlib import Path
+
+from data_utils.data_tribology import TribologyDataset
+from utils.experiment_utils import get_model, get_name, get_logger, train, evaluate, evaluate_vote
+from utils.arg_utils import get_args
+
+def main(args):
+    '''Reproducibility'''
+    SEED = args.seed
+    random.seed(SEED)
+    np.random.seed(SEED)
+    torch.manual_seed(SEED)
+    torch.cuda.manual_seed(SEED)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+    '''Folder Creation'''
+    basepath=os.getcwd()
+    experiment_dir = Path(os.path.join(basepath,'experiments',args.model,args.resolution,args.magnification,args.modality,args.pretrained,args.frozen,args.vote))
+    experiment_dir.mkdir(parents=True, exist_ok=True)
+    checkpoint_dir = Path(os.path.join(experiment_dir,'checkpoints'))
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+
+    '''Logging'''
+    model_name = get_name(args)
+    print(model_name, 'STARTED')
+    if os.path.exists(checkpoint_dir / 'epoch10.pth'):
+        print('CHECKPOINT FOUND')
+        print('TERMINATING TRAINING')
+        return 0 # terminate training if checkpoint exists
+    
+    logger = get_logger(experiment_dir, model_name)
+
+    '''Data Loading'''
+    train_csv_path = f"./LUA_Dataset/CSV/{args.resolution}_{args.magnification}_6w_train.csv"
+    test_csv_path = f"./LUA_Dataset/CSV/{args.resolution}_{args.magnification}_6w_test.csv"
+    img_path = f"./LUA_Dataset/{args.resolution}/{args.magnification}/{args.modality}"
+
+    # results_acc_1 = {}
+    # results_acc_3 = {}
+    # classes_num = 6
+    BATCHSIZE = args.batch_size
+    train_dataset = TribologyDataset(csv_path = train_csv_path, img_path = img_path)
+    test_dataset = TribologyDataset(csv_path = test_csv_path, img_path = img_path)
+
+    # prepare the data augmentation
+    means, stds = train_dataset.get_statistics()
+    train_dataset.prepare_transform(means, stds, mode='train')
+    test_dataset.prepare_transform(means, stds, mode='test')
+
+    VALID_RATIO = 0.1
+
+    num_train = len(train_dataset)
+    num_valid = int(VALID_RATIO * num_train)
+    train_dataset, valid_dataset = data.random_split(train_dataset, [num_train - num_valid, num_valid])
+    logger.info(f'Number of training samples: {len(train_dataset)}')
+    logger.info(f'Number of validation samples: {len(valid_dataset)}')
+    train_iterator = torch.utils.data.DataLoader(train_dataset, 
+                                                 batch_size=BATCHSIZE, 
+                                                 num_workers=4, 
+                                                 shuffle=True, 
+                                                 pin_memory=True,
+                                                 drop_last=False)
+    
+    valid_iterator = torch.utils.data.DataLoader(valid_dataset, 
+                                                 batch_size=BATCHSIZE, 
+                                                 num_workers=4, 
+                                                 shuffle=True, 
+                                                 pin_memory=True,
+                                                 drop_last=False)
+    test_iterator = torch.utils.data.DataLoader(test_dataset,
+                                                batch_size=BATCHSIZE, 
+                                                num_workers=4, 
+                                                shuffle=False, 
+                                                pin_memory=True,
+                                                drop_last=False)
+    print('DATA LOADED')
+
+    # Define model 
+    model = get_model(args)
+    print('MODEL LOADED')
+
+    # Define optimizer and scheduler
+    START_LR = args.start_lr
+    optimizer = optim.Adam(model.parameters(), lr=START_LR)
+    STEPS_PER_EPOCH = len(train_iterator)
+    print('STEPS_PER_EPOCH:', STEPS_PER_EPOCH)
+    print('VALIDATION STEPS:', len(valid_iterator))
+    scheduler = lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=max(STEPS_PER_EPOCH,STEPS_PER_EPOCH//10))
+
+    # Define loss function
+    criterion = nn.CrossEntropyLoss()
+
+    # Define device
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = model.to(device)
+    criterion = criterion.to(device)
+
+    EPOCHS = args.epochs
+
+    print('SETUP DONE')
+    # train our model
+
+    print('TRAINING STARTED')
+    for epoch in tqdm.tqdm(range(EPOCHS)):
+
+        train_loss, train_acc_1, train_acc_3 = train(model, train_iterator, optimizer, criterion, scheduler, device)
+        
+        torch.cuda.empty_cache() # clear cache between train and val
+
+        valid_loss, valid_acc_1, valid_acc_3 = evaluate(model, valid_iterator, criterion, device)
+
+        torch.save(model.state_dict(), checkpoint_dir / f'epoch{epoch+1}.pth')
+
+        logger.info(f'Epoch: {epoch + 1:02}')
+        logger.info(f'\tTrain Loss: {train_loss:.3f} | Train Acc @1: {train_acc_1 * 100:6.2f}% | ' \
+                f'Train Acc @3: {train_acc_3 * 100:6.2f}%')
+        logger.info(f'\tValid Loss: {valid_loss:.3f} | Valid Acc @1: {valid_acc_1 * 100:6.2f}% | ' \
+                f'Valid Acc @3: {valid_acc_3 * 100:6.2f}%')
+
+    logger.info('-------------------End of Training-------------------')
+    print('TRAINING DONE')
+    logger.info('-------------------Beginning of Testing-------------------')
+    print('TESTING STARTED')
+    for epoch in tqdm.tqdm(range(EPOCHS)):
+        model.load_state_dict(torch.load(checkpoint_dir / f'epoch{epoch+1}.pth'))
+
+        if args.vote == 'vote':
+            test_acc = evaluate_vote(model, test_iterator, device)
+            logger.info(f'Test Acc @1: {test_acc * 100:6.2f}%')
+        else:
+            test_loss, test_acc_1, test_acc_3 = evaluate(model, test_iterator, criterion, device)
+
+            logger.info(f'Test Acc @1: {test_acc_1 * 100:6.2f}% | ' \
+                    f'Test Acc @3: {test_acc_3 * 100:6.2f}%')
+    logger.info('-------------------End of Testing-------------------')
+    print('TESTING DONE')
+
+
+if __name__ == '__main__':
+    args = get_args()
+    main(args)

run_gradio.py

@@ -0,0 +1,55 @@
+import gradio as gr
+import torch
+import torchvision
+from utils.experiment_utils import get_model
+
+
+# 加载DINOv2模型
+def load_model():
+    class Args:
+        model = 'DINOv2'
+        pretrained = 'pretrained'
+        frozen = 'unfrozen'
+
+    args = Args()
+    model = get_model(args)
+    model.eval()
+    return model
+
+
+model = load_model()
+
+
+# 预测函数，返回每个类别的概率
+def predict(image):
+    transform = torchvision.transforms.Compose([
+        torchvision.transforms.Resize((224, 224)),
+        torchvision.transforms.ToTensor(),
+        torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+    image = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        output = model(image)
+        probabilities = torch.nn.functional.softmax(output, dim=1).squeeze().tolist()
+
+    # 类别名称列表
+    class_names = ["ANTLER", "BEECHWOOD", "BEFOREUSE", "BONE", "IVORY", "SPRUCEWOOD"]
+
+    # 将类别和对应的概率配对
+    results = {class_names[i]: probabilities[i] for i in range(len(class_names))}
+
+    return results
+
+
+# 创建Gradio界面
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=len(["ANTLER", "BEECHWOOD", "BEFOREUSE", "BONE", "IVORY", "SPRUCEWOOD"])),
+    title="LUWA DINOv2 Prediction",
+    description="Upload an image to get the probabilities for each class using the DINOv2 model."
+)
+
+if __name__ == "__main__":
+    interface.launch(share=True)

svm_pipeline.py

@@ -0,0 +1,100 @@
+import numpy as np
+from sklearn.svm import LinearSVC
+
+from skimage.feature import fisher_vector, learn_gmm
+
+import numpy as np
+import random
+import os
+from pathlib import Path
+
+from data_utils.data_tribology import TribologyDataset
+from utils.arg_utils import get_args
+from utils.experiment_utils import get_name, get_logger, SIFT_extraction, conduct_voting
+from utils.visualization_utils import plot_confusion_matrix
+from vis_confusion_mtx import generate_confusion_matrix
+
+def main(args):
+    '''Reproducibility'''
+    SEED = args.seed
+    random.seed(SEED)
+    np.random.seed(SEED)
+
+    '''Folder Creation'''
+    basepath=os.getcwd()
+    experiment_dir = Path(os.path.join(basepath,'experiments',args.model,args.resolution,args.magnification,args.modality,args.vote))
+    experiment_dir.mkdir(parents=True, exist_ok=True)
+    checkpoint_dir = Path(os.path.join(experiment_dir,'checkpoints'))
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+
+    '''Logging'''
+    model_name = get_name(args)
+    print(model_name, 'STARTED', flush=True)
+    logger = get_logger(experiment_dir, model_name)
+
+    '''Data Loading'''
+    train_csv_path = f"./LUA_Dataset/CSV/{args.resolution}_{args.magnification}_6w_train.csv"
+    test_csv_path = f"./LUA_Dataset/CSV/{args.resolution}_{args.magnification}_6w_test.csv"
+    img_path = f"./LUA_Dataset/{args.resolution}/{args.magnification}/{args.modality}"
+
+    BATCHSIZE = args.batch_size
+    train_dataset = TribologyDataset(csv_path = train_csv_path, img_path = img_path)
+    test_dataset = TribologyDataset(csv_path = test_csv_path, img_path = img_path)
+
+    # prepare the data augmentation
+    means, stds = train_dataset.get_statistics()
+    train_dataset.prepare_transform(means, stds, mode='train')
+    test_dataset.prepare_transform(means, stds, mode='test')
+
+    VALID_RATIO = 0.1
+
+    num_train = len(train_dataset)
+    num_valid = int(VALID_RATIO * num_train)
+    # train_dataset, valid_dataset = data.random_split(train_dataset, [num_train - num_valid, num_valid])
+    # logger.info(f'Number of training samples: {len(train_dataset)}')
+    # logger.info(f'Number of validation samples: {len(valid_dataset)}')
+
+    train_names, train_descriptor, train_labels = SIFT_extraction(train_dataset)
+    test_names, test_descriptor, test_labels = SIFT_extraction(test_dataset)
+    # val_descriptor, val_labels = SIFT_extraction(valid_dataset)
+    print('DATA LOADED', flush=True)
+
+    print('TRAINING STARTED', flush=True)
+
+    # Train a K-mode GMM
+    k = 16
+    gmm = learn_gmm(train_descriptor, n_modes=k)
+
+    # Compute the Fisher vectors
+    training_fvs = np.array([
+        fisher_vector(descriptor_mat, gmm)
+        for descriptor_mat in train_descriptor
+    ])
+
+    testing_fvs = np.array([
+        fisher_vector(descriptor_mat, gmm)
+        for descriptor_mat in test_descriptor
+    ])
+    
+    svm = LinearSVC().fit(training_fvs, train_labels)
+
+    logger.info('-------------------End of Training-------------------')
+    print('TRAINING DONE')
+    logger.info('-------------------Beginning of Testing-------------------')
+    print('TESTING STARTED')
+    predictions = svm.predict(testing_fvs)
+    conduct_voting(test_names, predictions)
+    plot_confusion_matrix('visualization_results/SIFT+FVs_confusion_mtx.png', predictions, test_labels,classes=["ANTLER", "BEECHWOOD", "BEFOREUSE", "BONE", "IVORY","SPRUCEWOOD"])
+    correct = 0
+    for i in range(len(predictions)):
+        if predictions[i] == test_labels[i]:
+            correct += 1
+    test_acc = float(correct)/len(predictions)
+    logger.info(f'Test Acc @1: {test_acc * 100:6.2f}%')
+
+    logger.info('-------------------End of Testing-------------------')
+    print('TESTING DONE')
+
+if __name__ == '__main__':
+    args = get_args()
+    main(args)

utils/MAE.py

@@ -0,0 +1,253 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+from timm.models.vision_transformer import PatchEmbed, Block
+
+from utils.model_utils import get_2d_sincos_pos_embed
+
+
+class MaskedAutoencoderViT(nn.Module):
+    """ Masked Autoencoder with VisionTransformer backbone
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3,
+                 embed_dim=1024, depth=24, num_heads=16,
+                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+                 mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=False):
+        super().__init__()
+
+        # --------------------------------------------------------------------------
+        # MAE encoder specifics
+        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False)  # fixed sin-cos embedding
+
+        self.blocks = nn.ModuleList([
+            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+        # --------------------------------------------------------------------------
+
+        # --------------------------------------------------------------------------
+        # MAE decoder specifics
+        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
+
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
+
+        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, decoder_embed_dim), requires_grad=False)  # fixed sin-cos embedding
+
+        self.decoder_blocks = nn.ModuleList([
+            Block(decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
+            for i in range(decoder_depth)])
+
+        self.decoder_norm = norm_layer(decoder_embed_dim)
+        self.decoder_pred = nn.Linear(decoder_embed_dim, patch_size**2 * in_chans, bias=True) # decoder to patch
+        # --------------------------------------------------------------------------
+
+        self.norm_pix_loss = norm_pix_loss
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # initialization
+        # initialize (and freeze) pos_embed by sin-cos embedding
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        decoder_pos_embed = get_2d_sincos_pos_embed(self.decoder_pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)
+        self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))
+
+        # initialize patch_embed like nn.Linear (instead of nn.Conv2d)
+        w = self.patch_embed.proj.weight.data
+        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+
+        # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
+        torch.nn.init.normal_(self.cls_token, std=.02)
+        torch.nn.init.normal_(self.mask_token, std=.02)
+
+        # initialize nn.Linear and nn.LayerNorm
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            # we use xavier_uniform following official JAX ViT:
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def patchify(self, imgs):
+        """
+        imgs: (N, 3, H, W)
+        x: (N, L, patch_size**2 *3)
+        """
+        p = self.patch_embed.patch_size[0]
+        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0
+
+        h = w = imgs.shape[2] // p
+        x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))
+        x = torch.einsum('nchpwq->nhwpqc', x)
+        x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3))
+        return x
+
+    def unpatchify(self, x):
+        """
+        x: (N, L, patch_size**2 *3)
+        imgs: (N, 3, H, W)
+        """
+        p = self.patch_embed.patch_size[0]
+        h = w = int(x.shape[1]**.5)
+        assert h * w == x.shape[1]
+        
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p))
+        return imgs
+
+    def random_masking(self, x, mask_ratio):
+        """
+        Perform per-sample random masking by per-sample shuffling.
+        Per-sample shuffling is done by argsort random noise.
+        x: [N, L, D], sequence
+        """
+        N, L, D = x.shape  # batch, length, dim
+        len_keep = int(L * (1 - mask_ratio))
+        
+        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]
+        
+        # sort noise for each sample
+        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        # keep the first subset
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+        # generate the binary mask: 0 is keep, 1 is remove
+        mask = torch.ones([N, L], device=x.device)
+        mask[:, :len_keep] = 0
+        # unshuffle to get the binary mask
+        mask = torch.gather(mask, dim=1, index=ids_restore)
+
+        return x_masked, mask, ids_restore
+
+    def forward_encoder(self, x, mask_ratio):
+        # embed patches
+        x = self.patch_embed(x)
+
+        # add pos embed w/o cls token
+        x = x + self.pos_embed[:, 1:, :]
+
+        # masking: length -> length * mask_ratio
+        x, mask, ids_restore = self.random_masking(x, mask_ratio)
+
+        # append cls token
+        cls_token = self.cls_token + self.pos_embed[:, :1, :]
+        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # apply Transformer blocks
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+
+        return x, mask, ids_restore
+
+    def forward_decoder(self, x, ids_restore):
+        # embed tokens
+        x = self.decoder_embed(x)
+
+        # append mask tokens to sequence
+        mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)
+        x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1)  # no cls token
+        x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))  # unshuffle
+        x = torch.cat([x[:, :1, :], x_], dim=1)  # append cls token
+
+        # add pos embed
+        x = x + self.decoder_pos_embed
+
+        # apply Transformer blocks
+        for blk in self.decoder_blocks:
+            x = blk(x)
+        x = self.decoder_norm(x)
+
+        # predictor projection
+        x = self.decoder_pred(x)
+
+        # remove cls token
+        x = x[:, 1:, :]
+
+        return x
+
+    def forward_loss(self, imgs, pred, mask):
+        """
+        imgs: [N, 3, H, W]
+        pred: [N, L, p*p*3]
+        mask: [N, L], 0 is keep, 1 is remove, 
+        """
+        target = self.patchify(imgs)
+        if self.norm_pix_loss:
+            mean = target.mean(dim=-1, keepdim=True)
+            var = target.var(dim=-1, keepdim=True)
+            target = (target - mean) / (var + 1.e-6)**.5
+
+        loss = (pred - target) ** 2
+        loss = loss.mean(dim=-1)  # [N, L], mean loss per patch
+
+        loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches
+        return loss
+
+    def forward(self, imgs, mask_ratio=0.75):
+        latent, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
+        # pred = self.forward_decoder(latent, ids_restore)  # [N, L, p*p*3]
+        # loss = self.forward_loss(imgs, pred, mask)
+        # return loss, pred, mask
+        print(latent.shape)
+        return latent
+        
+
+
+def mae_vit_base_patch16_dec512d8b(**kwargs):
+    model = MaskedAutoencoderViT(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def mae_vit_large_patch16_dec512d8b(**kwargs):
+    model = MaskedAutoencoderViT(
+        patch_size=16, embed_dim=1024, depth=24, num_heads=16,
+        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def mae_vit_huge_patch14_dec512d8b(**kwargs):
+    model = MaskedAutoencoderViT(
+        patch_size=14, embed_dim=1280, depth=32, num_heads=16,
+        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+# set recommended archs
+mae_vit_base_patch16 = mae_vit_base_patch16_dec512d8b  # decoder: 512 dim, 8 blocks
+mae_vit_large_patch16 = mae_vit_large_patch16_dec512d8b  # decoder: 512 dim, 8 blocks
+mae_vit_huge_patch14 = mae_vit_huge_patch14_dec512d8b  # decoder: 512 dim, 8 blocks

utils/__init__.py

@@ -0,0 +1 @@
+from .util_function import epoch_time, plot_lr_finder, plot_confusion_matrix, plot_most_incorrect, get_pca, plot_representations, plot_filtered_images, plot_filters

utils/__pycache__/MAE.cpython-38.pyc:Zone.Identifier

@@ -0,0 +1,3 @@
+[ZoneTransfer]
+ZoneId=3
+ReferrerUrl=C:\Users\Daniel\Desktop\LUWA-main.zip

utils/__pycache__/__init__.cpython-310.pyc:Zone.Identifier

@@ -0,0 +1,3 @@
+[ZoneTransfer]
+ZoneId=3
+ReferrerUrl=C:\Users\Daniel\Desktop\LUWA-main.zip

utils/__pycache__/__init__.cpython-38.pyc:Zone.Identifier

@@ -0,0 +1,3 @@
+[ZoneTransfer]
+ZoneId=3
+ReferrerUrl=C:\Users\Daniel\Desktop\LUWA-main.zip

utils/__pycache__/__init__.cpython-39.pyc:Zone.Identifier

@@ -0,0 +1,3 @@
+[ZoneTransfer]
+ZoneId=3
+ReferrerUrl=C:\Users\Daniel\Desktop\LUWA-main.zip

utils/__pycache__/arg_utils.cpython-38.pyc:Zone.Identifier

@@ -0,0 +1,3 @@
+[ZoneTransfer]
+ZoneId=3
+ReferrerUrl=C:\Users\Daniel\Desktop\LUWA-main.zip

utils/__pycache__/arg_utils.cpython-39.pyc:Zone.Identifier

@@ -0,0 +1,3 @@
+[ZoneTransfer]
+ZoneId=3
+ReferrerUrl=C:\Users\Daniel\Desktop\LUWA-main.zip

utils/__pycache__/experiment_utils.cpython-38.pyc:Zone.Identifier

@@ -0,0 +1,3 @@
+[ZoneTransfer]
+ZoneId=3
+ReferrerUrl=C:\Users\Daniel\Desktop\LUWA-main.zip

utils/__pycache__/experiment_utils.cpython-39.pyc:Zone.Identifier

@@ -0,0 +1,3 @@
+[ZoneTransfer]
+ZoneId=3
+ReferrerUrl=C:\Users\Daniel\Desktop\LUWA-main.zip

utils/__pycache__/model_utils.cpython-38.pyc:Zone.Identifier

@@ -0,0 +1,3 @@
+[ZoneTransfer]
+ZoneId=3
+ReferrerUrl=C:\Users\Daniel\Desktop\LUWA-main.zip

utils/__pycache__/util_function.cpython-310.pyc:Zone.Identifier

@@ -0,0 +1,3 @@
+[ZoneTransfer]
+ZoneId=3
+ReferrerUrl=C:\Users\Daniel\Desktop\LUWA-main.zip

utils/__pycache__/util_function.cpython-38.pyc:Zone.Identifier

@@ -0,0 +1,3 @@
+[ZoneTransfer]
+ZoneId=3
+ReferrerUrl=C:\Users\Daniel\Desktop\LUWA-main.zip

utils/__pycache__/util_function.cpython-39.pyc:Zone.Identifier

@@ -0,0 +1,3 @@
+[ZoneTransfer]
+ZoneId=3
+ReferrerUrl=C:\Users\Daniel\Desktop\LUWA-main.zip

utils/arg_utils.py

@@ -0,0 +1,18 @@
+import argparse
+
+def get_args():
+    # Training settings
+    parser = argparse.ArgumentParser('train')
+    
+    parser.add_argument('--resolution', type=str, default='256', help='Resolution of input image')
+    parser.add_argument('--magnification', type=str, default='20x', help='Magnification of input image')
+    parser.add_argument('--modality', type=str, default='texture', help='Modality of input image')
+    parser.add_argument('--model', type=str, default='ResNet50', help='Model to use')
+    parser.add_argument('--pretrained', type=str, default='pretrained', help='Use pretrained model')
+    parser.add_argument('--frozen', type=str, default='unfrozen', help='Freeze pretrained model')
+    parser.add_argument('--vote', type=str, default='vote', help='Conduct voting')
+    parser.add_argument('--epochs', type=int, default=2, help='Number of epochs to train')
+    parser.add_argument('--batch_size', type=int, default=100, help='Batch size')
+    parser.add_argument('--start_lr', type=float, default=0.01, help='Learning rate')
+    parser.add_argument('--seed', type=int, default=1234, help='Random seed')
+    return parser.parse_args()

utils/experiment_utils.py

@@ -0,0 +1,298 @@
+import torch
+import torchvision
+import torch.nn as nn
+import torch.nn.functional as F
+import logging
+from collections import Counter
+from utils.MAE import mae_vit_large_patch16_dec512d8b as MAE_large 
+
+def get_model(args) -> nn.Module:
+    if 'ResNet' in args.model:
+        # resnet family
+        if args.model == 'ResNet50':
+            if args.pretrained == 'pretrained':
+                model = torchvision.models.resnet50(weights='IMAGENET1K_V2')
+            else:
+                model = torchvision.models.resnet50()
+        elif args.model == 'ResNet152':
+            if args.pretrained == 'pretrained':
+                model = torchvision.models.resnet152(weights='IMAGENET1K_V2')
+            else:
+                model = torchvision.models.resnet152()
+        else:
+            raise NotImplementedError
+        if args.frozen == 'frozen':
+            model = freeze_backbone(model)
+        model.fc = nn.Linear(model.fc.in_features, 6)
+    
+    elif 'ConvNext' in args.model:
+        if args.model == 'ConvNext_Tiny':
+            if args.pretrained == 'pretrained':
+                model = torchvision.models.convnext_tiny(weights='IMAGENET1K_V1')
+            else:
+                model = torchvision.models.convnext_tiny()
+        elif args.model == 'ConvNext_Large':
+            if args.pretrained == 'pretrained':
+                model = torchvision.models.convnext_large(weights='IMAGENET1K_V1')
+            else:
+                model = torchvision.models.convnext_large()
+        else:
+            raise NotImplementedError
+        if args.frozen == 'frozen':
+            model = freeze_backbone(model)
+        num_ftrs = model.classifier[2].in_features
+        model.classifier[2] = nn.Linear(int(num_ftrs), 6) 
+
+    elif 'ViT' in args.model:
+        if args.pretrained == 'pretrained':
+            model = torchvision.models.vit_h_14(weights='IMAGENET1K_SWAG_LINEAR_V1')
+        else:
+            raise NotImplementedError('ViT does not support training from scratch')
+        if args.frozen == 'frozen':
+            model = freeze_backbone(model)
+        model.heads[0] = torch.nn.Linear(model.heads[0].in_features, 6)
+
+    elif 'DINOv2' in args.model:
+        if args.pretrained == 'pretrained':
+            model  = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitg14_reg_lc')
+        else:
+            raise NotImplementedError('DINOv2 does not support training from scratch')
+        if args.frozen == 'frozen':
+            model = freeze_backbone(model)
+        model.linear_head = torch.nn.Linear(model.linear_head.in_features, 6)
+    
+    elif 'MAE' in args.model:
+        if args.pretrained == 'pretrained':
+            model = MAE_large()
+            model.load_state_dict(torch.load('/scratch/zf540/LUWA/workspace/utils/pretrained_weights/mae_visualize_vit_large.pth')['model'])
+        else:
+            raise NotImplementedError('MAE does not support training from scratch')
+        if args.frozen == 'frozen':
+            model = freeze_backbone(model)
+        model = nn.Sequential(model, nn.Linear(1024, 6))
+        print(model)
+    else:
+        raise NotImplementedError
+    return model
+
+
+def freeze_backbone(model):
+    # freeze backbone
+    # we will replace the classifier at the end with a trainable one anyway, so we freeze the default here as well
+    for param in model.parameters():
+        param.requires_grad = False
+    return model
+
+def get_name(args):
+    name = args.model
+    name += '_'+str(args.resolution)
+    name += '_'+args.magnification
+    name += '_'+args.modality
+    if args.pretrained == 'pretrained':
+        name += '_pretrained'
+    else:
+        name += '_scratch'
+    if args.frozen == 'frozen':
+        name += '_frozen'
+    else:
+        name += '_unfrozen'
+    if args.vote == 'vote':
+        name += '_vote'
+    else:
+        name += '_novote'
+    return name
+
+def get_logger(path, name):
+    # set up logger
+
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.INFO)
+    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    file_handler = logging.FileHandler(path.joinpath(f'{name}_log.txt'))
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(formatter)
+    logger.addHandler(file_handler)
+    logger.info('---------------------------------------------------TRANING---------------------------------------------------')
+    
+    return logger
+
+def calculate_topk_accuracy(y_pred, y, k = 3):
+    with torch.no_grad():
+        batch_size = y.shape[0]
+        _, top_pred = y_pred.topk(k, 1)
+        top_pred = top_pred.t()
+        correct = top_pred.eq(y.view(1, -1).expand_as(top_pred))
+        correct_1 = correct[:1].reshape(-1).float().sum(0, keepdim = True)
+        correct_k = correct[:k].reshape(-1).float().sum(0, keepdim = True)
+        acc_1 = correct_1 / batch_size
+        acc_k = correct_k / batch_size
+    return acc_1, acc_k
+
+def train(model, iterator, optimizer, criterion, scheduler, device):
+    epoch_loss = 0
+    epoch_acc_1 = 0
+    epoch_acc_3 = 0
+
+    model.train()
+
+    for image, label, image_name in iterator:
+        x = image.to(device)
+        y = label.to(device)
+
+        optimizer.zero_grad()
+
+        y_pred = model(x)
+        print(y_pred.shape)
+        print(y.shape)
+        loss = criterion(y_pred, y)
+
+        acc_1, acc_3 = calculate_topk_accuracy(y_pred, y)
+
+        loss.backward()
+
+        optimizer.step()
+
+        scheduler.step()
+
+        epoch_loss += loss.item()
+        epoch_acc_1 += acc_1.item()
+        epoch_acc_3 += acc_3.item()
+
+    epoch_loss /= len(iterator)
+    epoch_acc_1 /= len(iterator)
+    epoch_acc_3 /= len(iterator)
+
+    return epoch_loss, epoch_acc_1, epoch_acc_3
+
+
+def evaluate(model, iterator, criterion, device):
+    epoch_loss = 0
+    epoch_acc_1 = 0
+    epoch_acc_3 = 0
+
+    model.eval()
+
+    with torch.no_grad():
+        for image, label, image_name in iterator:
+            x = image.to(device)
+            y = label.to(device)
+
+            y_pred = model(x)
+            loss = criterion(y_pred, y)
+
+            acc_1, acc_3 = calculate_topk_accuracy(y_pred, y)
+            
+            epoch_loss += loss.item()
+            epoch_acc_1 += acc_1.item()
+            epoch_acc_3 += acc_3.item()
+    
+    epoch_loss /= len(iterator)
+    epoch_acc_1 /= len(iterator)
+    epoch_acc_3 /= len(iterator)
+
+    return epoch_loss, epoch_acc_1, epoch_acc_3
+
+def evaluate_vote(model, iterator, device):
+
+    model.eval()
+
+    image_names = []
+    labels = []
+    predictions = []
+
+    with torch.no_grad():
+
+        for image, label, image_name in iterator:
+
+            x = image.to(device)
+
+            y_pred = model(x)
+            y_prob = F.softmax(y_pred, dim = -1)
+            top_pred = y_prob.argmax(1, keepdim = True)
+
+            image_names.extend(image_name)
+            labels.extend(label.numpy())
+            predictions.extend(top_pred.cpu().squeeze().numpy())
+
+    conduct_voting(image_names, predictions)
+
+    correct_count = 0
+    for i in range(len(labels)):
+        if labels[i] == predictions[i]:
+            correct_count += 1
+    accuracy = correct_count/len(labels)
+    return accuracy
+
+def conduct_voting(image_names, predictions):
+    # we need to do this because not all stones have the same number of partition
+    last_stone = image_names[0][:-8] # the name of the stone of the last image
+    voting_list = []
+    for i in range(len(image_names)):
+        image_area_name = image_names[i][:-8]
+        if image_area_name != last_stone:
+            # we have run through all the images of the last stone. We start voting
+            vote(voting_list, predictions, i)
+            voting_list = [] # reset the voting list
+        voting_list.append(predictions[i])
+        last_stone = image_area_name # update the last stone name
+    
+    # vote for the last stone
+    vote(voting_list, predictions, len(image_names))
+
+def vote(voting_list, predictions, i):
+    vote_result = Counter(voting_list).most_common(1)[0][0] # the most common prediction in the list
+    predictions[i-len(voting_list):i] = [vote_result]*len(voting_list) # replace the predictions of the last stone with the vote result
+        
+        
+
+
+# def get_predictions(model, iterator):
+
+#     model.eval()
+
+#     images = []
+#     labels = []
+#     probs = []
+
+#     with torch.no_grad():
+
+#         for (x, y) in iterator:
+
+#             x = x.to(device)
+
+#             y_pred = model(x)
+
+#             y_prob = F.softmax(y_pred, dim = -1)
+#             top_pred = y_prob.argmax(1, keepdim = True)
+
+#             images.append(x.cpu())
+#             labels.append(y.cpu())
+#             probs.append(y_prob.cpu())
+
+#     images = torch.cat(images, dim = 0)
+#     labels = torch.cat(labels, dim = 0)
+#     probs = torch.cat(probs, dim = 0)
+
+#     return images, labels, probs
+
+
+# def get_representations(model, iterator):
+#     model.eval()
+
+#     outputs = []
+#     intermediates = []
+#     labels = []
+
+#     with torch.no_grad():
+#         for (x, y) in iterator:
+#             x = x.to(device)
+
+#             y_pred = model(x)
+
+#             outputs.append(y_pred.cpu())
+#             labels.append(y)
+
+#     outputs = torch.cat(outputs, dim=0)
+#     labels = torch.cat(labels, dim=0)
+
+#     return outputs, labels

utils/model_utils.py

@@ -0,0 +1,96 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# Position embedding utils
+# --------------------------------------------------------
+
+import numpy as np
+
+import torch
+
+# --------------------------------------------------------
+# 2D sine-cosine position embedding
+# References:
+# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
+# MoCo v3: https://github.com/facebookresearch/moco-v3
+# --------------------------------------------------------
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=float)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+# --------------------------------------------------------
+# Interpolate position embeddings for high-resolution
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+def interpolate_pos_embed(model, checkpoint_model):
+    if 'pos_embed' in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model['pos_embed']
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model['pos_embed'] = new_pos_embed

utils/util_function.py

@@ -0,0 +1,238 @@
+import cv2
+from sklearn.manifold import TSNE
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import torch.nn.functional as F
+from sklearn.metrics import confusion_matrix
+from sklearn.metrics import ConfusionMatrixDisplay
+from sklearn import decomposition
+import itertools
+
+def normalize_image(image):
+    image_min = image.min()
+    image_max = image.max()
+    image.clamp_(min = image_min, max = image_max)
+    image.add_(-image_min).div_(image_max - image_min + 1e-5)
+    return image
+
+
+def plot_lr_finder(fig_name, lrs, losses, skip_start=5, skip_end=5):
+    if skip_end == 0:
+        lrs = lrs[skip_start:]
+        losses = losses[skip_start:]
+    else:
+        lrs = lrs[skip_start:-skip_end]
+        losses = losses[skip_start:-skip_end]
+
+    fig = plt.figure(figsize=(16, 8))
+    ax = fig.add_subplot(1, 1, 1)
+    ax.plot(lrs, losses)
+    ax.set_xscale('log')
+    ax.set_xlabel('Learning rate')
+    ax.set_ylabel('Loss')
+    ax.grid(True, 'both', 'x')
+    plt.show()
+    plt.savefig(fig_name)
+
+def epoch_time(start_time, end_time):
+    elapsed_time = end_time - start_time
+    elapsed_mins = int(elapsed_time / 60)
+    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
+    return elapsed_mins, elapsed_secs
+
+
+def plot_confusion_matrix(fig_name, labels, pred_labels, classes):
+    fig = plt.figure(figsize=(50, 50));
+    ax = fig.add_subplot(1, 1, 1);
+    cm = confusion_matrix(labels, pred_labels);
+    cm = ConfusionMatrixDisplay(cm, display_labels=classes);
+    cm.plot(values_format='d', cmap='Blues', ax=ax)
+    fig.delaxes(fig.axes[1])  # delete colorbar
+    plt.xticks(rotation=90, fontsize=50)
+    plt.yticks(fontsize=50)
+    plt.rcParams.update({'font.size': 50})
+    plt.xlabel('Predicted Label', fontsize=50)
+    plt.ylabel('True Label', fontsize=50)
+    plt.savefig(fig_name)
+
+def plot_confusion_matrix_SVM(fig_name, true_labels, predicted_labels, classes):
+    fig = plt.figure(figsize=(100, 100))
+    ax = fig.add_subplot(1, 1, 1)
+    
+    cm = confusion_matrix(true_labels, predicted_labels)
+    cm_display = ConfusionMatrixDisplay(cm, display_labels=classes)
+    
+    cm_display.plot(values_format='d', cmap='Blues', ax=ax)
+    
+    fig.delaxes(fig.axes[1])  # delete colorbar
+    plt.xticks(rotation=90, fontsize=50)
+    plt.yticks(fontsize=50)
+    plt.rcParams.update({'font.size': 50})
+    plt.xlabel('Predicted Label', fontsize=50)
+    plt.ylabel('True Label', fontsize=50)
+    plt.savefig(fig_name)
+
+
+def plot_most_incorrect(fig_name, incorrect, classes, n_images, normalize=True):
+    rows = int(np.sqrt(n_images))
+    cols = int(np.sqrt(n_images))
+
+    fig = plt.figure(figsize=(25, 20))
+
+    for i in range(rows * cols):
+
+        ax = fig.add_subplot(rows, cols, i + 1)
+
+        image, true_label, probs = incorrect[i]
+        image = image.permute(1, 2, 0)
+        true_prob = probs[true_label]
+        incorrect_prob, incorrect_label = torch.max(probs, dim=0)
+        true_class = classes[true_label]
+        incorrect_class = classes[incorrect_label]
+
+        if normalize:
+            image = normalize_image(image)
+
+        ax.imshow(image.cpu().numpy())
+        ax.set_title(f'true label: {true_class} ({true_prob:.3f})\n' \
+                     f'pred label: {incorrect_class} ({incorrect_prob:.3f})')
+        ax.axis('off')
+
+    fig.subplots_adjust(hspace=0.4)
+    plt.savefig(fig_name)
+
+def get_pca(data, n_components = 2):
+    pca = decomposition.PCA()
+    pca.n_components = n_components
+    pca_data = pca.fit_transform(data)
+    return pca_data
+
+
+def plot_representations(fig_name, data, labels, classes, n_images=None):
+    if n_images is not None:
+        data = data[:n_images]
+        labels = labels[:n_images]
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(111)
+    scatter = ax.scatter(data[:, 0], data[:, 1], c=labels, cmap='hsv')
+    # handles, _ = scatter.legend_elements(num = None)
+    # legend = plt.legend(handles = handles, labels = classes)
+    plt.savefig(fig_name)
+
+
+def plot_filtered_images(fig_name, images, filters, n_filters = None, normalize = True):
+
+    images = torch.cat([i.unsqueeze(0) for i in images], dim = 0).cpu()
+    filters = filters.cpu()
+
+    if n_filters is not None:
+        filters = filters[:n_filters]
+
+    n_images = images.shape[0]
+    n_filters = filters.shape[0]
+
+    filtered_images = F.conv2d(images, filters)
+
+    fig = plt.figure(figsize = (30, 30))
+
+    for i in range(n_images):
+
+        image = images[i]
+
+        if normalize:
+            image = normalize_image(image)
+
+        ax = fig.add_subplot(n_images, n_filters+1, i+1+(i*n_filters))
+        ax.imshow(image.permute(1,2,0).numpy())
+        ax.set_title('Original')
+        ax.axis('off')
+
+        for j in range(n_filters):
+            image = filtered_images[i][j]
+
+            if normalize:
+                image = normalize_image(image)
+
+            ax = fig.add_subplot(n_images, n_filters+1, i+1+(i*n_filters)+j+1)
+            ax.imshow(image.numpy(), cmap = 'bone')
+            ax.set_title(f'Filter {j+1}')
+            ax.axis('off');
+
+    fig.subplots_adjust(hspace = -0.7)
+    plt.savefig(fig_name)
+
+
+def plot_filters(fig_name, filters, normalize=True):
+    filters = filters.cpu()
+
+    n_filters = filters.shape[0]
+
+    rows = int(np.sqrt(n_filters))
+    cols = int(np.sqrt(n_filters))
+
+    fig = plt.figure(figsize=(30, 15))
+
+    for i in range(rows * cols):
+
+        image = filters[i]
+
+        if normalize:
+            image = normalize_image(image)
+
+        ax = fig.add_subplot(rows, cols, i + 1)
+        ax.imshow(image.permute(1, 2, 0))
+        ax.axis('off')
+
+    fig.subplots_adjust(wspace=-0.9)
+    plt.savefig(fig_name)
+
+def plot_tsne(fig_name, all_features, all_labels):
+    tsne = TSNE(n_components=2, random_state=42)
+    tsne_results = tsne.fit_transform(all_features)
+    plt.figure(figsize=(10, 7))
+    scatter = plt.scatter(tsne_results[:, 0], tsne_results[:, 1], c=all_labels, cmap='viridis', s=5)
+    plt.colorbar(scatter)
+    plt.title('t-SNE Visualization')
+    plt.show()
+    plt.savefig(fig_name)
+
+
+def plot_grad_cam(images, cams, predicted_labels, true_labels, classes, path):
+    fig, axs = plt.subplots(nrows=2, ncols=len(images), figsize=(20, 10))
+    
+    for i, (img, cam, pred_label, true_label) in enumerate(zip(images, cams, predicted_labels, true_labels)):
+        # Display the original image on the top row
+        axs[0, i].imshow(img.permute(1,2,0).cpu().numpy())
+        pred_class_name = classes[pred_label]
+        true_class_name = classes[true_label]
+        axs[0, i].set_title(f"Predicted: {pred_class_name}\nTrue: {true_class_name}", fontsize=12)
+        axs[0, i].axis('off')
+
+        # Add label to the leftmost plot
+        if i == 0:
+            axs[0, i].set_ylabel("Original Image", fontsize=14, rotation=90, labelpad=10)
+        
+        # Convert the original image to grayscale
+        grayscale_img = cv2.cvtColor(img.permute(1,2,0).cpu().numpy(), cv2.COLOR_RGB2GRAY)
+        grayscale_img = cv2.cvtColor(grayscale_img, cv2.COLOR_GRAY2RGB)
+
+        # Overlay the Grad-CAM heatmap on the grayscale image
+        heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)
+        heatmap = np.float32(heatmap) / 255
+        cam_img = heatmap + np.float32(grayscale_img)
+        cam_img = cam_img / np.max(cam_img)
+
+        # Display the Grad-CAM image on the bottom row
+        axs[1, i].imshow(cam_img)
+        axs[1, i].axis('off')
+
+        # Add label to the leftmost plot
+        if i == 0:
+            axs[1, i].set_ylabel("Grad-CAM", fontsize=14, rotation=90, labelpad=10)
+    
+    plt.tight_layout()
+    plt.savefig(path)
+    plt.close()
+

vis_confusion_mtx.py

@@ -0,0 +1,54 @@
+import torch
+
+import os
+from pathlib import Path
+
+from data_utils.data_tribology import TribologyDataset
+from utils.experiment_utils import get_model, get_prediction
+from utils.arg_utils import get_args
+from utils.visualization_utils import plot_confusion_matrix
+
+def generate_confusion_matrix(image_name, model, iterator, device):
+    labels, predictions = get_prediction(model, iterator, device)
+    plot_confusion_matrix('visualization_results/'+image_name+'_confusion_mtx.png', labels, predictions, classes=["ANTLER", "BEECHWOOD", "BEFOREUSE", "BONE", "IVORY","SPRUCEWOOD"])
+
+
+def main(args):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    model = get_model(args)
+
+    basepath=os.getcwd()
+    experiment_dir = Path(os.path.join(basepath,'experiments',args.model,args.resolution,args.magnification,args.modality,args.pretrained,args.frozen,args.vote))
+    if args.model == 'ViT':
+        experiment_dir = Path(os.path.join(basepath,'experiments','ViT_H',args.resolution,args.magnification,args.modality,args.pretrained,args.frozen,args.vote))
+    checkpoint_dir = Path(os.path.join(experiment_dir,'checkpoints'))
+    checkpoint_path = checkpoint_dir / f'epoch{str(args.epochs)}.pth'
+    model.load_state_dict(torch.load(checkpoint_path))
+    model = model.to(device)
+
+    train_csv_path = f"./LUA_Dataset/CSV/{args.resolution}_{args.magnification}_6w_train.csv"
+    test_csv_path = f"./LUA_Dataset/CSV/{args.resolution}_{args.magnification}_6w_test.csv"
+    img_path = f"./LUA_Dataset/{args.resolution}/{args.magnification}/{args.modality}"
+    BATCHSIZE = args.batch_size
+    train_dataset = TribologyDataset(csv_path = train_csv_path, img_path = img_path)
+    test_dataset = TribologyDataset(csv_path = test_csv_path, img_path = img_path)
+
+    means, stds = train_dataset.get_statistics()
+    train_dataset.prepare_transform(means, stds, mode='train')
+    test_dataset.prepare_transform(means, stds, mode='test')
+
+    test_iterator = torch.utils.data.DataLoader(test_dataset,
+                                            batch_size=BATCHSIZE, 
+                                            num_workers=4, 
+                                            shuffle=False, 
+                                            pin_memory=True,
+                                            drop_last=False)
+
+
+    generate_confusion_matrix(args.model, model, test_iterator, device)
+
+if __name__ == "__main__":
+    args = get_args()
+    main(args)
+

vote_analysis.py

@@ -0,0 +1,107 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import torch.optim as optim
+import torch.optim.lr_scheduler as lr_scheduler
+
+import torch.utils.data as data
+
+import numpy as np
+import random
+import tqdm
+import os
+from pathlib import Path
+
+from data_utils.data_tribology import TribologyDataset
+from utils.experiment_utils import get_model, get_name, get_logger, train, evaluate, evaluate_vote, evaluate_vote_analysis
+from utils.arg_utils import get_args
+
+def main(args):
+    '''Reproducibility'''
+    SEED = args.seed
+    random.seed(SEED)
+    np.random.seed(SEED)
+    torch.manual_seed(SEED)
+    torch.cuda.manual_seed(SEED)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+    '''Folder Creation'''
+    basepath=os.getcwd()
+    experiment_dir = Path(os.path.join(basepath,'experiments',args.model,args.resolution,args.magnification,args.modality,args.pretrained,args.frozen,args.vote))
+    experiment_dir.mkdir(parents=True, exist_ok=True)
+    checkpoint_dir = Path(os.path.join(experiment_dir,'checkpoints'))
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+
+    '''Logging'''
+    model_name = get_name(args)
+    print(model_name, 'STARTED')
+    
+    logger = get_logger(experiment_dir, 'vote_analysis')
+
+    '''Data Loading'''
+    train_csv_path = f"./LUA_Dataset/CSV/{args.resolution}_{args.magnification}_6w_train.csv"
+    test_csv_path = f"./LUA_Dataset/CSV/{args.resolution}_{args.magnification}_6w_test.csv"
+    img_path = f"./LUA_Dataset/{args.resolution}/{args.magnification}/{args.modality}"
+
+    # results_acc_1 = {}
+    # results_acc_3 = {}
+    # classes_num = 6
+    BATCHSIZE = args.batch_size
+    train_dataset = TribologyDataset(csv_path = train_csv_path, img_path = img_path)
+    test_dataset = TribologyDataset(csv_path = test_csv_path, img_path = img_path)
+
+    # prepare the data augmentation
+    means, stds = train_dataset.get_statistics()
+    train_dataset.prepare_transform(means, stds, mode='train')
+    test_dataset.prepare_transform(means, stds, mode='test')
+
+    VALID_RATIO = 0.1
+
+    num_train = len(train_dataset)
+    num_valid = int(VALID_RATIO * num_train)
+    train_dataset, valid_dataset = data.random_split(train_dataset, [num_train - num_valid, num_valid])
+    logger.info(f'Number of training samples: {len(train_dataset)}')
+    logger.info(f'Number of validation samples: {len(valid_dataset)}')
+
+    test_iterator = torch.utils.data.DataLoader(test_dataset,
+                                                batch_size=BATCHSIZE, 
+                                                num_workers=4, 
+                                                shuffle=False, 
+                                                pin_memory=True,
+                                                drop_last=False)
+    print('DATA LOADED')
+
+    # Define model 
+    model = get_model(args)
+    print('MODEL LOADED')
+
+    # Define device
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = model.to(device)
+
+
+    print('SETUP DONE')
+    # train our model
+
+    print('TRAINING STARTED')
+
+    model.load_state_dict(torch.load(checkpoint_dir / f'epoch{args.epochs}.pth'))
+    logger.info('-------------------Beginning of Testing-------------------')
+    print('TESTING STARTED')
+
+    vote_accuracy, correct_case_accuracy, incorrect_case_accuracy, incorrect_most_common, novote_accuracy = evaluate_vote_analysis(model, test_iterator, device)
+    logger.info(f'Test Acc @1: {vote_accuracy * 100:6.2f}%')
+    logger.info(f'No Vote Accuracy @1: {novote_accuracy * 100:6.2f}%')
+    logger.info(f'Correct Case Consistency @1: {correct_case_accuracy * 100:6.2f}%')
+    logger.info(f'Incorrect Case Consistency @1: {incorrect_case_accuracy * 100:6.2f}%')
+    logger.info(f'Incorrect Most Common: {incorrect_most_common* 100:6.2f}%')
+
+    logger.info('-------------------End of Testing-------------------')
+    print('TESTING DONE')
+
+
+if __name__ == '__main__':
+    args = get_args()
+    main(args)