def create_cam(config):
if not os.path.exists(config.result_path):
os.mkdir(config.result_path)
test_loader, num_class = utils.get_testloader(config.dataset,
config.dataset_path,
config.img_size)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
mdl = model.NLayerDiscriminator()
#mdl = nn.DataParallel(mdl)
cnn = mdl.to(device)
state_dict = torch.load(os.path.join(config.model_path, config.model_name))
cnn.load_state_dict(state_dict)
finalconv_name = 'conv'
# hook
feature_blobs = []
def hook_feature(module, input, output):
feature_blobs.append(output.cpu().data.numpy())
cnn._modules.get(finalconv_name).register_forward_hook(hook_feature)
params = list(cnn.parameters())
# get weight only from the last layer(linear)
weight_softmax = np.squeeze(params[-2].cpu().data.numpy())
def returnCAM(feature_conv, weight_softmax, class_idx):
size_upsample = (config.img_size, config.img_size)
_, nc, h, w = feature_conv.shape
output_cam = []
cam = weight_softmax[class_idx].dot(feature_conv.reshape((nc, h * w)))
cam = cam.reshape(h, w)
cam = cam - np.min(cam)
cam_img = cam / np.max(cam)
cam_img = np.uint8(255 * cam_img)
output_cam.append(cv2.resize(cam_img, size_upsample))
return output_cam
for i, (image_tensor, label) in enumerate(test_loader):
image_PIL = transforms.ToPILImage()(image_tensor[0])
image_PIL.save(os.path.join(config.result_path, 'img%d.png' % (i + 1)))
image_tensor = image_tensor.to(device)
logit, _ = cnn(image_tensor)
h_x = F.softmax(logit, dim=1).data.squeeze()
probs, idx = h_x.sort(0, True)
print("True label : %d, Predicted label : %d, Probability : %.2f" %
(label.item(), idx[0].item(), probs[0].item()))
CAMs = returnCAM(feature_blobs[0], weight_softmax, [idx[0].item()])
img = cv2.imread(
os.path.join(config.result_path, 'img%d.png' % (i + 1)))
height, width, _ = img.shape
heatmap = cv2.applyColorMap(cv2.resize(CAMs[0], (width, height)),
cv2.COLORMAP_JET)
result = heatmap * 0.3 + img * 0.5
cv2.imwrite(os.path.join(config.result_path, 'cam%d.png' % (i + 1)),
result)
if i + 1 == config.num_result:
break
feature_blobs.clear()
def load_models(directory, batch_num):
generator = model.GlobalGenerator()
discriminator = model.NLayerDiscriminator(input_nc=3)
gen_name = os.path.join(directory, '%05d_generator.pth' % batch_num)
dis_name = os.path.join(directory, '%05d_discriminator.pth' % batch_num)
if os.path.isfile(gen_name) and os.path.isfile(dis_name):
gen_dict = torch.load(gen_name)
dis_dict = torch.load(dis_name)
generator.load_state_dict(gen_dict)
discriminator.load_state_dict(dis_dict)
print('Models loaded, resume training from batch %05d...' % batch_num)
else:
print('Cannot find saved models, start training from scratch...')
batch_num = 0
return generator, discriminator, batch_num
def load_models(directory, batch_num):
# 20180924: smaller network.
generator = model.GlobalGenerator(n_downsampling=2, n_blocks=6)
discriminator = model.NLayerDiscriminator(input_nc=3, n_layers=3) # 48 input
gen_name = os.path.join(directory, '%05d_generator.pth' % batch_num)
dis_name = os.path.join(directory, '%05d_discriminator.pth' % batch_num)
if os.path.isfile(gen_name) and os.path.isfile(dis_name):
gen_dict = torch.load(gen_name)
dis_dict = torch.load(dis_name)
generator.load_state_dict(gen_dict)
discriminator.load_state_dict(dis_dict)
print('Models loaded, resume training from batch %05d...' % batch_num)
else:
print('Cannot find saved models, start training from scratch...')
batch_num = 0
return generator, discriminator, batch_num
torch.manual_seed(0)
if torch.cuda.is_available():
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# prepare neural networks
model_g = model.Generator(n_res_blocks=opt.num_res_blocks)
print('# generator parameters:',
sum(param.numel() for param in model_g.parameters()))
model_d = model.Discriminator(kernel_size=opt.kernel_size,
gaussian=opt.gaussian,
wgan=opt.wgan,
highpass=opt.highpass)
print('# discriminator parameters:',
sum(param.numel() for param in model_d.parameters()))
model_d_save = model.NLayerDiscriminator(3, n_layers=2)
DWT2 = DWTForward(J=1, mode='zero', wave='haar').cuda()
g_loss_module = loss.GeneratorLoss(**vars(opt))
# filters are used for generating validation images
filter_low_module = model.FilterLow(kernel_size=opt.kernel_size,
gaussian=opt.gaussian,
include_pad=False)
filter_high_module = model.FilterHigh(kernel_size=opt.kernel_size,
gaussian=opt.gaussian,
include_pad=False)
if torch.cuda.is_available():
model_g = model_g.cuda()
model_d = model_d.cuda()
model_d_save = model_d_save.cuda()
