def cluster_by_masked_area(root_dir, args):
logger.info(f'Clustering {root_dir}')
clustered_dir = root_dir + '_clustered'
make_dirs(clustered_dir)
radius = 5
# all masks with ratio in x +- radius will be stored in sub-directory x
clustered_centors = np.arange(radius, 100, radius * 2)
clustered_subdirs = []
for c in clustered_centors:
# make sub-directories for each ratio range
clustered_subdirs.append(make_dir_under_root(clustered_dir, str(c)))
for i, filename in enumerate(get_everything_under(root_dir)):
if i % 100 == 0:
logger.info(f'clustering {filename}')
if args.image_masks:
ratio = get_masked_ratio(Image.open(filename))
else:
# filename is a diretory containing multiple mask files
ratio = np.mean([
get_masked_ratio(Image.open(f))
for f in get_everything_under(filename, pattern='*.png')
])
# find the nearest centor
for i, c in enumerate(clustered_centors):
if c - radius <= ratio * 100 <= c + radius:
shutil.move(filename, clustered_subdirs[i])
break
shutil.rmtree(root_dir)
os.rename(clustered_dir, root_dir)
def copy_masks_without_boarder(root_dir, args):
def erase_mask_boarder(mask, gap):
pix = np.asarray(mask).astype('uint8') * 255
pix[:gap, :] = 255
pix[-gap:, :] = 255
pix[:, :gap] = 255
pix[:, -gap:] = 255
return Image.fromarray(pix).convert('1')
wo_boarder_dir = root_dir + '_noBoarder'
logger.info('Copying all masks')
shutil.copytree(root_dir, wo_boarder_dir)
for i, filename in enumerate(get_everything_under(wo_boarder_dir)):
if i % 100 == 0:
logger.info(f'Erasing {filename}\'s boarder')
if args.image_masks:
mask = Image.open(filename)
mask_wo_boarder = erase_mask_boarder(mask,
args.leave_boarder_unmasked)
mask_wo_boarder.save(filename)
else:
# filename is a diretory containing multiple mask files
for f in get_everything_under(filename, pattern='*.png'):
mask = Image.open(f)
mask_wo_boarder = erase_mask_boarder(
mask, args.leave_boarder_unmasked)
mask_wo_boarder.save(f)
return wo_boarder_dir
def main(args, vid_lens):
preset = get_stroke_preset(args.stroke_preset)
make_dirs(args.output_dir)
if args.redo_without_generation:
assert (len(get_everything_under(args.output_dir)) > 0)
# put back clustered masks
for clustered_subdir in get_everything_under(args.output_dir):
if not os.path.isdir(clustered_subdir):
continue
for f in get_everything_under(clustered_subdir):
shutil.move(f, args.output_dir)
os.rmdir(clustered_subdir)
else:
if args.image_masks:
for i in range(args.n):
if i % 100 == 0:
logger.info(f'Generating mask number {i:07d}')
nStroke = decide_nStroke(args)
mask = get_video_masks_by_moving_random_stroke(
video_len=1,
imageWidth=args.image_width,
imageHeight=args.image_height,
nStroke=nStroke,
**preset)[0]
mask.save(os.path.join(args.output_dir, f'{i:07d}.png'))
else:
for i in range(args.n):
mask_dir = make_dir_under_root(args.output_dir, f'{i:05d}')
mask_reader = MaskReader(mask_dir, read=False)
nStroke = decide_nStroke(args)
masks = get_video_masks_by_moving_random_stroke(
imageWidth=args.image_width,
imageHeight=args.image_height,
video_len=vid_lens[i],
nStroke=nStroke,
**preset)
mask_reader.set_files(masks)
mask_reader.save_files(output_dir=mask_reader.dir_name)
if args.leave_boarder_unmasked is not None:
logger.info(
f'Create a copy of all output and erase the copies\' s boarder '
f'by {args.leave_boarder_unmasked} pixels')
dir_leave_boarder = copy_masks_without_boarder(args.output_dir, args)
if args.cluster_by_area:
cluster_by_masked_area(dir_leave_boarder, args)
if args.cluster_by_area:
cluster_by_masked_area(args.output_dir, args)
def evaluate_all_videos(
root_gt_dir,
root_result_dir,
flownet_checkpoint_path,
root_mask_dir,
test_num,
result_postfix="",
):
total_warping_error = 0
total_error = 0
total_psnr = 0
total_ssim = 0
total_p_dist = 0
total_length = 0
result_dirs = get_everything_under(root_result_dir,
only_dirs=True)[:test_num]
gt_dirs = get_everything_under(root_gt_dir, only_dirs=True)[:test_num]
mask_dirs = get_everything_under(root_mask_dir, only_dirs=True)[:test_num]
for i, (result_dir, gt_dir) in enumerate(zip(result_dirs, gt_dirs)):
result_dir = os.path.join(result_dir, result_postfix)
mask_dir = mask_dirs[i]
logger.info(f"Processing {result_dir}, mask {mask_dir}")
warping_error, error, psnr_value, ssim_value, p_dist, length = \
evaluate_video(result_dir, gt_dir, mask_dir, flownet_checkpoint_path)
total_warping_error += warping_error
total_error += error
total_ssim += ssim_value
total_psnr += psnr_value
total_p_dist += p_dist
total_length += length
avg_warping_error = total_warping_error / len(gt_dirs)
avg_mse_error = total_error / total_length
avg_ssim = total_ssim / total_length
avg_psnr = total_psnr / total_length
avg_p_dist = (total_p_dist / total_length)[0]
logger.info(f"Total avg warping error {avg_warping_error:.8f}")
logger.info(f"Total avg error {avg_mse_error:.5f}")
logger.info(f"Total avg ssim {avg_ssim:.4f}")
logger.info(f"Total avg pSNR {avg_psnr:.2f}")
logger.info(f"Total avg Perceptual distance {avg_p_dist:.4f}")
logger.info(f"Total length {total_length}")
logger.info(f"Video num: {len(result_dirs)}")
return (avg_warping_error, avg_mse_error, avg_ssim, avg_psnr, avg_p_dist)
def __init__(
self,
rids: RootInputDirectories,
rods: RootOutputDirectories,
args: dict,
):
super().__init__(rids, rods, args)
self.image_dir = rids.root_videos_dir
self.mask_dir = rids.root_masks_dir
self.max_num = args.get('max_num', None)
self.image_filenames = get_everything_under(
self.image_dir, pattern="*/*.jpg")[:self.max_num]
self.mask_filenames = get_everything_under(
self.mask_dir, pattern="*/*.png")[:self.max_num]
if len(self.image_filenames) > len(self.mask_filenames):
logger.warning(
f"image num {len(self.image_filenames)} > mask num {len(self.mask_filenames)}"
)
self.data_len = len(self.image_filenames)
self.mask_len = len(self.mask_filenames)
def evaluate_fid_score(root_gt_dir, root_result_dir, result_postfix):
to_tensors = transforms.Compose([
Stack(),
ToTorchFormatTensor(),
])
result_dirs = get_everything_under(root_result_dir, only_dirs=True)
gt_dirs = get_everything_under(root_gt_dir, only_dirs=True)
# worker = torchvision.transforms.Resize(SIZE, Image.BILINEAR)
output_i3d_activations = []
real_i3d_activations = []
with torch.no_grad():
for i, (result_dir, gt_dir) in enumerate(zip(result_dirs, gt_dirs)):
if i % 20 == 0:
logger.info(f"Getting {i}th i3d activations")
result_dir = os.path.join(result_dir, result_postfix)
result_frame_reader = FrameReader(result_dir, evaluate_image=True).files
gt_frame_reader = FrameReader(gt_dir, evaluate_image=True).files
# result_frame_reader = [worker(img) for img in result_frame_reader]
# gt_frame_reader = [worker(img) for img in gt_frame_reader]
# Unsqueeze batch dimension
outputs = to_tensors(result_frame_reader).unsqueeze(0).to(torch.device('cuda:0'))
targets = to_tensors(gt_frame_reader).unsqueeze(0).to(torch.device('cuda:0'))
# get i3d activation
output_i3d_activations.append(get_i3d_activations(outputs).cpu().numpy())
real_i3d_activations.append(get_i3d_activations(targets).cpu().numpy())
# concat and evaluate fid score
output_i3d_activations = np.concatenate(output_i3d_activations, axis=0)
real_i3d_activations = np.concatenate(real_i3d_activations, axis=0)
fid_score = get_fid_score(real_i3d_activations, output_i3d_activations)
logger.info(f"Video num: {len(result_dirs)}")
logger.info(f"FID score: {fid_score}")
return fid_score
def gen_masks(in_direc, out_direc, args):
assert in_direc != out_direc, 'input and output directory should not be the same'
video_dirs = get_everything_under(in_direc)
assert all([os.path.isdir(d) for d in video_dirs])
for i, video_dir in enumerate(video_dirs):
video_name = os.path.basename(video_dir)
length = len(glob(os.path.join(video_dir, '*')))
assert length > 0, f'Got video in "{video_dir}" with no frames'
if i % 50 == 0:
logger.info(
f'generating masks of video #{i}, name: {video_name}, len: {length}'
)
mask_dir = os.path.join(out_direc, video_name)
if not os.path.exists(mask_dir):
os.makedirs(mask_dir, mode=0o755)
mask = get_random_mask(img_size=args.img_size)
mask_len = length if args.n_masks is None else args.n_masks
for j in range(mask_len):
mask_name = f'{j:05d}.png'
mask.save(os.path.join(mask_dir, mask_name))
