vit_test.py

import glob
import time

import numpy as np
import pandas as pd
import torch
from PIL import Image
from matplotlib import pyplot as plt
from sklearn.metrics import precision_score, recall_score, confusion_matrix, ConfusionMatrixDisplay
from torchvision import transforms
from tqdm import tqdm

from src.model import ClassificationModel


def run_model(input_tensor):
    with torch.no_grad():
        out_val = model(input_tensor)
        out_val = torch.softmax(out_val, dim=1)
        out_val = out_val.detach().cpu().numpy().squeeze()

    return np.argmax(out_val, axis=-1)


start_time = time.time()
np.set_printoptions(suppress=True)
assert torch.cuda.is_available(), '!!!you are not using cuda!!!'

class_names = np.array(
    ['box', 'circle', 'clean', 'fall', 'run', 'walk']
)

exp_type = 'cir'
test_root = f'data/{exp_type}_dataset/test/'
test_list = glob.glob(f'{test_root}*/*.jpg')
assert len(test_list) > 0, 'test folder is empty'
print(len(test_list))

exp_name = f'{exp_type}_model'
ckpt_path = glob.glob(f'{exp_name}/lightning_logs/version_0/checkpoints/best*')[0]
model = ClassificationModel.load_from_checkpoint(checkpoint_path=ckpt_path)
model.cuda().eval()
print('loaded model', time.time() - start_time)

transforms_test = transforms.Compose(
    [
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
    ]
)

batch_img, gt_labels, pred_labels = [], [], []
for idx, test_file in tqdm(enumerate(test_list), total=len(test_list)):
    test_img = Image.open(test_file)
    test_img = transforms_test(test_img).unsqueeze(dim=0).cuda()
    batch_img.append(test_img)

    label = test_file.split('/')[-2]
    gt_labels.append(label)

    if len(batch_img) == 512 or idx == len(test_list) - 1:
        data = torch.cat(batch_img)
        class_idx = run_model(data)
        class_idx = class_names[class_idx]
        pred_labels.extend(class_idx)
        batch_img.clear()

y_true, y_pred = np.array(gt_labels), np.array(pred_labels)
precision = precision_score(y_true=y_true, y_pred=y_pred, average=None)
recall = recall_score(y_true=y_true, y_pred=y_pred, average=None)

results = []
for idx, label in enumerate(class_names):
    results.append((label, precision[idx], recall[idx]))
    print(results[-1])
results.append(('Mean', precision.mean(), recall.mean()))
print(results[-1])

results = pd.DataFrame(results, columns=['Label', 'Precision', 'Recall'])
results.to_csv(f'{exp_name}/{exp_name}_results.csv')

cm = confusion_matrix(y_true=y_true, y_pred=y_pred)
cm = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=class_names)
cm.plot(xticks_rotation='vertical', colorbar=False)
plt.tight_layout()
plt.savefig(f'{exp_name}/{exp_name}_cm.png', dpi=300, bbox_inches='tight')
print(exp_name, time.time() - start_time)