Exemple #1
0
def main(args):
    if args.im1 is None or args.im2 is None:
        print("Please enter two image paths through cmd")
        exit(-1)
    print("Comparing {} and {}".format(args.im1, args.im2))
    im1 = Image.open(args.im1)
    im2 = Image.open(args.im2)
    im1 = np.array(im1)
    im2 = np.array(im2)
    psnr = evaluate.my_psnr(im1, im2)
    ssim = evaluate.ssim(im1, im2)
    print("PSNR {:4f}, SSIM {:4f}".format(psnr, ssim))

    if args.save:
        diff = image_warping.get_diff_image(im1, im2)
        file_name = 'DifferenceImage{}{}.png'.format(args.im1[-10:-4],
                                                     args.im2[-10:-4])
        image_warping.save_array_as_image(diff, file_name)
def main(args, config, sample_index):
    cuda = cnn_utils.check_cuda(config)
    model = cnn_utils.load_model_and_weights(args, config)
    if cuda:
        model = model.cuda()

    model.eval()

    # Create output directory
    base_dir = os.path.join(config['PATH']['output_dir'], 'warped')
    if not os.path.isdir(base_dir):
        pathlib.Path(base_dir).mkdir(parents=True, exist_ok=True)
    save_dir = get_sub_dir_for_saving(base_dir)

    start_time = time.time()
    file_path = os.path.join(config['PATH']['hdf5_dir'],
                             config['PATH']['hdf5_name'])
    with h5py.File(file_path, mode='r', libver='latest') as hdf5_file:
        depth_grp = hdf5_file['val']['disparity']
        SNUM = sample_index
        depth_images = torch.squeeze(
            torch.tensor(depth_grp['images'][SNUM], dtype=torch.float32))

        colour_grp = hdf5_file['val']['colour']
        colour_images = torch.tensor(colour_grp['images'][SNUM],
                                     dtype=torch.float32)

        sample = {
            'depth': depth_images,
            'colour': colour_images,
            'grid_size': depth_images.shape[0]
        }

        warped = data_transform.center_normalise(sample)
        im_input = warped['inputs'].unsqueeze_(0)

        if cuda:
            im_input = im_input.cuda()

        output = model(im_input)

        time_taken = time.time() - start_time
        print("Time taken was {:4f}s".format(time_taken))
        grid_size = 64

        psnr_accumulator = (0, 0, 0)
        ssim_accumulator = (0, 0, 0)

        print("Saving output to", save_dir)

        output = torch.squeeze(denormalise_lf(output))
        cpu_output = np.around(output.cpu().detach().numpy()).astype(np.uint8)

        if (not args.no_eval) or args.get_diff:
            ground_truth = np.around(
                denormalise_lf(colour_images).numpy()).astype(np.uint8)

        grid_len = int(math.sqrt(grid_size))
        for i in range(grid_size):
            row, col = i // grid_len, i % grid_len

            file_name = 'Colour{}{}.png'.format(row, col)
            save_location = os.path.join(save_dir, file_name)
            if i == 0:
                print("Saving images of size ", cpu_output[i].shape)
            image_warping.save_array_as_image(cpu_output[i], save_location)

            if args.get_diff:
                colour = ground_truth[i]
                diff = image_warping.get_diff_image(colour, cpu_output[i])
                #diff = get_diff_image_floatint(res, colour)
                file_name = 'Diff{}{}.png'.format(row, col)
                save_location = os.path.join(save_dir, file_name)
                image_warping.save_array_as_image(diff, save_location)

            if not args.no_eval:
                img = ground_truth[i]
                file_name = 'GT_Colour{}{}.png'.format(row, col)
                save_location = os.path.join(save_dir, file_name)
                image_warping.save_array_as_image(img, save_location)
                psnr = evaluate.my_psnr(cpu_output[i], img)
                ssim = evaluate.ssim(cpu_output[i], img)
                psnr_accumulator = welford.update(psnr_accumulator, psnr)
                ssim_accumulator = welford.update(ssim_accumulator, ssim)

        psnr_mean, psnr_var, _ = welford.finalize(psnr_accumulator)
        ssim_mean, ssim_var, _ = welford.finalize(ssim_accumulator)
        print("For cnn, psnr average {:5f}, stddev {:5f}".format(
            psnr_mean, math.sqrt(psnr_var)))
        print("For cnn, ssim average {:5f}, stddev {:5f}".format(
            ssim_mean, math.sqrt(ssim_var)))
