decompositioner.py

import os
import numpy as np
from scipy.spatial import ConvexHull
from sklearn.cluster import MiniBatchKMeans
from tricks import *
import cv2


ksd = 8
mbc = MiniBatchKMeans(ksd)


def get_theme(img):
 images = np.reshape(cv2.resize(img, (256, 256)), (256 * 256, 3))
 hull = ConvexHull(images)
 return hull.points[hull.vertices]


def simplify_points(points, img):
 labels = mbc.fit(points)
 new_points = []
 all_center = np.mean(labels.cluster_centers_, axis=0)
 distances = np.sum((points - all_center) ** 2, axis=1) ** 0.5

 for idx in range(ksd):
 candidates = points[labels.labels_ == idx]
 scores = distances[labels.labels_ == idx]
 best_id = np.argmax(scores)
 new_points.append(candidates[best_id])

 new_points.sort(key=np.sum, reverse=True)

 new_points = np.stack(new_points, axis=0)
 return new_points.clip(0, 255).astype(np.uint8)


def get_ini_layers(miku, points):
 results = []
 final_target = miku.astype(np.float32)
 bg = np.zeros_like(final_target, dtype=np.float32) + points[0]
 results.append(np.concatenate([bg, np.zeros_like(bg, dtype=np.float32) + 255], axis=2)[:, :, 0:4])
 current_result = bg.copy()
 for layer_index in range(1, ksd):
 current_base = current_result.astype(np.float32)
 current_color = np.zeros_like(final_target, dtype=np.float32) + points[layer_index]
 overall_direction = final_target - current_base
 avaliable_direction = current_color - current_base
 current_alpha = np.sum(overall_direction * avaliable_direction, axis=2, keepdims=True) / np.sum(
 avaliable_direction * avaliable_direction, axis=2, keepdims=True)
 current_alpha = current_alpha.clip(0, 1)
 current_result = (current_color * current_alpha + current_base * (1 - current_alpha)).clip(0, 255)
 results.append(np.concatenate([current_color, current_alpha * 255.0], axis=2))
 return results


def make_reconstruction(layers):
 bg = np.zeros_like(layers[0], dtype=np.float32)[:, :, 0:3] + 255
 for item in layers:
 current_alpha = item[:, :, 3:4] / 255.0
 bg = item[:, :, 0:3] * current_alpha + bg * (1 - current_alpha)
 return bg


def improve_layers(layers, miku):
 reconstruction = make_reconstruction(layers)
 b = miku - reconstruction
 new_layers = []
 for item in layers:
 new_item = item.copy()
 new_item[:, :, 0:3] = (new_item[:, :, 0:3] + b).clip(0, 255)
 new_layers.append(new_item)
 return new_layers


def cluster_all(labeled_array, num_features):
 xs = [[] for _ in range(num_features)]
 ys = [[] for _ in range(num_features)]
 M = labeled_array.shape[0]
 N = labeled_array.shape[1]
 for x in range(M):
 for y in range(N):
 i = labeled_array[x, y]
 xs[i].append(x)
 ys[i].append(y)
 result = []
 for _ in range(num_features):
 result.append((np.array(xs[_]), np.array(ys[_])))
 return result


def meder(x):
 y = x.copy()
 y = cv2.medianBlur(y, 5)
 y = cv2.medianBlur(y, 5)
 y = cv2.medianBlur(y, 3)
 y = cv2.medianBlur(y, 3)
 return y


def re_med(s_2048):

 sample_2048 = s_2048.astype(np.float32)
 sample_1024 = cv2.pyrDown(sample_2048)
 sample_512 = cv2.pyrDown(sample_1024)
 sample_256 = cv2.pyrDown(sample_512)

 gradient_2048 = sample_2048 - cv2.pyrUp(sample_1024)
 gradient_1024 = sample_1024 - cv2.pyrUp(sample_512)
 gradient_512 = sample_512 - cv2.pyrUp(sample_256)

 rec_256 = meder(sample_256)
 rec_512 = cv2.pyrUp(rec_256) + meder(gradient_512)
 rec_1024 = cv2.pyrUp(rec_512) + meder(gradient_1024)
 rec_2048 = cv2.pyrUp(rec_1024) + meder(gradient_2048)
 return rec_2048


