diff options
Diffstat (limited to 'scripts/builders/canny-dataset.pl')
| -rw-r--r-- | scripts/builders/canny-dataset.pl | 39 |
1 files changed, 39 insertions, 0 deletions
diff --git a/scripts/builders/canny-dataset.pl b/scripts/builders/canny-dataset.pl new file mode 100644 index 0000000..6d1085d --- /dev/null +++ b/scripts/builders/canny-dataset.pl @@ -0,0 +1,39 @@ +import os +from shutil import move, copyfile +from PIL import Image, ImageOps +from shutil import copyfile, rmtree +import numpy as np +import cv2 + +work_dir = "woodsized/" +render_dir = "woodcanny/" + +if os.path.exists(render_dir): + rmtree(render_dir) + +os.makedirs(render_dir) +os.makedirs(render_dir + "A/") +os.makedirs(render_dir + "A/train/") +os.makedirs(render_dir + "A/test/") +os.makedirs(render_dir + "A/val/") +os.makedirs(render_dir + "B/") +os.makedirs(render_dir + "B/train/") +os.makedirs(render_dir + "B/test/") +os.makedirs(render_dir + "B/val/") + +for i,fn in enumerate(sorted(os.listdir(work_dir))): + pil_image = Image.open(work_dir + fn).convert('RGB') + opencv_image = np.array(pil_image) + opencv_image = opencv_image[:, :, ::-1].copy() + opencv_image = cv2.GaussianBlur(opencv_image, (3,3), 1) + opencv_image = cv2.Canny(opencv_image, 100, 200) + ren = "frame_{:05d}.png".format(i) + if (i % 10) == 3: + wd = "test/" + elif (i % 10) == 6: + wd = "val/" + else: + wd = "train/" + cv2.imwrite(render_dir + "A/" + wd + ren, opencv_image) + copyfile(work_dir + fn, render_dir + "B/" + wd + ren) + |