Spaces:
Running
on
Zero
Running
on
Zero
File size: 8,320 Bytes
05fb4ab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
import glob
import os.path
from datasets.listdataset import ListDataset, Aug_ListDataset, Aug_Doc3d_ListDataset
from datasets.util import split2list
def make_doc3d_dataset_list(dir, split=None):
dataset_list = []
for sample_name in os.listdir(dir):
# flow_map = os.path.join(dir, os.path.basename(flow_map))
# root_filename = os.path.basename(flow_map)[:-9]
img = os.path.join(dir, sample_name, 'img.png') # source image
flow_map = os.path.join(dir, sample_name, 'bm.mat') # flow_map
abd = os.path.join(dir, sample_name, 'recon.png') # recon
# img2 = os.path.join(dir, root_filename + '_img_2.jpg') # target image
dataset_list.append([img, flow_map, abd])
return split2list(dataset_list, split, default_split=0.97)
def make_doc_dataset_list(dir, split=None):
dataset_list = []
for sample_name in os.listdir(dir):
# flow_map = os.path.join(dir, os.path.basename(flow_map))
# root_filename = os.path.basename(flow_map)[:-9]
img = os.path.join(dir, sample_name, 'warped_document.png') # source image
flow_map = os.path.join(dir, sample_name, 'warped_BM.npz') # flow_map
abd = os.path.join(dir, sample_name, 'warped_recon.png') # recon
# img2 = os.path.join(dir, root_filename + '_img_2.jpg') # target image
dataset_list.append([img, flow_map, abd])
return split2list(dataset_list, split, default_split=0.97)
def make_dataset(dir, get_mapping, split=None, dataset_name=None):
"""
Will search for triplets that go by the pattern '[name]_img1.ppm [name]_img2.ppm in folder images and
[name]_flow.flo' in folder flow """
images = []
if get_mapping: # false
flow_dir = 'mapping'
# flow_dir is actually mapping dir in that case, it is always normalised to [-1,1]
else:
flow_dir = 'flow'
image_dir = 'images'
# Make sure that the folders exist
if not os.path.isdir(dir):
raise ValueError("the training directory path that you indicated does not exist ! ")
if not os.path.isdir(os.path.join(dir, flow_dir)):
raise ValueError("the training directory path that you indicated does not contain the flow folder ! "
"Check your directories.")
if not os.path.isdir(os.path.join(dir, image_dir)):
raise ValueError("the training directory path that you indicated does not contain the images folder ! "
"Check your directories.")
for flow_map in sorted(glob.glob(os.path.join(dir, flow_dir, '*_flow.flo'))):
flow_map = os.path.join(flow_dir, os.path.basename(flow_map))
root_filename = os.path.basename(flow_map)[:-9]
img1 = os.path.join(image_dir, root_filename + '_img_1.jpg') # source image
img2 = os.path.join(image_dir, root_filename + '_img_2.jpg') # target image
if not (os.path.isfile(os.path.join(dir,img1)) or os.path.isfile(os.path.join(dir,img2))):
continue
if dataset_name is not None: # false
images.append([[os.path.join(dataset_name, img1),
os.path.join(dataset_name, img2)],
os.path.join(dataset_name, flow_map)])
else: # true
images.append([[img1, img2], flow_map])
return split2list(images, split, default_split=0.97)
def assign_default(default_dict, dict):
if dict is None:
dall = default_dict
else:
dall = {}
dall.update(default_dict)
dall.update(dict)
return dall
def Doc_Dataset(root, source_image_transform=None, target_image_transform=None, flow_transform=None,
co_transform=None, split=None, get_mapping=False, compute_mask_zero_borders=False,
add_discontinuity=False, min_nbr_perturbations=5, max_nbr_perturbations=6,
parameters_v2=None):
"""
Builds a dataset from existing image pairs and corresponding ground-truth flow fields and optionally add
some flow perturbations.
Args:
root: path to root folder
source_image_transform: image transformations to apply to source images
target_image_transform: image transformations to apply to target images
flow_transform: flow transformations to apply to ground-truth flow fields
co_transform: transformations to apply to both images and ground-truth flow fields
split: split (float) between training and testing, 0 means all pairs are in test_dataset
get_mapping: output mapping instead of flow in __getittem__ ?
compute_mask_zero_borders: output mask of zero borders ?
add_discontinuity: add discontinuity to image pairs and corresponding ground-truth flow field ?
min_nbr_perturbations:
max_nbr_perturbations:
parameters_v2: parameters of v2
Returns:
train_dataset
test_dataset
"""
train_list, test_list = make_doc_dataset_list(root, split) # get list [[sample1],[sample2],...]
train_dataset = ListDataset(root, train_list, source_image_transform=source_image_transform,
target_image_transform=target_image_transform,
flow_transform=flow_transform, co_transform=co_transform, get_mapping=get_mapping,
compute_mask_zero_borders=compute_mask_zero_borders)
test_dataset = ListDataset(root, test_list, source_image_transform=source_image_transform,
target_image_transform=target_image_transform,
flow_transform=flow_transform, co_transform=co_transform, get_mapping=get_mapping,
compute_mask_zero_borders=compute_mask_zero_borders)
return train_dataset, test_dataset
def Aug_Doc_Dataset(root, source_image_transform=None, target_image_transform=None, flow_transform=None,
co_transform=None, split=None, get_mapping=False, compute_mask_zero_borders=False,
add_discontinuity=False, min_nbr_perturbations=5, max_nbr_perturbations=6,
parameters_v2=None):
train_list, test_list = make_doc_dataset_list(root, split) # get list [[sample1],[sample2],...]
train_dataset = Aug_ListDataset(root, train_list, source_image_transform=source_image_transform,
target_image_transform=target_image_transform,
flow_transform=flow_transform, co_transform=co_transform, get_mapping=get_mapping,
compute_mask_zero_borders=compute_mask_zero_borders)
test_dataset = Aug_ListDataset(root, test_list, source_image_transform=source_image_transform,
target_image_transform=target_image_transform,
flow_transform=flow_transform, co_transform=co_transform, get_mapping=get_mapping,
compute_mask_zero_borders=compute_mask_zero_borders)
return train_dataset, test_dataset
def Doc3d_Dataset(root, source_image_transform=None, target_image_transform=None, flow_transform=None,
co_transform=None, split=None, get_mapping=False, compute_mask_zero_borders=False,
add_discontinuity=False, min_nbr_perturbations=5, max_nbr_perturbations=6,
parameters_v2=None):
train_list, test_list = make_doc3d_dataset_list(root, split) # get list [[sample1],[sample2],...]
train_dataset = Aug_Doc3d_ListDataset(root, train_list, source_image_transform=source_image_transform,
target_image_transform=target_image_transform,
flow_transform=flow_transform, co_transform=co_transform, get_mapping=get_mapping,
compute_mask_zero_borders=compute_mask_zero_borders)
test_dataset = Aug_Doc3d_ListDataset(root, test_list, source_image_transform=source_image_transform,
target_image_transform=target_image_transform,
flow_transform=flow_transform, co_transform=co_transform, get_mapping=get_mapping,
compute_mask_zero_borders=compute_mask_zero_borders)
return train_dataset, test_dataset
|