import cv2 from .thinning import * from .trappedball_fill import * from skimage.measure import block_reduce from skimage.morphology import disk, dilation, erosion from numba import njit def np_min_pool(x): return block_reduce(x, (2, 2), np.min) def np_max_pool(x): return block_reduce(x, (2, 2), np.max) def np_max_441(x): return block_reduce(x, (4, 4, 1), np.max) def np_max_pool_221(x): return block_reduce(x, (2, 2, 1), np.max) def np_max_pool_s(x, s): return block_reduce(x, (s, s, 1), np.max) def binarize(x): xp = x.copy() xp[xp < 250] = 0 xp[xp > 0] = 255 return xp def get_initial_fillmap(boundary, merge=True): fillmap = build_fill_map(boundary, flood_fill_multi(boundary, merge=merge)) return fillmap def up_propagate(small_fillmap, big_boundary): new_fillmap = cv2.resize(small_fillmap, (big_boundary.shape[1], big_boundary.shape[0]), interpolation=cv2.INTER_NEAREST) padded_fillmap = np.pad(new_fillmap, [[1, 1], [1, 1]], 'constant', constant_values=0) new_mask = np.ones_like(new_fillmap, dtype=np.uint8) * 255 new_mask[new_fillmap > 0] = 0 new_mask[big_boundary < 240] = 0 fills = flood_fill_multi(new_mask, merge=True) max_id = np.max(new_fillmap) 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) > 0: new_fillmap[item] = pointsets[np.argmax(pointcounts)] else: max_id += 1 new_fillmap[item] = max_id return new_fillmap def laplas_fill(b_512, b_256, b_128): b_512 = binarize(b_512) b_256 = binarize(b_256) b_128 = binarize(b_128) f128 = get_initial_fillmap(b_128) f256 = up_propagate(f128, b_256) f512 = up_propagate(f256, b_512) fin = thinning(f512) return fin @ njit def get_corner(x): corner = x.copy() s0 = corner.shape[0] s1 = corner.shape[1] for i0 in range(1, s0 - 1): for i1 in range(1, s1 - 1): if x[i0, i1] == 0: continue if x[i0, i1 - 1] == 0: if x[i0 - 1, i1 - 1] == 0: continue if x[i0 + 1, i1 - 1] == 0: continue corner[i0, i1] = 0 continue if x[i0, i1 + 1] == 0: if x[i0 - 1, i1 + 1] == 0: continue if x[i0 + 1, i1 + 1] == 0: continue corner[i0, i1] = 0 continue if x[i0 - 1, i1] == 0: if x[i0 - 1, i1 - 1] == 0: continue if x[i0 - 1, i1 + 1] == 0: continue corner[i0, i1] = 0 continue if x[i0 + 1, i1] == 0: if x[i0 + 1, i1 - 1] == 0: continue if x[i0 + 1, i1 + 1] == 0: continue corner[i0, i1] = 0 continue return corner def monogrouh(x): y = 255 - x y = dilation(y, disk(1)) y = dilation(y, disk(1)) y = erosion(y, disk(1)) y = erosion(y, disk(1)) y = 255 - y return y def corners(x): y = x.copy() y = monogrouh(y) y = get_corner(y) y = monogrouh(y) y = get_corner(y) y = monogrouh(y) return y def save_fill(name, fill): cv2.imwrite(name, show_fill_map(fill)) def double_fill(b_1024, b_512, b256): b256 = binarize(b256) b_512 = binarize(b_512) b_1024 = binarize(b_1024) b_1024 = corners(b_1024) b_512 = np.min(np.stack([b_512, np_min_pool(b_1024)], axis=2), axis=2) b_512 = corners(b_512) b_256 = np.min(np.stack([b256, np_min_pool(b_512)], axis=2), axis=2) b_256 = corners(b_256) b_128 = np_min_pool(b_256) b_128 = corners(b_128) b_64 = np_min_pool(b_128) f64 = get_initial_fillmap(b_64) print('get_initial_fillmap(b_64)') f128 = up_propagate(f64, b_128) print('up_propagate(f64, b_128)') f256 = up_propagate(f128, b_256) print('up_propagate(f128, b_256)') f512 = up_propagate(f256, b_512) print('up_propagate(f256, b_512)') f1024 = up_propagate(f512, b_1024) print('up_propagate(f512, b_1024)') fin = thinning(f1024) print('thinning(f1024)') # cv2.imwrite('b_64.png', b_64) # cv2.imwrite('b_128.png', b_128) # cv2.imwrite('b_256.png', b_256) # cv2.imwrite('b_512.png', b_512) # cv2.imwrite('b_1024.png', b_1024) # save_fill('f64.png', f64) # save_fill('f128.png', f128) # save_fill('f256.png', f256) # save_fill('f512.png', f512) # save_fill('f1024.png', f1024) # save_fill('fin.png', fin) return find_all(fin) def single_fill(b_2048, path): b_2048 = corners(binarize(b_2048)) f2048 = get_initial_fillmap(b_2048, merge=False) print(path + 'get_initial_fillmap(b_2048, merge=False)') fin = thinning(f2048) print(path + 'thinning(f2048)') # cv2.imwrite(path + 'b_2048.png', b_2048) # save_fill(path + 'f2048.png', f2048) # save_fill(path + 'fin.png', fin) return find_all(fin) def deatlize(x): x = cv2.GaussianBlur(x, (0, 0), 0.8) x = cv2.medianBlur(x, 3) return x def low_down(gradient_mask): return 1.0 - cv2.dilate(255 - gradient_mask, np.ones((3, 3), np.uint8), iterations=2).astype(np.float32) / 255.0 def cv2pyrDown(x): return cv2.pyrDown(cv2.medianBlur(cv2.medianBlur(x, 3), 3)) def cv2pyrUp(x): return cv2.pyrUp(cv2.medianBlur(cv2.medianBlur(x, 3), 3)) def re_deatlize(visulized, s1024): gradient_mask_1024 = binarize(s1024) gradient_mask_512 = np_min_pool(gradient_mask_1024) gradient_mask_256 = np_min_pool(gradient_mask_512) gradient_mask_128 = np_min_pool(gradient_mask_256) gradient_mask_64 = np_min_pool(gradient_mask_128) gradient_mask_1024 = low_down(gradient_mask_1024) gradient_mask_512 = low_down(gradient_mask_512) gradient_mask_256 = low_down(gradient_mask_256) gradient_mask_128 = low_down(gradient_mask_128) gradient_mask_64 = low_down(gradient_mask_64) sample_1024 = visulized.astype(np.float32) sample_512 = cv2pyrDown(sample_1024) sample_256 = cv2pyrDown(sample_512) sample_128 = cv2pyrDown(sample_256) sample_64 = cv2pyrDown(sample_128) sample_32 = cv2pyrDown(sample_64) gradient_1024 = sample_1024 - cv2pyrUp(sample_512) gradient_512 = sample_512 - cv2pyrUp(sample_256) gradient_256 = sample_256 - cv2pyrUp(sample_128) gradient_128 = sample_128 - cv2pyrUp(sample_64) gradient_64 = sample_64 - cv2pyrUp(sample_32) rec_32 = sample_32 rec_64 = cv2pyrUp(rec_32) + gradient_64 * (1 - gradient_mask_64[:, :, None]) rec_128 = cv2pyrUp(rec_64) + gradient_128 * (1 - gradient_mask_128[:, :, None]) rec_256 = cv2pyrUp(rec_128) + gradient_256 * (1 - gradient_mask_256[:, :, None]) rec_512 = cv2pyrUp(rec_256) + gradient_512 * (1 - gradient_mask_512[:, :, None]) rec_1024 = cv2pyrUp(rec_512) + gradient_1024 * (1 - gradient_mask_1024[:, :, None]) return rec_1024.clip(0, 255).astype(np.uint8) def tiny_deatlize(visulized, s2048): gradient_mask_2048 = s2048.copy() gradient_mask_1024 = np_min_pool(gradient_mask_2048) gradient_mask_512 = np_min_pool(gradient_mask_1024) gradient_mask_256 = np_min_pool(gradient_mask_512) gradient_mask_2048 = low_down(gradient_mask_2048) gradient_mask_1024 = low_down(gradient_mask_1024) gradient_mask_512 = low_down(gradient_mask_512) gradient_mask_256 = low_down(gradient_mask_256) sample_2048 = visulized.astype(np.float32) sample_1024 = cv2.pyrDown(sample_2048) sample_512 = cv2.pyrDown(sample_1024) sample_256 = cv2.pyrDown(sample_512) sample_128 = cv2.pyrDown(sample_256) gradient_2048 = sample_2048 - cv2.pyrUp(sample_1024) gradient_1024 = sample_1024 - cv2.pyrUp(sample_512) gradient_512 = sample_512 - cv2.pyrUp(sample_256) gradient_256 = sample_256 - cv2.pyrUp(sample_128) rec_128 = sample_128 rec_256 = cv2.pyrUp(rec_128) + gradient_256 * (1 - gradient_mask_256[:, :, None]) rec_512 = cv2.pyrUp(rec_256) + gradient_512 * (1 - gradient_mask_512[:, :, None]) rec_1024 = cv2.pyrUp(rec_512) + gradient_1024 * (1 - gradient_mask_1024[:, :, None]) rec_2048 = cv2.pyrUp(rec_1024) + gradient_2048 * (1 - gradient_mask_2048[:, :, None]) return rec_2048.clip(0, 255).astype(np.uint8) def adain(x, y): x_high = cv2.GaussianBlur(x, (0, 0), 3.0) y_high = cv2.GaussianBlur(y, (0, 0), 3.0) return (x.astype(np.float32) - x_high.astype(np.float32) + y_high.astype(np.float32)).clip(0, 255).astype(np.uint8) def corrupt(x, b128): float_sketch = x.astype(float) float_base = cv2.resize(float_sketch, (b128.shape[1], b128.shape[0]), cv2.INTER_AREA) alpha = b128[:, :, 0] / 255.0 float_base = alpha * float_base + (1 - alpha) * np.mean(float_base) float_base = cv2.GaussianBlur(float_base, (0, 0), 8.0) float_base = cv2.resize(float_base, (x.shape[1], x.shape[0]), cv2.INTER_CUBIC) result = float_sketch / (float_base + 1e-10) result = result.clip(0, 1) result -= np.min(result) result /= np.max(result) return (result * 255.0).clip(0, 255).astype(np.uint8) def fuse_sketch(color, sketch, fills, fixer, points_arr, colors_arr): sketch = cv2.resize(sketch, (color.shape[1], color.shape[0])) fills = cv2.resize(fills, (color.shape[1], color.shape[0]), interpolation=cv2.INTER_NEAREST) fill_id = np.unique(fills.flatten()) bg = np.zeros_like(color, dtype=np.uint8) checking_result = np.zeros(dtype=np.int32, shape=(np.max(fills) + 1,)) - 1 length_points = int(len(points_arr)) for _ in range(length_points): checking_result[fills[points_arr[_][0], points_arr[_][1]]] = _ for id in fill_id: points = np.where(fills == id) if len(points[0]) > 0: color_id = checking_result[id] if color_id > -1: bg[points] = np.array(colors_arr[color_id]) else: bg[points] = np.median(color[points], axis=0) fixed = adain(fixer(sketch, bg), bg) result = (fixed.astype(np.float32) + sketch[:, :, None].astype(np.float32) - 255.0).clip(0, 255).astype(np.uint8) return result, fixed, bg def balance_fill(color, fills, points, sizer): color = cv2.resize(color, (sizer.shape[1], sizer.shape[0]), interpolation=cv2.INTER_NEAREST) points = cv2.resize(points, (sizer.shape[1], sizer.shape[0]), interpolation=cv2.INTER_NEAREST) bg = np.zeros_like(color, dtype=np.uint8) for region in fills: if len(region[0]) > 0: region_points = points[region] region_points = region_points[region_points[:, 3] > 0] if region_points.shape[0] > 0: points_color, points_color_count = np.unique(region_points, return_counts=True, axis=0) bg[region] = points_color[np.argmax(points_color_count)][0:3] else: bg[region] = np.median(color[region], axis=0) return bg def shade_fill(color, fills, points, sizer): color = cv2.resize(color, (sizer.shape[1], sizer.shape[0]), interpolation=cv2.INTER_NEAREST) points = cv2.resize(points, (sizer.shape[1], sizer.shape[0]), interpolation=cv2.INTER_NEAREST) bg = np.zeros_like(color, dtype=np.uint8) for region in fills: if len(region[0]) > 0: region_points = points[region] region_points = region_points[region_points[:, 3] > 0] if region_points.shape[0] > 0: points_color, points_color_count = np.unique(region_points, return_counts=True, axis=0) c = points_color[np.argmax(points_color_count)][0:3] r = c[0] g = c[1] b = c[2] if r == 1 and g == 233 and b == 0: bg[region] = 255 elif r == 0 and g == 233 and b == 1: bg[region] = 0 else: bg[region] = np.median(color[region], axis=0) else: bg[region] = np.median(color[region], axis=0) return bg def get_alpha_piece(points): padded_points = np.pad(points, [[1, 1], [1, 1], [0, 0]], 'constant', constant_values=127) lines = 255 - padded_points[:, :, 3] lines[lines < 240] = 0 fills = flood_fill_multi(lines, merge=True) result = np.zeros_like(padded_points) for item in fills: points0 = padded_points[(item[0], item[1] + 1)] points1 = padded_points[(item[0], item[1] - 1)] points2 = padded_points[(item[0] + 1, item[1])] points3 = padded_points[(item[0] - 1, item[1])] all_points = np.concatenate([points0, points1, points2, points3], axis=0) all_points = all_points[all_points[:, 3] > 0] all_points = np.unique(all_points, axis=0) if all_points.shape[0] == 1: result[item] = all_points[0] piece = result[1:-1, 1:-1, :] piece = np.maximum(piece, points) return piece, points def fin_deatlize(color, sketch): cf = color.astype(np.float32) alpha = sketch.astype(np.float32)[:, :, None] / 255.0 plain = cf * alpha lines = cf * (1 - alpha) plain = cv2.medianBlur(plain, 5) plain = cv2.medianBlur(plain, 3) fin = plain + lines return fin.clip(0, 255).astype(np.uint8)