demo_refocus.py

import torch
import torch.nn.functional as F
from torchvision import transforms

import PIL
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

import argparse
import os.path
from pathlib import Path
import glob
import sys
import pdb

from data.refocus_augmentation import RefocusImageAugmentation
from data.transforms import get_transform

parser = argparse.ArgumentParser(description='Visualize 3D refocus augmentation')

parser.add_argument('--num_quantiles', dest='num_quantiles', \
    help="number of qualtiles to use in blur stack. More is better, but slower", required=False)
parser.set_defaults(num_quantiles=10)

parser.add_argument('--min_aperture', dest='min_aperture', help="smallest aperture to use", required=False)
parser.set_defaults(min_aperture=0.001)

parser.add_argument('--max_aperture', dest='max_aperture', help="largest aperture to use", required=False)
parser.set_defaults(max_aperture=6)

parser.add_argument('--input_path', dest='input_path', \
    help="path to folder containing rgb and depth_euclidean")
parser.set_defaults(im_name='NONE')

parser.add_argument('--output_path', dest='output_path',\
    help="path to where refocused rgb should be stored")
    
parser.set_defaults(store_name='NONE')

args = parser.parse_args()

transform_rgb = get_transform('rgb', image_size=512)
transform_depth = get_transform('depth_euclidean', image_size=512)
trans_topil = transforms.ToPILImage()

os.system(f"mkdir -p {args.output_path}")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

refocus_aug = RefocusImageAugmentation(
    args.num_quantiles, float(args.min_aperture), float(args.max_aperture), return_segments=False)


def save_outputs(img_path, output_file_name):
    if not output_file_name.__contains__('rgb'): return
    rgb_path = img_path
    depth_path = img_path.replace('rgb', 'depth_euclidean')
    save_path = args.output_path+'/'+output_file_name+'_refocused'+'.png'

    print(f'Reading input {img_path} ...')
    rgb = Image.open(rgb_path)
    rgb_tensor = transform_rgb(rgb)[:3].unsqueeze(0).to(device)
    depth = Image.open(depth_path)
    depth_tensor = transform_depth(depth)[:1].unsqueeze(0).to(device)

    if rgb_tensor.shape[1] == 1:
        rgb_tensor = rgb_tensor.repeat_interleave(3,1)

    augmented_rgb = refocus_aug(rgb_tensor, depth_tensor)

    print(f'Writing output {save_path} ...')
    trans_topil(augmented_rgb[0]).save(save_path)


input_path = Path(args.input_path)

if input_path.is_dir():
    for f in glob.glob(args.input_path+'/*'):
        save_outputs(f, os.path.splitext(os.path.basename(f))[0])
else:
    print("invalid file path!")
    sys.exit()