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)