wehub-resource-sync
377 lines
14 KiB
Python
377 lines
14 KiB
Python
from __future__ import print_function, division
|
|
|
|
import time
|
|
import multiprocessing
|
|
|
|
import numpy as np
|
|
from skimage import data
|
|
|
|
import imgaug as ia
|
|
import imgaug.multicore as multicore
|
|
from imgaug import augmenters as iaa
|
|
|
|
|
|
class PoolWithMarkedWorker(multicore.Pool):
|
|
def __init__(self, *args, **kwargs):
|
|
super(PoolWithMarkedWorker, self).__init__(*args, **kwargs)
|
|
|
|
@classmethod
|
|
def _worker(cls, batch_idx, batch):
|
|
process_name = multiprocessing.current_process().name
|
|
# print("[_worker] called %s. images in batch: %d" % (process_name, len(batch.images_unaug),))
|
|
if "-1" in process_name:
|
|
for image in batch.images_unaug:
|
|
image[::4, ::4, :] = [255, 255, 255]
|
|
return multicore.Pool._worker(batch_idx, batch)
|
|
|
|
|
|
def main():
|
|
augseq = iaa.Sequential([
|
|
iaa.Fliplr(0.5),
|
|
iaa.CoarseDropout(p=0.1, size_percent=0.1)
|
|
])
|
|
|
|
def func_images(images, random_state, parents, hooks):
|
|
time.sleep(0.2)
|
|
return images
|
|
|
|
def func_heatmaps(heatmaps, random_state, parents, hooks):
|
|
return heatmaps
|
|
|
|
def func_keypoints(keypoints_on_images, random_state, parents, hooks):
|
|
return keypoints_on_images
|
|
|
|
augseq_slow = iaa.Sequential([
|
|
iaa.Fliplr(0.5),
|
|
iaa.Lambda(
|
|
func_images=func_images,
|
|
func_heatmaps=func_heatmaps,
|
|
func_keypoints=func_keypoints
|
|
)
|
|
])
|
|
|
|
print("------------------")
|
|
print(".pool()")
|
|
print("------------------")
|
|
with augseq.pool() as pool:
|
|
time_start = time.time()
|
|
batches = list(load_images())
|
|
batches_aug = pool.map_batches(batches)
|
|
images_aug = []
|
|
keypoints_aug = []
|
|
for batch_aug in batches_aug:
|
|
images_aug.append(batch_aug.images_aug)
|
|
keypoints_aug.append(batch_aug.keypoints_aug)
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
# ia.imshow(draw_grid(images_aug, keypoints_aug))
|
|
|
|
print("------------------")
|
|
print("Pool.map_batches(batches)")
|
|
print("------------------")
|
|
with multicore.Pool(augseq) as pool:
|
|
time_start = time.time()
|
|
batches = list(load_images())
|
|
batches_aug = pool.map_batches(batches)
|
|
images_aug = []
|
|
keypoints_aug = []
|
|
for batch_aug in batches_aug:
|
|
images_aug.append(batch_aug.images_aug)
|
|
keypoints_aug.append(batch_aug.keypoints_aug)
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
# ia.imshow(draw_grid(images_aug, keypoints_aug))
|
|
|
|
print("------------------")
|
|
print("Pool.imap_batches(batches)")
|
|
print("------------------")
|
|
with multicore.Pool(augseq) as pool:
|
|
time_start = time.time()
|
|
batches_aug = pool.imap_batches(load_images())
|
|
images_aug = []
|
|
keypoints_aug = []
|
|
for batch in batches_aug:
|
|
images_aug.append(batch.images_aug)
|
|
keypoints_aug.append(batch.keypoints_aug)
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
# ia.imshow(draw_grid(images_aug, keypoints_aug))
|
|
|
|
print("------------------")
|
|
print("Pool.imap_batches(batches, chunksize=32)")
|
|
print("------------------")
|
|
with multicore.Pool(augseq) as pool:
|
|
time_start = time.time()
|
|
batches_aug = pool.imap_batches(load_images(n_batches=1000), chunksize=32)
|
|
count = 0
|
|
for batch in batches_aug:
|
|
count += 1
|
|
assert count == 1000
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
|
|
print("------------------")
|
|
print("Pool.imap_batches(batches, chunksize=2)")
|
|
print("------------------")
|
|
with multicore.Pool(augseq) as pool:
|
|
time_start = time.time()
|
|
batches_aug = pool.imap_batches(load_images(n_batches=1000), chunksize=2)
|
|
count = 0
|
|
for batch in batches_aug:
|
|
count += 1
|
|
assert count == 1000
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
|
|
print("------------------")
|
|
print("Pool.imap_batches(batches, chunksize=1)")
|
|
print("------------------")
|
|
with multicore.Pool(augseq) as pool:
|
|
time_start = time.time()
|
|
batches_aug = pool.imap_batches(load_images(n_batches=1000), chunksize=1)
|
|
count = 0
|
|
for batch in batches_aug:
|
|
count += 1
|
|
assert count == 1000
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
|
|
print("------------------")
|
|
print("Pool.map_batches(batches, chunksize=32)")
|
|
print("------------------")
|
|
with multicore.Pool(augseq) as pool:
|
|
time_start = time.time()
|
|
batches_aug = pool.map_batches(list(load_images(n_batches=1000)), chunksize=32)
|
|
assert len(batches_aug) == 1000
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
|
|
print("------------------")
|
|
print("Pool.map_batches chunksize with fast aug")
|
|
print("------------------")
|
|
def test_fast(processes, chunksize):
|
|
augseq = iaa.Dropout(0.1)
|
|
with multicore.Pool(augseq, processes=processes) as pool:
|
|
batches = list(load_images(n_batches=10000, draw_text=False))
|
|
time_start = time.time()
|
|
batches_aug = pool.map_batches(batches, chunksize=chunksize)
|
|
assert len(batches_aug) == 10000
|
|
print("chunksize=%d, worker=%s, time=%.4fs" % (chunksize, processes, time.time() - time_start))
|
|
|
|
test_fast(-4, 1)
|
|
test_fast(1, 1)
|
|
test_fast(None, 1)
|
|
test_fast(1, 4)
|
|
test_fast(None, 4)
|
|
test_fast(1, 32)
|
|
test_fast(None, 32)
|
|
|
|
print("------------------")
|
|
print("Pool.imap_batches chunksize with fast aug")
|
|
print("------------------")
|
|
def test_fast_imap(processes, chunksize):
|
|
augseq = iaa.Dropout(0.1)
|
|
with multicore.Pool(augseq, processes=processes) as pool:
|
|
time_start = time.time()
|
|
batches_aug = pool.imap_batches(load_images(n_batches=10000, draw_text=False), chunksize=chunksize)
|
|
batches_aug = list(batches_aug)
|
|
assert len(batches_aug) == 10000
|
|
print("chunksize=%d, worker=%s, time=%.4fs" % (chunksize, processes, time.time() - time_start))
|
|
|
|
test_fast_imap(-4, 1)
|
|
test_fast_imap(1, 1)
|
|
test_fast_imap(None, 1)
|
|
test_fast_imap(1, 4)
|
|
test_fast_imap(None, 4)
|
|
test_fast_imap(1, 32)
|
|
test_fast_imap(None, 32)
|
|
|
|
print("------------------")
|
|
print("Pool.map_batches with computationally expensive aug")
|
|
print("------------------")
|
|
def test_heavy(processes, chunksize):
|
|
augseq_heavy = iaa.PiecewiseAffine(scale=0.2, nb_cols=8, nb_rows=8)
|
|
with multicore.Pool(augseq_heavy, processes=processes) as pool:
|
|
batches = list(load_images(n_batches=500, draw_text=False))
|
|
time_start = time.time()
|
|
batches_aug = pool.map_batches(batches, chunksize=chunksize)
|
|
assert len(batches_aug) == 500
|
|
print("chunksize=%d, worker=%s, time=%.4fs" % (chunksize, processes, time.time() - time_start))
|
|
|
|
test_heavy(-4, 1)
|
|
test_heavy(1, 1)
|
|
test_heavy(None, 1)
|
|
test_heavy(1, 4)
|
|
test_heavy(None, 4)
|
|
test_heavy(1, 32)
|
|
test_heavy(None, 32)
|
|
|
|
print("------------------")
|
|
print("Pool.imap_batches(batches), slow loading")
|
|
print("------------------")
|
|
with multicore.Pool(augseq) as pool:
|
|
time_start = time.time()
|
|
batches_aug = pool.imap_batches(load_images(n_batches=100, sleep=0.2))
|
|
images_aug = []
|
|
keypoints_aug = []
|
|
for batch in batches_aug:
|
|
images_aug.append(batch.images_aug)
|
|
keypoints_aug.append(batch.keypoints_aug)
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
|
|
print("------------------")
|
|
print("Pool.imap_batches(batches), maxtasksperchild=4")
|
|
print("------------------")
|
|
with multicore.Pool(augseq, maxtasksperchild=4) as pool:
|
|
time_start = time.time()
|
|
batches_aug = pool.imap_batches(load_images(n_batches=100))
|
|
images_aug = []
|
|
keypoints_aug = []
|
|
for batch in batches_aug:
|
|
images_aug.append(batch.images_aug)
|
|
keypoints_aug.append(batch.keypoints_aug)
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
ia.imshow(draw_grid(images_aug, keypoints_aug))
|
|
|
|
print("------------------")
|
|
print("Pool.imap_batches(batches), seed=1")
|
|
print("------------------")
|
|
# we color here the images of the first worker to see in the grids which images belong to one worker
|
|
with PoolWithMarkedWorker(augseq, seed=1) as pool:
|
|
time_start = time.time()
|
|
batches_aug = pool.imap_batches(load_images(n_batches=4))
|
|
images_aug = []
|
|
keypoints_aug = []
|
|
for batch in batches_aug:
|
|
images_aug.append(batch.images_aug)
|
|
keypoints_aug.append(batch.keypoints_aug)
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
grid_a = draw_grid(images_aug, keypoints_aug)
|
|
|
|
with multicore.Pool(augseq, seed=1) as pool:
|
|
time_start = time.time()
|
|
batches_aug = pool.imap_batches(load_images(n_batches=4))
|
|
images_aug = []
|
|
keypoints_aug = []
|
|
for batch in batches_aug:
|
|
images_aug.append(batch.images_aug)
|
|
keypoints_aug.append(batch.keypoints_aug)
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
grid_b = draw_grid(images_aug, keypoints_aug)
|
|
|
|
grid_b[:, 0:2, 0] = 0
|
|
grid_b[:, 0:2, 1] = 255
|
|
grid_b[:, 0:2, 2] = 0
|
|
ia.imshow(np.hstack([grid_a, grid_b]))
|
|
|
|
print("------------------")
|
|
print("Pool.imap_batches(batches), seed=None")
|
|
print("------------------")
|
|
with multicore.Pool(augseq, seed=None) as pool:
|
|
time_start = time.time()
|
|
batches_aug = pool.imap_batches(load_images(n_batches=4))
|
|
images_aug = []
|
|
keypoints_aug = []
|
|
for batch in batches_aug:
|
|
images_aug.append(batch.images_aug)
|
|
keypoints_aug.append(batch.keypoints_aug)
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
grid_a = draw_grid(images_aug, keypoints_aug)
|
|
|
|
with multicore.Pool(augseq, seed=None) as pool:
|
|
time_start = time.time()
|
|
batches_aug = pool.imap_batches(load_images(n_batches=4))
|
|
images_aug = []
|
|
keypoints_aug = []
|
|
for batch in batches_aug:
|
|
images_aug.append(batch.images_aug)
|
|
keypoints_aug.append(batch.keypoints_aug)
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
grid_b = draw_grid(images_aug, keypoints_aug)
|
|
|
|
ia.imshow(np.hstack([grid_a, grid_b]))
|
|
|
|
print("------------------")
|
|
print("Pool.imap_batches(batches), maxtasksperchild=4, seed=1")
|
|
print("------------------")
|
|
with multicore.Pool(augseq, maxtasksperchild=4, seed=1) as pool:
|
|
time_start = time.time()
|
|
batches_aug = pool.imap_batches(load_images(n_batches=100))
|
|
images_aug = []
|
|
keypoints_aug = []
|
|
for batch in batches_aug:
|
|
images_aug.append(batch.images_aug)
|
|
keypoints_aug.append(batch.keypoints_aug)
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
ia.imshow(draw_grid(images_aug, keypoints_aug))
|
|
|
|
for augseq_i in [augseq, augseq_slow]:
|
|
print("------------------")
|
|
print("Many very small runs (batches=1)")
|
|
print("------------------")
|
|
with multicore.Pool(augseq_i) as pool:
|
|
time_start = time.time()
|
|
for i in range(100):
|
|
_ = pool.map_batches(list(load_images(n_batches=1)))
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
|
|
print("------------------")
|
|
print("Many very small runs (batches=2)")
|
|
print("------------------")
|
|
with multicore.Pool(augseq_i) as pool:
|
|
time_start = time.time()
|
|
for i in range(100):
|
|
_ = pool.map_batches(list(load_images(n_batches=2)))
|
|
print("Done in %.4fs" % (time.time() - time_start,))
|
|
|
|
|
|
def load_images(n_batches=10, sleep=0.0, draw_text=True):
|
|
batch_size = 4
|
|
astronaut = data.astronaut()
|
|
astronaut = ia.imresize_single_image(astronaut, (64, 64))
|
|
kps = ia.KeypointsOnImage([ia.Keypoint(x=15, y=25)], shape=astronaut.shape)
|
|
|
|
counter = 0
|
|
for i in range(n_batches):
|
|
if draw_text:
|
|
batch_images = []
|
|
batch_kps = []
|
|
for b in range(batch_size):
|
|
astronaut_text = ia.draw_text(astronaut, x=0, y=0, text="%d" % (counter,), color=[0, 255, 0], size=16)
|
|
batch_images.append(astronaut_text)
|
|
batch_kps.append(kps)
|
|
counter += 1
|
|
batch = ia.Batch(
|
|
images=np.array(batch_images, dtype=np.uint8),
|
|
keypoints=batch_kps
|
|
)
|
|
else:
|
|
if i == 0:
|
|
batch_images = np.array([np.copy(astronaut) for _ in range(batch_size)], dtype=np.uint8)
|
|
|
|
batch = ia.Batch(
|
|
images=np.copy(batch_images),
|
|
keypoints=[kps.deepcopy() for _ in range(batch_size)]
|
|
)
|
|
yield batch
|
|
if sleep > 0:
|
|
time.sleep(sleep)
|
|
|
|
|
|
def draw_grid(images_aug, keypoints_aug):
|
|
if keypoints_aug is None:
|
|
keypoints_aug = []
|
|
for bidx in range(len(images_aug)):
|
|
keypoints_aug.append([None for image in images_aug[bidx]])
|
|
|
|
images_kps_batches = []
|
|
for bidx in range(len(images_aug)):
|
|
images_kps_batch = []
|
|
for image, kps in zip(images_aug[bidx], keypoints_aug[bidx]):
|
|
if kps is None:
|
|
image_kps = image
|
|
else:
|
|
image_kps = kps.draw_on_image(image, size=5, color=[255, 0, 0])
|
|
images_kps_batch.append(image_kps)
|
|
images_kps_batches.extend(images_kps_batch)
|
|
|
|
grid = ia.draw_grid(images_kps_batches, cols=len(images_aug[0]))
|
|
return grid
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|