1. import os
    2. import logging
    3. import warnings
    4. import torch
    5. import torch.nn as nn
    6. import torch.nn.functional as F
    7. from torchvision import transforms
    8. from torch.utils.data import DataLoader
    9. from tqdm import tqdm
    11. from PIL import Image
    12. import numpy as np
    13. import torch.utils.data as data
    14. from dataset_fine_02 import dataset
    15. from model.resnet_p2p_02 import ResNet50
    16. from eval_model_02 import eval_turn
    17. from ZYJ_utils.utils_02 import calc_map_k, label2onehot, calc_train_codes, batch_augment, generate_heatmap, aug_augment
    18. from ZYJ_utils.utils_1 import NCESoftmaxLoss
    19. from utils import show_cam_on_image
    20. # settings#################################################
    21. assert torch.cuda.is_available()
    22. os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    23. device = torch.device("cuda")
    24. torch.backends.cudnn.benchmark = True
    25. savepath = './visualize/'
    26. weights_path = './training_checkpoint_bit_class/weights_0.pth'
    27. dataset_name = 'cub_bird'
    28. img_path = "bird3.jpg"
    29. labels = torch.tensor([0])
    30. batch_size = 1
    31. os.makedirs(savepath, exist_ok=True)
    32. ToPILImage = transforms.ToPILImage()
    33. MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
    34. STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
    35. ###########################################################
    37. def main():
    38.     logging.basicConfig(
    39.         format='%(asctime)s: %(levelname)s: [%(filename)s:%(lineno)d]: %(message)s',
    40.         level=logging.INFO)
    41.     warnings.filterwarnings("ignore")
    42.     # Dataset for testing#################################
    43.     if dataset_name == 'cub_bird':
    44.         classes = 200
    45.         data_dir = './datasets/cub_bird/'
    46.     elif dataset_name == 'stanford_dog':
    47.         classes = 120
    48.         data_dir = './datasets/stanford_dog/'
    49.     elif dataset_name == 'aircraft':
    50.         classes = 100
    51.         data_dir = './datasets/aircraft/'
    52.     elif dataset_name == 'vegfru':
    53.         classes = 292
    54.         data_dir = './datasets/vegfru/'
    55.     else:
    56.         print('undefined dataset ! ')
    57.     base_set = dataset(dataset_name, root_dir=data_dir, train=True)
    58.     print('basa_set', len(base_set))
    59.     dataloader = {}
    60.     dataloader['base'] = data.DataLoader(base_set, batch_size=batch_size, shuffle=False, num_workers=8, pin_memory=True)
    61.     #########################################################
    62.     ##################################
    63.     # Initialize model
    64.     ##################################
    65.     net = ResNet50(32, classes, 0.7)
    66.     net.load_state_dict(torch.load(weights_path, map_location=device), strict=False)
    67.     net.to(device)
    68.     if torch.cuda.device_count() > 1:
    69.         net = nn.DataParallel(net)
    70.     net.eval()
    71.     # load_image########################################################
    72.     assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
    73.     img = Image.open(img_path).convert('RGB')
    74.     img_np = transforms.RandomResizedCrop(224)(img)
    75.     img_np = np.array(img_np, dtype=np.uint8)  # (224, 224, 3)
    77.     # [C, H, W]
    78.     img_tensor, aug_sample, sample, i, j, h, w, hor_flip = aug_augment(img)
    79.     # expand batch dimension
    80.     # [C, H, W] -> [N, C, H, W]
    81.     input_tensor = torch.unsqueeze(img_tensor, dim=0)
    82.     aug_input_tensor = torch.unsqueeze(aug_sample, dim=0)
    83.     sample_input_tensor = torch.unsqueeze(sample, dim=0)
    84.     X = input_tensor.cuda()  # torch.Size([1, 3, 223, 320])
    85.     aug_X = aug_input_tensor.cuda()  # torch.Size([1, 3, 223, 320])
    86.     sample_X = sample_input_tensor.cuda()
    87.     ##########################################################
    88.     feature_map, y3, code_q, c3, alpha1, alpha2 = net(X)
    89.     with torch.no_grad():
    90.         drop_images = batch_augment(X, labels, feature_map, mode='drop')
    91.     with torch.no_grad():
    92.         crop_images, attention_maps = batch_augment(X, labels, feature_map, mode='crop')
    93.     with torch.no_grad():
    94.         zoom_images = batch_augment(X, labels, feature_map, mode='zoom')
    95.     raw_image = X.cpu() * STD + MEAN
    96.     sample_X_image = sample_X.cpu() * STD + MEAN
    97.     aug_raw_image = aug_X.cpu() * STD + MEAN
    98.     crop_images = crop_images.cpu() * STD + MEAN
    99.     zoom_images = zoom_images.cpu() * STD + MEAN
    100.     drop_images = drop_images.cpu() * STD + MEAN
    101.     rimg = ToPILImage(raw_image[0])  # pic should be 2/3 dimensional. Got 4 dimensions.
    102.     sampleimg = ToPILImage(sample_X_image[0])
    103.     augrimg = ToPILImage(aug_raw_image[0])
    104.     crimg = ToPILImage(crop_images[0])
    105.     zoimg = ToPILImage(zoom_images[0])
    106.     drimg = ToPILImage(drop_images[0])
    107.     rimg.save(os.path.join(savepath, '_raw'+img_path))
    108.     sampleimg.save(os.path.join(savepath, 'ogimg'+img_path))
    109.     augrimg.save(os.path.join(savepath, '_augrimg'+img_path))
    110.     crimg.save(os.path.join(savepath, img_path+'_raw_atten_cropp.jpg'))
    111.     zoimg.save(os.path.join(savepath, img_path+'_raw_atten_zoom.jpg'))
    112.     drimg.save(os.path.join(savepath, img_path+'_raw_atten_drop.jpg'))
    113.     print('ok')
    115. if __name__ == '__main__':
    116.     main()