Derived from [DanbooRegion project](https://github.com/lllyasviel/DanbooRegion). Create a flat color image from an original image and a skeleton map image. Set `DANBOO_REGION_DIR` variable to the directory of a cloned DanbooRegion project.

flatten.py

#    Copyright 2020 lllyasviel
#    Copyright 2021 kosuke1701
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at

#        http://www.apache.org/licenses/LICENSE-2.0

#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.
#
#    Note:
#    This file is derived from https://raw.githubusercontent.com/lllyasviel/DanbooRegion/cabdb23255955ca73c81376fecaba65698d551e0/code/segment.py
#    Function get_fill and up_fill are copied from the original file.
#    Function flatten is derived from function segment of the original file.

import os
import sys

import cv2
import numpy as np
from PIL import Image

DANBOO_REGION_DIR = os.environ["DANBOO_REGION_DIR"]
sys.path.append(f"{DANBOO_REGION_DIR}/code")
from tricks import *

def pil2cv(img):
    img = np.array(img)
    return cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

def get_fill(image):
    labeled_array, num_features = label(image / 255)
    filled_area = find_all(labeled_array)
    return filled_area


def up_fill(fills, cur_fill_map):
    new_fillmap = cur_fill_map.copy()
    padded_fillmap = np.pad(cur_fill_map, [[1, 1], [1, 1]], 'constant', constant_values=0)
    max_id = np.max(cur_fill_map)
    for item in fills:
        points0 = padded_fillmap[(item[0] + 1, item[1] + 0)]
        points1 = padded_fillmap[(item[0] + 1, item[1] + 2)]
        points2 = padded_fillmap[(item[0] + 0, item[1] + 1)]
        points3 = padded_fillmap[(item[0] + 2, item[1] + 1)]
        all_points = np.concatenate([points0, points1, points2, points3], axis=0)
        pointsets, pointcounts = np.unique(all_points[all_points > 0], return_counts=True)
        if len(pointsets) == 1 and item[0].shape[0] < 128:
            new_fillmap[item] = pointsets[0]
        else:
            max_id += 1
            new_fillmap[item] = max_id
    return new_fillmap

def flatten(skeleton_map, orig_image):
    orig_size = orig_image.size
    height = pil2cv(skeleton_map).astype(np.float32)
    orig_image = pil2cv(orig_image).astype(np.float32)
    height = np.mean(height, axis=2)

    height += (height - cv2.GaussianBlur(height, (0, 0), 3.0)) * 10.0
    height = height.clip(0, 255).astype(np.uint8)
    marker = height.copy()
    marker[marker > 135] = 255
    marker[marker < 255] = 0
    fills = get_fill(marker / 255)
    for fill in fills:
        if fill[0].shape[0] < 64:
            marker[fill] = 0
    filter = np.array([
        [0, 1, 0],
        [1, 1, 1],
        [0, 1, 0]],
        dtype=np.uint8)
    big_marker = cv2.erode(marker, filter, iterations=5)
    fills = get_fill(big_marker / 255)
    for fill in fills:
        if fill[0].shape[0] < 64:
            big_marker[fill] = 0
    big_marker = cv2.dilate(big_marker, filter, iterations=5)
    small_marker = marker.copy()
    small_marker[big_marker > 127] = 0
    fin_labels, nil = label(big_marker / 255)
    fin_labels = up_fill(get_fill(small_marker), fin_labels)
    water = cv2.watershed(orig_image.clip(0, 255).astype(np.uint8), fin_labels.astype(np.int32)) + 1
    water = thinning(water)
    all_region_indices = find_all(water)
    regions = np.zeros_like(orig_image, dtype=np.uint8)
    region_label = np.zeros((orig_image.shape[0], orig_image.shape[1]), dtype=np.uint32)
    i_label = 0
    for region_indices in all_region_indices:
        regions[region_indices] = np.random.randint(low=0, high=255, size=(3,)).clip(0, 255).astype(np.uint8)
        region_label[region_indices] = i_label
        i_label += 1
    result = np.zeros_like(orig_image, dtype=np.uint8)
    for region_indices in all_region_indices:
        result[region_indices] = np.median(orig_image[region_indices], axis=0)
    result = result.clip(0, 255)

    return Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)).resize(orig_size), region_label

if __name__=="__main__":
    img, label = flatten(Image.open("20210421_skeleton.png"), Image.open("20210421.png"))
    img.save("20210421_flatten.png")
    print(label)