def process_ctx(sketch, solid, render):
 solid = solid.astype(np.float32)
 sketch = d_resize(cv2.cvtColor(sketch, cv2.COLOR_GRAY2RGB), solid.shape).astype(np.float32)
 render = d_resize(render, solid.shape).astype(np.float32)
 alpha = sketch / 255.0
 all_diff = render - solid
 all_lines = render.copy()
 all_lines = cv2.erode(all_lines, np.ones((3,3), np.uint8)) * 0.618
 all_diff = re_med(all_diff)
 all_lines = re_med(all_lines)
 recon = solid + all_diff
 recon = recon * alpha + all_lines * (1 - alpha)
 recon2 = (solid + all_diff) * alpha + re_med(solid) * (1 - alpha)
 recon3 = reason_blending(recon2, sketch)
 return recon.clip(0, 255).astype(np.uint8), recon2.clip(0, 255).astype(np.uint8), recon3.clip(0, 255).astype(np.uint8)


def process_psd(sketch, solid, render, path='./'):
 recon = process_ctx(sketch, solid, render)
 points = get_theme(solid)
 points = simplify_points(points, solid)
 compositions = get_ini_layers(solid, points)
 compositions = improve_layers(compositions, solid)
 for _ in range(ksd):
 cv2.imwrite(path + str(_ + 1) + '.color.png', compositions[_].clip(0, 255).astype(np.uint8))
 solid = make_reconstruction(compositions).clip(0, 255).astype(np.uint8)
 os.makedirs(path, exist_ok=True)
 alpha = 1 - sketch.astype(np.float32) / 255.0
 now = solid
 now = (now.astype(np.float32) + sketch.astype(np.float32) - 255.0).clip(0, 255)
 sketch = 255 + now - solid
 cv2.imwrite(path + '9.sketch.png', sketch.clip(0, 255).astype(np.uint8))
 all_diff = recon.astype(np.float32) - now
 all_light = all_diff.copy()
 all_shadow = - all_diff.copy()
 all_light[all_light < 0] = 0
 all_shadow[all_shadow < 0] = 0
 sketch_color = all_light * alpha
 light = all_light * (1 - alpha)
 all_shadow = 255 - all_shadow
 cv2.imwrite(path + '10.sketch_color.png', sketch_color.clip(0, 255).astype(np.uint8))
 cv2.imwrite(path + '11.light.png', light.clip(0, 255).astype(np.uint8))
 cv2.imwrite(path + '12.shadow.png', all_shadow.clip(0, 255).astype(np.uint8))
 return recon


def process_albedo(albedo, composition, sketch):
 DEL = albedo.astype(np.float32)
 HSV = cv2.cvtColor(albedo, cv2.COLOR_RGB2HSV).astype(np.float32)
 YUV = cv2.cvtColor(albedo, cv2.COLOR_RGB2YUV).astype(np.float32)
 solid = composition.astype(np.float32)
 light = sketch[:, :, None].astype(np.float32)

 DEL = DEL * light / 255.0 + solid * (1 - light / 255.0)
 HSV[:, :, 2:3] = np.minimum(HSV[:, :, 2:3], light)
 YUV[:, :, 0:1] = np.minimum(YUV[:, :, 0:1], light)

 DEL = DEL.clip(0, 255).astype(np.uint8)
 HSV = HSV.clip(0, 255).astype(np.uint8)
 YUV = YUV.clip(0, 255).astype(np.uint8)

 return cv2.cvtColor(HSV, cv2.COLOR_HSV2RGB), cv2.cvtColor(YUV, cv2.COLOR_YUV2RGB), DEL


def process_overlay(composition, sketch):
 RGB = composition.astype(np.float32)
 alpha = sketch[:, :, None].astype(np.float32) / 255.0
 return (RGB * alpha).clip(0, 255).astype(np.uint8)