mogrify script, flow scripts, dataset options

1 files changed, 81 insertions, 0 deletions

diff --git a/scripts/flow/flow-dir.py b/scripts/flow/flow-dir.py
new file mode 100644
index 0000000..60c1dd9
--- /dev/null
+++ b/scripts/flow/flow-dir.py
@@ -0,0 +1,81 @@
+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 = "IMG_1734/"
+render_dir = "IMG_1734/diff/"
+
+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/")
+os.makedirs(render_dir + "out/")
+
+hsv = []
+prev = None
+prev_bgr = None
+
+for i,fn in enumerate(sorted(os.listdir(work_dir))):
+  if os.path.basename(fn).startswith('.') or not os.path.isfile(work_dir + fn):
+    continue
+
+  # 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 prev is None:
+    prev = im
+    hsv = np.zeros((256,512,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_L2)
+
+  # convert this HSL to BGR
+  bgr = cv2.cvtColor(hsv.astype('uint8'), cv2.COLOR_HSV2BGR)
+
+  if prev_bgr is None:
+    prev_bgr = np.copy(bgr)
+    continue
+
+  if (i % 10) == 3:
+    wd = "test/"
+  elif (i % 10) == 6:
+    wd = "val/"
+  else:
+    wd = "train/"
+
+  ren = "frame_{:05d}.png".format(i)
+
+  #cv2.imwrite(render_dir + "A/" + wd + ren, prev_bgr)
+  #cv2.imwrite(render_dir + "B/" + wd + ren, bgr)
+  cv2.imwrite(render_dir + "out/" + ren, bgr)
+  # copyfile(work_dir + fn, render_dir + "B/" + wd + ren)
+  prev = im
+  prev_bgr = np.copy(bgr)
+