visualize.py 文件源码

def main():

    # Model class options
    model_parser = argparse.ArgumentParser(description='Model Parameters', add_help=False)
    model_parser.add_argument('--model_name', type=str, help='Model name {"SimpleCNN", "MiddleCNN"}')
    model_parser.add_argument('--init_model', type=str, help='Initialize the model from given file')
    model_parser.add_argument('--n_classes', type=int, default=48, help='Number of classes')
    model_args, remaining_argv = model_parser.parse_known_args()

    # Model runtime options
    runtime_parser = argparse.ArgumentParser(description='Runtime Parameters', add_help=False)
    runtime_parser.add_argument('--gpu', type=int, default=-1, help='GPU ID (negative value indicates CPU')
    runtime_parser.add_argument('--input_image', type=str, help='/path/to/input_image.jpg')
    runtime_parser.add_argument('--target_label', type=str, help='If not specified, predicted label is used as target.')    
    runtime_parser.add_argument('--save_dir', type=str, default='./grad_cam', help='Save directory')
    runtime_args, remaining_argv = runtime_parser.parse_known_args(remaining_argv)

    # merge options
    parser = argparse.ArgumentParser(
        description='Visualize Saliency',
        parents=[model_parser, runtime_parser])
    parser.add_argument('--debug', action='store_true', help='if specified, using chainer.set_debug()')
    args = parser.parse_args()

    chainer.set_debug(args.debug)
    assert runtime_args.input_image is not None, "input_image must be specified."
    assert model_args.init_model is not None, "init_model must be specified."


    ''' Visualization '''
    # read image
    img = imread(args.input_image, mode='RGB').astype(np.float32)

    # visualize
    visualize(
        img,
        args.target_label,
        args.model_name,
        args.init_model,
        args.n_classes,
        args.save_dir,
        args.gpu
    )