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)
|