diff options
Diffstat (limited to 'rpc/img_ops.py')
| -rw-r--r-- | rpc/img_ops.py | 151 |
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 |