grad_cam.py 文件源码

def main():
    args = parse_args()

    if not os.path.exists(args.out_dir):
        os.makedirs(args.out_dir)

    target_layer_names = ['35']
    target_index = None

    # Prepare input image
    if args.img:
        img = cv2.imread(args.img, 1)
    else:
        img = misc.face()
    img = np.float32(cv2.resize(img, (224, 224))) / 255
    preprocessed_img = preprocess_image(img, args.cuda)

    model = vgg19(pretrained=True)
    if args.cuda:
        model.cuda()

    # Prediction
    output = model(preprocessed_img)
    pred_index = np.argmax(output.data.cpu().numpy())
    print('Prediction: {}'.format(IMAGENET_LABELS[pred_index]))

    # Prepare grad cam
    grad_cam = GradCam(
        pretrained_model=model,
        target_layer_names=target_layer_names,
        cuda=args.cuda)

        # Compute grad cam
    mask = grad_cam(preprocessed_img, target_index)

    save_cam_image(img, mask, os.path.join(args.out_dir, 'grad_cam.jpg'))
    print('Saved Grad-CAM image')

    # Reload preprocessed image
    preprocessed_img = preprocess_image(img)

    # Compute guided backpropagation
    guided_backprop = GuidedBackpropGrad(
        pretrained_model=model, cuda=args.cuda)
    guided_backprop_saliency = guided_backprop(preprocessed_img, index=target_index)

    cam_mask = np.zeros(guided_backprop_saliency.shape)
    for i in range(guided_backprop_saliency.shape[0]):
        cam_mask[i, :, :] = mask

    cam_guided_backprop = np.multiply(cam_mask, guided_backprop_saliency)
    save_as_gray_image(
        cam_guided_backprop,
        os.path.join(args.out_dir, 'guided_grad_cam.jpg'))
    print('Saved Guided Grad-CAM image')