def dataflow(centroids, num_reference=3, num_process=16, shuffle=True):
ds = Kinetics('/data/public/rw/datasets/videos/kinetics',
num_frames=num_reference + 1,
skips=[0, 4, 4, 8][:num_reference + 1],
shuffle=shuffle)
ds = df.MapDataComponent(ds, ImageProcess.resize(small_axis=256), index=1)
ds = df.MapDataComponent(ds, ImageProcess.crop(shape=(256, 256)), index=1)
#ds = df.MapDataComponent(ds, lambda images: [cv2.resize(image, (256, 256)) for image in images], index=1)
ds = df.MapData(
ds, lambda dp: [
dp[1][:num_reference],
copy.deepcopy(dp[1][:num_reference]), dp[1][num_reference:],
copy.deepcopy(dp[1][num_reference:])
])
# for images (ref, target)
for idx in [0, 2]:
ds = df.MapDataComponent(
ds,
lambda images: [
cv2.cvtColor(image, cv2.COLOR_BGR2GRAY).reshape(256, 256, 1)
for image in images
],
index=idx)
# for labels (ref, target)
for idx in [1, 3]:
ds = df.MapDataComponent(
ds,
lambda images: [cv2.resize(image, (32, 32)) for image in images],
index=idx)
ds = df.MapDataComponent(ds,
lambda images: [
cv2.cvtColor(np.float32(image / 255.0),
cv2.COLOR_BGR2Lab)[:, :, 1:]
for image in images
],
index=idx)
ds = df.MapDataComponent(
ds,
lambda images: [
np.array([
np.argmin(np.linalg.norm(centroids - v, axis=1))
for v in image.reshape((-1, 2))
]).reshape((32, 32, 1)) for image in images
],
index=idx)
# stack for tensor
ds = df.MapData(
ds, lambda dp:
[np.stack(dp[0] + dp[2], axis=0),
np.stack(dp[1] + dp[3], axis=0)])
ds = df.MapData(ds, tuple) # for tensorflow.data.dataset
ds = df.MultiProcessPrefetchData(ds, nr_prefetch=256, nr_proc=num_process)
ds = df.PrefetchDataZMQ(ds, nr_proc=1)
return ds
def test_resizw_with_crop_center():
images = [np.zeros((480, 640, 3)), np.zeros((640, 480, 3))]
outputs = ImageProcess.resize(small_axis=256, copy=True)(images)
outputs = ImageProcess.crop((256, 256), copy=True)(outputs)
assert images[0].shape == (480, 640, 3)
assert images[1].shape == (640, 480, 3)
assert outputs[0].shape == (256, 256, 3)
assert outputs[1].shape == (256, 256, 3)
def dataflow(name='davis', scale=1, split='val'):
if name == 'davis':
ds = Davis('/data/zubin/videos/davis', name=split, num_frames=1, shuffle=False)
elif name == 'kinetics':
ds = Kinetics('/data/public/rw/datasets/videos/kinetics', num_frames=1, skips=[0], shuffle=False)
else:
raise Exception('not support dataset %s' % name)
if name != 'davis':
ds = df.MapData(ds, lambda dp: [dp[0], dp[1], dp[1]])
ds = df.MapData(ds, lambda dp: [
dp[0], # index
dp[1], # original
dp[2], # mask
dp[3], # name
])
feature_size = int(256 * scale)
size = (feature_size, feature_size)
ds = df.MapDataComponent(ds, ImageProcess.resize(small_axis=feature_size), index=1)
ds = df.MapDataComponent(ds, lambda images: cv2.resize(images[0], size), index=2)
ds = df.MapData(ds, lambda dp: [
dp[0], # index
dp[1][0], # original small axis 256 x scale
cv2.cvtColor(cv2.resize(dp[1][0], size), cv2.COLOR_BGR2GRAY).reshape((size[0], size[1], 1)), # gray (256xscale)x(256xscale)x1
dp[2], # annotated mask 256xscale x 256xscale
dp[3], # name
])
ds = df.MultiProcessPrefetchData(ds, nr_prefetch=32, nr_proc=1)
return ds
def test_resize_small_axis():
images = [np.zeros((480, 640, 3)), np.zeros((640, 480, 3))]
outputs = ImageProcess.resize(small_axis=256, copy=True)(images)
assert images[0].shape == (480, 640, 3)
assert images[1].shape == (640, 480, 3)
assert outputs[0].shape == (256, 341, 3)
assert outputs[1].shape == (341, 256, 3)
def test_crop_center():
images = [np.zeros((480, 640, 3)), np.zeros((640, 480, 3))]
outputs = ImageProcess.crop((200, 300), copy=True)(images)
assert images[0].shape == (480, 640, 3)
assert images[1].shape == (640, 480, 3)
assert outputs[0].shape == (200, 300, 3)
assert outputs[1].shape == (200, 300, 3)
def dataflow(name='davis', scale=1):
if name == 'davis':
ds = Davis('/data/public/rw/datasets/videos/davis/trainval',
num_frames=1,
shuffle=False)
elif name == 'kinetics':
ds = Kinetics('/data/public/rw/datasets/videos/kinetics',
num_frames=1,
skips=[0],
shuffle=False)
else:
raise Exception('not support dataset %s' % name)
if name != 'davis':
ds = df.MapData(ds, lambda dp: [dp[0], dp[1], dp[1]])
ds = df.MapData(
ds,
lambda dp: [
dp[0], # index
dp[1], # original
dp[2], # mask
])
size = (256 * scale, 256 * scale)
ds = df.MapDataComponent(ds,
ImageProcess.resize(small_axis=256 * scale),
index=1)
ds = df.MapDataComponent(ds,
lambda images: cv2.resize(images[0], size),
index=2)
ds = df.MapData(
ds,
lambda dp: [
dp[0], # index
dp[1][0], # original
cv2.cvtColor(cv2.resize(dp[1][0], size), cv2.COLOR_BGR2GRAY).
reshape((size[0], size[1], 1)), # gray
dp[2], # mask
])
ds = df.MultiProcessPrefetchData(ds, nr_prefetch=32, nr_proc=1)
return ds