def do_one_demo(args, config, hdf5_file, model, sample_num, cuda):
    depth_grp = hdf5_file['val']['disparity']
    # Create output directory

    if not args.no_save:
        base_dir = os.path.join(config['PATH']['output_dir'], 'warped')
        if not os.path.isdir(base_dir):
            pathlib.Path(base_dir).mkdir(parents=True, exist_ok=True)
        save_dir = get_sub_dir_for_saving(base_dir)

    model.eval()
    SNUM = sample_num
    start_time = time.time()
    print("Working on image", SNUM)
    depth_images = torch.squeeze(
        torch.tensor(depth_grp['images'][SNUM], dtype=torch.float32))

    colour_grp = hdf5_file['val']['colour']
    colour_images = torch.tensor(colour_grp['images'][SNUM],
                                 dtype=torch.float32)

    sample = {
        'depth': depth_images,
        'colour': colour_images,
        'grid_size': depth_images.shape[0]
    }

    warped = data_transform.transform_to_warped(sample)
    warped = data_transform.stack(warped, 65)
    im_input = warped['inputs'].unsqueeze_(0)

    if cuda:
        im_input = im_input.cuda()

    output = model(im_input)
    im_input = im_input[:, :-3]
    output += im_input
    output = torch.clamp(output, 0.0, 1.0)

    time_taken = time.time() - start_time
    print("Time taken was {:.0f}s".format(time_taken))
    grid_size = 64

    psnr_accumulator = (0, 0, 0)
    ssim_accumulator = (0, 0, 0)

    if not args.no_save:
        print("Saving output to", save_dir)
        no_cnn_dir = os.path.join(save_dir, "no_cnn")
        cnn_dir = os.path.join(save_dir, "cnn")
        os.mkdir(cnn_dir)
        os.mkdir(no_cnn_dir)

    output = torch.squeeze(denormalise_lf(output))
    output = data_transform.torch_unstack(output)
    im_input = im_input.cpu().detach()
    cpu_output = np.around(output.cpu().detach().numpy()).astype(np.uint8)

    if (not args.no_eval) or args.get_diff:
        ground_truth = np.around(denormalise_lf(colour_images).numpy()).astype(
            np.uint8)
    grid_len = int(math.sqrt(grid_size))
    for i in range(grid_size):
        row, col = i // grid_len, i % grid_len

        if not args.no_save:
            file_name = 'Colour{}{}.png'.format(row, col)
            save_location = os.path.join(cnn_dir, file_name)
            if i == 0:
                print("Saving images of size ", cpu_output[i].shape)
            image_warping.save_array_as_image(cpu_output[i], save_location)

        if args.get_diff and not args.no_save:
            colour = ground_truth[i]
            diff = image_warping.get_diff_image(colour, cpu_output[i])
            #diff = get_diff_image_floatint(res, colour)
            file_name = 'Diff{}{}.png'.format(row, col)
            save_location = os.path.join(cnn_dir, file_name)
            image_warping.save_array_as_image(diff, save_location)

        if not args.no_eval:
            img = ground_truth[i]

            if not args.no_save:
                file_name = 'GT_Colour{}{}.png'.format(row, col)
                save_location = os.path.join(save_dir, file_name)
                image_warping.save_array_as_image(img, save_location)
            psnr = evaluate.my_psnr(cpu_output[i], img)
            ssim = evaluate.ssim(cpu_output[i], img)
            psnr_accumulator = welford.update(psnr_accumulator, psnr)
            ssim_accumulator = welford.update(ssim_accumulator, ssim)

    psnr_mean, psnr_var, _ = welford.finalize(psnr_accumulator)
    ssim_mean, ssim_var, _ = welford.finalize(ssim_accumulator)
    print("For cnn, psnr average {:5f}, stddev {:5f}".format(
        psnr_mean, math.sqrt(psnr_var)))
    print("For cnn, ssim average {:5f}, stddev {:5f}".format(
        ssim_mean, math.sqrt(ssim_var)))
    psnr1 = psnr_mean
    ssim1 = ssim_mean

    psnr_accumulator = (0, 0, 0)
    ssim_accumulator = (0, 0, 0)

    psnr2, ssim2 = 0, 0
    if args.no_cnn:
        squeeze_input = torch.squeeze(denormalise_lf(im_input))
        squeeze_input = data_transform.torch_unstack(squeeze_input)
        cpu_input = np.around(squeeze_input.numpy()).astype(np.uint8)
        for i in range(grid_size):
            row, col = i // grid_len, i % grid_len

            if not args.no_save:
                file_name = 'Colour{}{}.png'.format(row, col)
                save_location = os.path.join(no_cnn_dir, file_name)
                if i == 0:
                    print("Saving images of size ", cpu_input[i].shape)
                image_warping.save_array_as_image(cpu_input[i], save_location)

            if args.get_diff and not args.no_save:
                colour = ground_truth[i]
                diff = image_warping.get_diff_image(colour, cpu_output[i])
                #diff = get_diff_image_floatint(res, colour)
                file_name = 'Diff{}{}.png'.format(row, col)
                save_location = os.path.join(no_cnn_dir, file_name)
                image_warping.save_array_as_image(diff, save_location)

            if not args.no_eval:
                img = ground_truth[i]
                psnr = evaluate.my_psnr(cpu_input[i], img)
                ssim = evaluate.ssim(cpu_input[i], img)
                psnr_accumulator = welford.update(psnr_accumulator, psnr)
                ssim_accumulator = welford.update(ssim_accumulator, ssim)

        psnr_mean, psnr_var, _ = welford.finalize(psnr_accumulator)
        ssim_mean, ssim_var, _ = welford.finalize(ssim_accumulator)
        print("For no cnn, psnr average {:5f}, stddev {:5f}".format(
            psnr_mean, math.sqrt(psnr_var)))
        print("For no cnn, ssim average {:5f}, stddev {:5f}".format(
            ssim_mean, math.sqrt(ssim_var)))
        psnr2, ssim2 = psnr_mean, ssim_mean

    return psnr1, ssim1, psnr2, ssim2