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)