def main():
parser = argparse.ArgumentParser(
description='Explainable VAE MNIST Example')
parser.add_argument('--result_dir',
type=str,
default='test_results',
metavar='DIR',
help='output directory')
parser.add_argument('--batch_size',
type=int,
default=128,
metavar='N',
help='input batch size for training (default: 128)')
parser.add_argument('--seed',
type=int,
default=1,
metavar='S',
help='random seed (default: 1)')
# model options
parser.add_argument('--latent_size',
type=int,
default=32,
metavar='N',
help='latent vector size of encoder')
parser.add_argument('--model_path',
type=str,
default='./ckpt/model_best.pth',
metavar='DIR',
help='pretrained model directory')
parser.add_argument('--one_class',
type=int,
default=8,
metavar='N',
help='outlier digit for one-class VAE testing')
args = parser.parse_args()
torch.manual_seed(args.seed)
kwargs = {'num_workers': 1, 'pin_memory': True} if cuda else {}
one_class = args.one_class # Choose the current outlier digit to be 8
one_mnist_test_dataset = OneClassMnist.OneMNIST(
'./data', one_class, train=False, transform=transforms.ToTensor())
test_loader = torch.utils.data.DataLoader(one_mnist_test_dataset,
batch_size=args.batch_size,
shuffle=False,
**kwargs)
model = ConvVAE(args.latent_size).to(device)
checkpoint = torch.load(args.model_path)
model.load_state_dict(checkpoint['state_dict'])
mu_avg, logvar_avg = 0, 1
gcam = GradCAM(model, target_layer='encoder.2', cuda=True)
test_index = 0
for batch_idx, (x, _) in enumerate(test_loader):
model.eval()
x = x.to(device)
x_rec, mu, logvar = gcam.forward(x)
model.zero_grad()
gcam.backward(mu, logvar, mu_avg, logvar_avg)
gcam_map = gcam.generate()
## Visualize and save attention maps ##
x = x.repeat(1, 3, 1, 1)
for i in range(x.size(0)):
raw_image = x[i] * 255.0
ndarr = raw_image.permute(1, 2, 0).cpu().byte().numpy()
im = Image.fromarray(ndarr.astype(np.uint8))
im_path = args.result_dir
if not os.path.exists(im_path):
os.mkdir(im_path)
im.save(
os.path.join(
im_path, "{}-{}-origin.png".format(test_index,
str(one_class))))
file_path = os.path.join(
im_path, "{}-{}-attmap.png".format(test_index, str(one_class)))
r_im = np.asarray(im)
save_cam(r_im, file_path, gcam_map[i].squeeze().cpu().data.numpy())
test_index += 1
def main(args):
# Load the synset words
file_name = 'synset_words.txt'
classes = list()
with open(file_name) as class_file:
for line in class_file:
classes.append(line.strip().split(' ', 1)[1].split(', ',
1)[0].replace(
' ', '_'))
print('Loading a model...')
model = torchvision.models.resnet152(pretrained=True)
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
print('\nGrad-CAM')
gcam = GradCAM(model=model,
target_layer='layer4.2',
n_class=1000,
cuda=args.cuda)
gcam.load_image(args.image, transform)
gcam.forward()
for i in range(0, 5):
gcam.backward(idx=gcam.idx[i])
cls_name = classes[gcam.idx[i]]
output = gcam.generate()
print('\t{:.5f}\t{}'.format(gcam.prob[i], cls_name))
gcam.save('results/{}_gcam.png'.format(cls_name), output)
print('\nBackpropagation')
bp = BackPropagation(model=model,
target_layer='conv1',
n_class=1000,
cuda=args.cuda)
bp.load_image(args.image, transform)
bp.forward()
for i in range(0, 5):
bp.backward(idx=bp.idx[i])
cls_name = classes[bp.idx[i]]
output = bp.generate()
print('\t{:.5f}\t{}'.format(bp.prob[i], cls_name))
bp.save('results/{}_bp.png'.format(cls_name), output)
print('\nGuided Backpropagation')
gbp = GuidedBackPropagation(model=model,
target_layer='conv1',
n_class=1000,
cuda=args.cuda)
gbp.load_image(args.image, transform)
gbp.forward()
for i in range(0, 5):
cls_idx = gcam.idx[i]
cls_name = classes[cls_idx]
gcam.backward(idx=cls_idx)
output_gcam = gcam.generate()
gbp.backward(idx=cls_idx)
output_gbp = gbp.generate()
output_gcam -= output_gcam.min()
output_gcam /= output_gcam.max()
output_gcam = cv2.resize(output_gcam, (224, 224))
output_gcam = cv2.cvtColor(output_gcam, cv2.COLOR_GRAY2BGR)
output = output_gbp * output_gcam
print('\t{:.5f}\t{}'.format(gbp.prob[i], cls_name))
gbp.save('results/{}_gbp.png'.format(cls_name), output_gbp)
gbp.save('results/{}_ggcam.png'.format(cls_name), output)
def main(args):
"""
Main Function for testing and saving attention maps.
Inputs:
args - Namespace object from the argument parser
"""
torch.manual_seed(args.seed)
# Load dataset
if args.dataset == 'mnist':
test_dataset = OneClassMnist.OneMNIST('./data',
args.one_class,
train=False,
transform=transforms.ToTensor())
elif args.dataset == 'ucsd_ped1':
test_dataset = Ped1_loader.UCSDAnomalyDataset('./data',
train=False,
resize=args.image_size)
elif args.dataset == 'mvtec_ad':
class_name = mvtec.CLASS_NAMES[args.one_class]
test_dataset = mvtec.MVTecDataset(class_name=class_name,
is_train=False,
grayscale=False,
root_path=args.data_path)
test_steps = len(test_dataset)
kwargs = {
'num_workers': args.num_workers,
'pin_memory': True
} if device == "cuda" else {}
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=args.batch_size,
shuffle=True,
**kwargs)
# Select a model architecture
if args.model == 'vanilla_mnist':
imshape = [1, 28, 28]
model = ConvVAE_mnist(args.latent_size).to(device)
elif args.model == 'vanilla_ped1':
imshape = [1, args.image_size, args.image_size]
model = ConvVAE_ped1(args.latent_size, args.image_size,
args.batch_norm).to(device)
elif args.model == 'resnet18_3':
imshape = [3, 256, 256]
model = ResNet18VAE_3(args.latent_size,
x_dim=imshape[-1],
nc=imshape[0]).to(device)
print("Layer is:", args.target_layer)
# Load model
checkpoint = torch.load(args.model_path)
model.load_state_dict(checkpoint['state_dict'])
mu_avg, logvar_avg = (0, 1)
gcam = GradCAM(model, target_layer=args.target_layer, device=device)
prediction_stack = np.zeros((test_steps, imshape[-1], imshape[-1]),
dtype=np.float32)
gt_mask_stack = np.zeros((test_steps, imshape[-1], imshape[-1]),
dtype=np.uint8)
# Generate attention maps
for batch_idx, (x, y) in enumerate(test_loader):
# print("batch_idx", batch_idx)
model.eval()
x = x.to(device)
x_rec, mu, logvar = gcam.forward(x)
model.zero_grad()
gcam.backward(mu, logvar, mu_avg, logvar_avg)
gcam_map = gcam.generate()
gcam_max = torch.max(gcam_map).item()
# If image has one channel, make it three channel(need for heatmap)
if x.size(1) == 1:
x = x.repeat(1, 3, 1, 1)
# Visualize and save attention maps
for i in range(x.size(0)):
x_arr = x[i].permute(1, 2, 0).cpu().numpy() * 255
x_im = Image.fromarray(x_arr.astype(np.uint8))
# Get the gradcam for this image
prediction = gcam_map[i].squeeze().cpu().data.numpy()
# Add prediction and mask to the stacks
prediction_stack[batch_idx * args.batch_size + i] = prediction
gt_mask_stack[batch_idx * args.batch_size + i] = y[i]
if save_gcam_image:
im_path = args.result_dir
if not os.path.exists(im_path):
os.mkdir(im_path)
x_im.save(
os.path.join(im_path,
"{}-{}-origin.png".format(batch_idx, i)))
file_path = os.path.join(
im_path, "{}-{}-attmap.png".format(batch_idx, i))
save_gradcam(x_arr, file_path, prediction, gcam_max=gcam_max)
# Stop of dataset is mnist because there aren't GTs available
if args.dataset != 'mnist':
# Compute area under the ROC score
auc = roc_auc_score(gt_mask_stack.flatten(),
prediction_stack.flatten())
print(f"AUROC score: {auc}")
fpr, tpr, thresholds = roc_curve(gt_mask_stack.flatten(),
prediction_stack.flatten())
if plot_ROC:
plt.plot(tpr, fpr, label="ROC")
plt.xlabel("FPR")
plt.ylabel("TPR")
plt.legend()
plt.savefig(
str(args.result_dir) + "auroc_" + str(args.model) +
str(args.target_layer) + str(args.one_class) + ".png")
# Compute IoU
if args.iou == True:
print(f"IoU score: {j_score}")
max_val = np.max(prediction_stack)
max_steps = 100
best_thres = 0
best_iou = 0
# Ge the IoU for 100 different thresholds
for i in range(1, max_steps):
thresh = i / max_steps * max_val
prediction_bin_stack = prediction_stack > thresh
iou = jaccard_score(gt_mask_stack.flatten(),
prediction_bin_stack.flatten())
if iou > best_iou:
best_iou = iou
best_thres = thresh
print("Best threshold;", best_thres)
print("Best IoU score:", best_iou)
return
def gradCAM():
# Chap 2 : Train Network
print('\n[Chapter 2] : Operate gradCAM with trained network')
# Phase 1 : Model Upload
print('\n[Phase 1] : Model Weight Upload')
use_gpu = torch.cuda.is_available()
# upload labels
data_dir = cf.test_dir
trainset_dir = cf.data_base.split("/")[-1] + os.sep
dsets = datasets.ImageFolder(data_dir, None)
H = datasets.ImageFolder(cf.aug_base + '/train/')
dset_classes = H.classes
def softmax(x):
return np.exp(x) / np.sum(np.exp(x), axis=0)
# return np.exp(x) / np.sum(np.exp(x), axis=1)
def getNetwork(opts):
if (opts.net_type == 'alexnet'):
file_name = 'alexnet'
elif (opts.net_type == 'vggnet'):
file_name = 'vgg-%s' % (opts.depth)
elif (opts.net_type == 'resnet'):
file_name = 'resnet-%s' % (opts.depth)
else:
print('[Error]: Network should be either [alexnet / vgget / resnet]')
sys.exit(1)
return file_name
def random_crop(image, dim):
if len(image.shape):
W, H, D = image.shape
w, h, d = dim
else:
W, H = image.shape
w, h = dim[0], dim[1]
left, top = np.random.randint(W - w + 1), np.random.randint(H - h + 1)
return image[left:left + w, top:top + h], left, top
# uploading the model
print("| Loading checkpoint model for grad-CAM...")
assert os.path.isdir('./path'), '[Error]: No checkpoint directory found!'
assert os.path.isdir('./path/' + trainset_dir), '[Error]: There is no model weight to upload!'
file_name = getNetwork(opts)
checkpoint = torch.load('./path/' + trainset_dir + file_name + '.t7')
model = checkpoint['model']
if use_gpu:
model.cuda()
cudnn.benchmark = True
model.eval()
sample_input = Variable(torch.randn(1, 3, 224, 224), volatile=False)
if use_gpu:
sampe_input = sample_input.cuda()
def is_image(f):
return f.endswith(".png") or f.endswith(".jpg")
test_transform = transforms.Compose([
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(cf.mean, cf.std)
])
"""
#@ Code for inference test
img = Image.open(cf.image_path)
if test_transform is not None:
img = test_transform(img)
inputs = img
inputs = Variable(inputs, volatile=False, requires_grad=True)
if use_gpu:
inputs = inputs.cuda()
inputs = inputs.view(1, inputs.size(0), inputs.size(1), inputs.size(2))
outputs = model(inputs)
softmax_res = softmax(outputs.data.cpu().numpy()[0])
index,score = max(enumerate(softmax_res), key=operator.itemgetter(1))
print('| Uploading %s' %(cf.image_path.split("/")[-1]))
print('| prediction = ' + dset_classes[index])
"""
# @ Code for extracting a grad-CAM region for a given class
gcam = GradCAM(list(model._modules.items())[0][1],
cuda=use_gpu) # model=model._modules.items()[0][1], cuda=use_gpu)
gbp = GuidedBackPropagation(model=list(model._modules.items())[0][1], cuda=use_gpu)
# print(dset_classes)
WBC_id = 5 # BHX class
print("Checking Activated Regions for " + dset_classes[WBC_id] + "...")
fileList = os.listdir('./samples/')
i = 1
for f in fileList:
file_name = './samples/' + f
print("Opening " + file_name + "...")
original_image = cv2.imread(file_name)
resize_ratio = 224. / min(original_image.shape[0:2])
resized = cv2.resize(original_image, (0, 0), fx=resize_ratio, fy=resize_ratio)
cropped, left, top = random_crop(resized, (224, 224, 3))
print(cropped.size)
if test_transform is not None:
img = test_transform(Image.fromarray(cropped, mode='RGB'))
# center_cropped = original_image[16:240, 16:240, :]
# expand the image based on the short side
inputs = img
inputs = Variable(inputs, requires_grad=True)
if use_gpu:
inputs = inputs.cuda()
inputs = inputs.view(1, inputs.size(0), inputs.size(1), inputs.size(2))
probs, idx = gcam.forward(inputs)
# probs, idx = gbp.forward(inputs)
# Grad-CAM
gcam.backward(idx=WBC_id)
if opts.depth == 18:
output = gcam.generate(target_layer='layer4.1')
else:
output = gcam.generate(target_layer='layer4.2') # a module name to be visualized (required)
# Guided Back Propagation
# gbp.backward(idx=WBC_id)
# feature = gbp.generate(target_layer='conv1')
# Guided Grad-CAM
# output = np.multiply(feature, region)
gcam.save('./results/%s.png' % str(i), output, cropped)
cv2.imwrite('./results/map%s.png' % str(i), cropped)
for j in range(3):
print('\t{:5f}\t{}\n'.format(probs[j], dset_classes[idx[j]]))
i += 1
"""