summaryrefslogtreecommitdiff
path: root/scripts/builders/canny-dir.py
blob: ae1bd62546b13eed3181ab3e5dcf0af7827c5321 (plain)
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
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)