def get_dataset_vgg_with_masks(annot_path, img_dir, batch_size, strict=False, keypoints_num = 22, connections_num = 19):
def gen(df):
def f():
for i in df:
yield tuple(i)
return f
df, size = get_dataflow_vgg(
annot_path=annot_path,
img_dir=img_dir,
strict=strict,
x_size=368,
y_size=46,
include_outputs_masks=True
)
df.reset_state()
ds = tf.data.Dataset.from_generator(
gen(df), (tf.float32, tf.float32, tf.float32, tf.float32, tf.float32),
output_shapes=(
tf.TensorShape([368, 368, 3]),
tf.TensorShape([46, 46, connections_num*2]),
tf.TensorShape([46, 46, keypoints_num]),
tf.TensorShape([46, 46, connections_num*2]),
tf.TensorShape([46, 46, keypoints_num])
)
)
ds = ds.map(lambda x0, x1, x2, x3, x4: ((x0, x1, x2), (x3, x4, x3, x4, x3, x4, x3, x4, x3, x4, x3, x4)))
ds = ds.batch(batch_size)
return ds, size
def get_dataset_vgg(annot_path, img_dir, batch_size, strict=False):
def gen(df):
def f():
for i in df:
yield tuple(i)
return f
df, size = get_dataflow_vgg(annot_path=annot_path,
img_dir=img_dir,
strict=strict,
x_size=368,
y_size=46,
include_outputs_masks=False)
df.reset_state()
ds = tf.data.Dataset.from_generator(
gen(df), (tf.float32, tf.float32, tf.float32),
output_shapes=(tf.TensorShape([368, 368,
3]), tf.TensorShape([46, 46, 38]),
tf.TensorShape([46, 46, 19])))
ds = ds.map(lambda x0, x1, x2:
(x0, (x1, x2, x1, x2, x1, x2, x1, x2, x1, x2, x1, x2)))
ds = ds.batch(batch_size)
return ds, size
def get_dataset_with_masks(annot_path, img_dir, batch_size, strict=False, x_size=368, y_size=46):
# def get_dataset_with_masks(annot_path, img_dir, batch_size, strict=False, x_size=48, y_size=64):
def gen(df):
def f():
for i in df:
yield tuple(i)
return f
df, size = get_dataflow_vgg(
annot_path=annot_path,
img_dir=img_dir,
strict=strict,
x_size=x_size,
y_size=y_size,
include_outputs_masks=True
)
df.reset_state()
ds = tf.data.Dataset.from_generator(
gen(df), (tf.float32, tf.float32, tf.float32, tf.float32, tf.float32),
output_shapes=(
tf.TensorShape([x_size, x_size, 3]),
tf.TensorShape([y_size, y_size, 12]),#TODO
tf.TensorShape([y_size, y_size, 6]),
tf.TensorShape([y_size, y_size, 12]),
tf.TensorShape([y_size, y_size, 6]),#TODO
)
)
ds = ds.map(lambda x0, x1, x2, x3, x4: ((x0, x1, x2), (x3, x4)))
ds = ds.batch(batch_size)
return ds, size