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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
|
import os
import sys
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 = "woodflow/"
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/")
hsv = []
for i,fn in enumerate(sorted(os.listdir(work_dir))):
# load image and convert to grayscale
pil_image = Image.open(work_dir + fn).convert('RGB')
im = np.array(pil_image)
im = im[:, :, ::-1].copy()
im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
# store first frame
if i == 0:
prev = im
hsv = np.zeros((256,256,3))
hsv[...,1] = 255
continue
# compute optical flow
flow = cv2.calcOpticalFlowFarneback(prev, im, None, 0.5, 3, 15, 3, 5, 1.2, 0)
# turn into magnitude/angle
mag, ang = cv2.cartToPolar(flow[...,0], flow[...,1])
# store angle as hue
hsv[...,0] = ang * 180 / np.pi / 2
# store magnitude as lum
hsv[...,2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
# convert this HSL to BGR
bgr = cv2.cvtColor(hsv.astype('uint8'), cv2.COLOR_HSV2BGR)
ren = "frame_{:05d}.png".format(i)
if (i % 10) == 3:
wd = "test/"
elif (i % 10) == 6:
wd = "val/"
else:
wd = "train/"
if i == 1:
prev_hsv = np.copy(hsv)
prev_bgr = np.copy(bgr)
continue
cv2.imwrite(render_dir + "A/" + wd + ren, prev_bgr)
cv2.imwrite(render_dir + "B/" + wd + ren, bgr)
# copyfile(work_dir + fn, render_dir + "B/" + wd + ren)
prev = im
prev_hsv = np.copy(hsv)
prev_bgr = np.copy(bgr)
|
