img ops!!

1 files changed, 151 insertions, 0 deletions

diff --git a/rpc/img_ops.py b/rpc/img_ops.py
new file mode 100644
index 0000000..07e0bf1
--- /dev/null
+++ b/rpc/img_ops.py
@@ -0,0 +1,151 @@
+import numpy as np
+import cv2
+import math
+
+def clamp(n,a,b):
+  return max(a, min(n, b))
+
+def lerp(n,a,b):
+  return (b-a)*n+a
+
+def process_image(opt, data_opt, im):
+  img = im[:, :, ::-1].copy()
+  processed = False
+
+  if data_opt.process_frac == 0:
+    return img
+
+  if data_opt.clahe is True:
+    processed = True
+    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
+    l, a, b = cv2.split(lab)
+    clahe = cv2.createCLAHE(clipLimit=data_opt.clip_limit, tileGridSize=(8,8))
+    l = clahe.apply(l)
+    limg = cv2.merge((l,a,b))
+    img = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
+  if data_opt.posterize is True:
+    processed = True
+    img = cv2.pyrMeanShiftFiltering(img, data_opt.spatial_window, data_opt.color_window)
+  if data_opt.grayscale is True:
+    processed = True
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+  if data_opt.blur is True:
+    processed = True
+    img = cv2.GaussianBlur(img, (data_opt.blur_radius, data_opt.blur_radius), data_opt.blur_sigma)
+  if data_opt.canny is True:
+    processed = True
+    img = cv2.Canny(img, data_opt.canny_lo, data_opt.canny_hi)
+    img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+
+  if processed is False:
+    return img
+
+  src_img = im[:, :, ::-1].copy()
+  frac_a = data_opt.process_frac
+  frac_b = 1.0 - frac_a
+  array_a = np.multiply(src_img.astype('float64'), frac_a)
+  array_b = np.multiply(img.astype('float64'), frac_b)
+  img = np.add(array_a, array_b).astype('uint8')
+  return img
+
+last_im = None
+def mix_next_image(opt, data_opt, im, i, sequence, sequence_i):
+  global last_im
+
+  if (i % 100) == 0:
+    print('%04d: process image...' % (i))
+
+  last_path = opt.render_dir + "frame_{:05d}.png".format(i)
+  tmp_path = opt.render_dir + "frame_{:05d}_tmp.png".format(i+1)
+  next_path = opt.render_dir + "frame_{:05d}.png".format(i+1)
+  current_path = opt.render_dir + "ren_{:05d}.png".format(i+1)
+  meta = { 'i': i, 'sequence_i': sequence_i, 'sequence_len': len(sequence) }
+  if data_opt.sequence and len(sequence):
+    sequence_path = sequence[sequence_i]
+    if sequence_i >= len(sequence)-1:
+      print('(((( sequence looped ))))')
+      sequence_i = 1
+    else:
+      sequence_i += 1
+
+  if data_opt.store_a is not True:
+    os.remove(last_path)
+
+  if data_opt.store_b is True:
+    image_pil = Image.fromarray(im, mode='RGB')
+    image_pil.save(tmp_path)
+    os.rename(tmp_path, current_path)
+
+  if data_opt.recursive and last_im is not None:
+    if data_opt.sequence and len(sequence):
+      A_img = Image.open(sequence_path).convert('RGB')
+      A_im = np.asarray(A_img)
+      frac_a = data_opt.recursive_frac
+      frac_b = data_opt.sequence_frac
+      frac_sum = frac_a + frac_b
+      if frac_sum > 1.0:
+        frac_a = frac_a / frac_sum
+        frac_b = frac_b / frac_sum
+      if data_opt.transition:
+        t = lerp(math.sin(i / data_opt.transition_period * math.pi * 2.0 ) / 2.0 + 0.5, data_opt.transition_min, data_opt.transition_max)
+        frac_a *= 1.0 - t
+        frac_b *= 1.0 - t
+      frac_c = 1.0 - frac_a - frac_b
+      array_a = np.multiply(last_im.astype('float64'), frac_a)
+      array_b = np.multiply(A_im.astype('float64'), frac_b)
+      array_c = np.multiply(im.astype('float64'), frac_c)
+      array_ab = np.add(array_a, array_b)
+      array_abc = np.add(array_ab, array_c)
+      next_im = array_abc.astype('uint8')
+
+    else:
+      frac_a = data_opt.recursive_frac
+      frac_b = 1.0 - frac_a
+      array_a = np.multiply(last_im.astype('float64'), frac_a)
+      array_b = np.multiply(im.astype('float64'), frac_b)
+      next_im = np.add(array_a, array_b).astype('uint8')
+
+    if data_opt.recurse_roll != 0:
+      last_im = np.roll(im, data_opt.recurse_roll, axis=data_opt.recurse_roll_axis)
+    else:
+      last_im = next_im.copy().astype('uint8')
+
+  elif data_opt.sequence and len(sequence):
+    A_img = Image.open(sequence_path).convert('RGB')
+    A_im = np.asarray(A_img)
+    frac_b = data_opt.sequence_frac
+    if data_opt.transition:
+      t = lerp(math.sin(i / data_opt.transition_period * math.pi * 2.0 ) / 2.0 + 0.5, data_opt.transition_min, data_opt.transition_max)
+      frac_b *= 1.0 - t
+    frac_c = 1.0 - frac_b
+    array_b = np.multiply(A_im.astype('float64'), frac_b)
+    array_c = np.multiply(im.astype('float64'), frac_c)
+    array_bc = np.add(array_b, array_c)
+    next_im = array_bc.astype('uint8')
+
+  else:
+    last_im = im.copy().astype('uint8')
+    next_im = im
+
+  next_img = process_image(opt, data_opt, next_im)
+
+  if data_opt.send_image == 'sequence':
+    rpc_client.send_pil_image("frame_{:05d}.png".format(i+1), meta, A_img, data_opt.output_format)
+  if data_opt.send_image == 'recursive':
+    pil_im = Image.fromarray(next_im)
+    rpc_client.send_pil_image("frame_{:05d}.png".format(i+1), meta, pil_im, data_opt.output_format)
+  if data_opt.send_image == 'a':
+    rgb_im = cv2.cvtColor(next_img, cv2.COLOR_BGR2RGB)
+    pil_im = Image.fromarray(rgb_im)
+    rpc_client.send_pil_image("frame_{:05d}.png".format(i+1), meta, pil_im, data_opt.output_format)
+  if data_opt.send_image == 'b':
+    image_pil = Image.fromarray(im, mode='RGB')
+    rpc_client.send_pil_image("frame_{:05d}.png".format(i+1), meta, image_pil, data_opt.output_format)
+
+  cv2.imwrite(tmp_path, next_img)
+  os.rename(tmp_path, next_path)
+
+  if (i % 20) == 0:
+    print("created {}".format(next_path))
+
+  return sequence_i
\ No newline at end of file