Duplicate from hylee/White-box-Cartoonization

Co-authored-by: hylee <[email protected]>

.gitattributes

@@ -0,0 +1,29 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text
+*.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,15 @@
+---
+python_version: 3.7
+title: White Box Cartoonization
+emoji: 📚
+colorFrom: purple
+colorTo: green
+sdk: gradio
+sdk_version: 2.9.4
+app_file: app.py
+pinned: false
+license: apache-2.0
+duplicated_from: hylee/White-box-Cartoonization
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference

app.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+import argparse
+import functools
+import os
+import pathlib
+import sys
+from typing import Callable
+import uuid
+
+import gradio as gr
+import huggingface_hub
+import numpy as np
+import PIL.Image
+
+from io import BytesIO
+from wbc.cartoonize import Cartoonize
+
+ORIGINAL_REPO_URL = 'https://github.com/SystemErrorWang/White-box-Cartoonization'
+TITLE = 'SystemErrorWang/White-box-Cartoonization'
+DESCRIPTION = f"""This is a demo for {ORIGINAL_REPO_URL}.
+
+"""
+ARTICLE = """
+
+"""
+
+SAFEHASH = [x for x in "0123456789-abcdefghijklmnopqrstuvwxyz_ABCDEFGHIJKLMNOPQRSTUVWXYZ"]
+def compress_UUID():
+    '''
+    根据http://www.ietf.org/rfc/rfc1738.txt,由uuid编码扩bai大字符域生成du串
+    包括:[0-9a-zA-Z\-_]共64个
+    长度:(32-2)/3*2=20
+    备注:可在地球上人zhi人都用,使用100年不重复(2^120)
+    :return:String
+    '''
+    row = str(uuid.uuid4()).replace('-', '')
+    safe_code = ''
+    for i in range(10):
+        enbin = "%012d" % int(bin(int(row[i * 3] + row[i * 3 + 1] + row[i * 3 + 2], 16))[2:], 10)
+        safe_code += (SAFEHASH[int(enbin[0:6], 2)] + SAFEHASH[int(enbin[6:12], 2)])
+    safe_code = safe_code.replace('-', '')
+    return safe_code
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', type=str, default='cpu')
+    parser.add_argument('--theme', type=str)
+    parser.add_argument('--live', action='store_true')
+    parser.add_argument('--share', action='store_true')
+    parser.add_argument('--port', type=int)
+    parser.add_argument('--disable-queue',
+                        dest='enable_queue',
+                        action='store_false')
+    parser.add_argument('--allow-flagging', type=str, default='never')
+    parser.add_argument('--allow-screenshot', action='store_true')
+    return parser.parse_args()
+
+def run(
+        image,
+        cartoonize : Cartoonize
+) -> tuple[PIL.Image.Image]:
+
+    out_path = compress_UUID()+'.png'
+    cartoonize.run_sigle(image.name, out_path)
+
+    return PIL.Image.open(out_path)
+
+
+def main():
+    gr.close_all()
+
+    args = parse_args()
+
+    cartoonize = Cartoonize(os.path.join(os.path.dirname(os.path.abspath(__file__)),'wbc/saved_models/'))
+
+    func = functools.partial(run, cartoonize=cartoonize)
+    func = functools.update_wrapper(func, run)
+
+    gr.Interface(
+        func,
+        [
+            gr.inputs.Image(type='file', label='Input Image'),
+        ],
+        [
+            gr.outputs.Image(
+                type='pil',
+                label='Result'),
+        ],
+        # examples=examples,
+        theme=args.theme,
+        title=TITLE,
+        description=DESCRIPTION,
+        article=ARTICLE,
+        allow_screenshot=args.allow_screenshot,
+        allow_flagging=args.allow_flagging,
+        live=args.live,
+    ).launch(
+        enable_queue=args.enable_queue,
+        server_port=args.port,
+        share=args.share,
+    )
+
+
+if __name__ == '__main__':
+    main()

packages.txt

@@ -0,0 +1,2 @@
+
+

requirements.txt

@@ -0,0 +1,5 @@
+opencv-python-headless==4.5.5.62
+Pillow==9.0.1
+scipy==1.7.3
+tensorflow-gpu==1.14.0
+scikit-image==0.14.5

wbc/cartoonize.py

@@ -0,0 +1,112 @@
+import os
+import cv2
+import numpy as np
+import tensorflow as tf
+import wbc.network as network
+import wbc.guided_filter as guided_filter
+from tqdm import tqdm
+
+
+def resize_crop(image):
+    h, w, c = np.shape(image)
+    if min(h, w) > 720:
+        if h > w:
+            h, w = int(720 * h / w), 720
+        else:
+            h, w = 720, int(720 * w / h)
+    image = cv2.resize(image, (w, h),
+                       interpolation=cv2.INTER_AREA)
+    h, w = (h // 8) * 8, (w // 8) * 8
+    image = image[:h, :w, :]
+    return image
+
+
+def cartoonize(load_folder, save_folder, model_path):
+    print(model_path)
+    input_photo = tf.placeholder(tf.float32, [1, None, None, 3])
+    network_out = network.unet_generator(input_photo)
+    final_out = guided_filter.guided_filter(input_photo, network_out, r=1, eps=5e-3)
+
+    all_vars = tf.trainable_variables()
+    gene_vars = [var for var in all_vars if 'generator' in var.name]
+    saver = tf.train.Saver(var_list=gene_vars)
+
+    config = tf.ConfigProto()
+    config.gpu_options.allow_growth = True
+    sess = tf.Session(config=config)
+
+    sess.run(tf.global_variables_initializer())
+    saver.restore(sess, tf.train.latest_checkpoint(model_path))
+    name_list = os.listdir(load_folder)
+    for name in tqdm(name_list):
+        try:
+            load_path = os.path.join(load_folder, name)
+            save_path = os.path.join(save_folder, name)
+            image = cv2.imread(load_path)
+            image = resize_crop(image)
+            batch_image = image.astype(np.float32) / 127.5 - 1
+            batch_image = np.expand_dims(batch_image, axis=0)
+            output = sess.run(final_out, feed_dict={input_photo: batch_image})
+            output = (np.squeeze(output) + 1) * 127.5
+            output = np.clip(output, 0, 255).astype(np.uint8)
+            cv2.imwrite(save_path, output)
+        except:
+            print('cartoonize {} failed'.format(load_path))
+
+
+class Cartoonize:
+    def __init__(self, model_path):
+        print(model_path)
+        self.input_photo = tf.placeholder(tf.float32, [1, None, None, 3])
+        network_out = network.unet_generator(self.input_photo)
+        self.final_out = guided_filter.guided_filter(self.input_photo, network_out, r=1, eps=5e-3)
+
+        all_vars = tf.trainable_variables()
+        gene_vars = [var for var in all_vars if 'generator' in var.name]
+        saver = tf.train.Saver(var_list=gene_vars)
+
+        config = tf.ConfigProto()
+        config.gpu_options.allow_growth = True
+        self.sess = tf.Session(config=config)
+
+        self.sess.run(tf.global_variables_initializer())
+        saver.restore(self.sess, tf.train.latest_checkpoint(model_path))
+
+    def run(self, load_folder, save_folder):
+        name_list = os.listdir(load_folder)
+        for name in tqdm(name_list):
+            try:
+                load_path = os.path.join(load_folder, name)
+                save_path = os.path.join(save_folder, name)
+                image = cv2.imread(load_path)
+                image = resize_crop(image)
+                batch_image = image.astype(np.float32) / 127.5 - 1
+                batch_image = np.expand_dims(batch_image, axis=0)
+                output = self.sess.run(self.final_out, feed_dict={self.input_photo: batch_image})
+                output = (np.squeeze(output) + 1) * 127.5
+                output = np.clip(output, 0, 255).astype(np.uint8)
+                cv2.imwrite(save_path, output)
+            except:
+                print('cartoonize {} failed'.format(load_path))
+
+    def run_sigle(self, load_path, save_path):
+        try:
+            image = cv2.imread(load_path)
+            image = resize_crop(image)
+            batch_image = image.astype(np.float32) / 127.5 - 1
+            batch_image = np.expand_dims(batch_image, axis=0)
+            output = self.sess.run(self.final_out, feed_dict={self.input_photo: batch_image})
+            output = (np.squeeze(output) + 1) * 127.5
+            output = np.clip(output, 0, 255).astype(np.uint8)
+            cv2.imwrite(save_path, output)
+        except:
+            print('cartoonize {} failed'.format(load_path))
+
+
+if __name__ == '__main__':
+    model_path = 'saved_models'
+    load_folder = 'test_images'
+    save_folder = 'cartoonized_images'
+    if not os.path.exists(save_folder):
+        os.mkdir(save_folder)
+    cartoonize(load_folder, save_folder, model_path)

wbc/guided_filter.py

@@ -0,0 +1,87 @@
+import tensorflow as tf
+import numpy as np
+
+
+
+
+def tf_box_filter(x, r):
+    k_size = int(2*r+1)
+    ch = x.get_shape().as_list()[-1]
+    weight = 1/(k_size**2)
+    box_kernel = weight*np.ones((k_size, k_size, ch, 1))
+    box_kernel = np.array(box_kernel).astype(np.float32)
+    output = tf.nn.depthwise_conv2d(x, box_kernel, [1, 1, 1, 1], 'SAME')
+    return output
+
+
+
+def guided_filter(x, y, r, eps=1e-2):
+    
+    x_shape = tf.shape(x)
+    #y_shape = tf.shape(y)
+
+    N = tf_box_filter(tf.ones((1, x_shape[1], x_shape[2], 1), dtype=x.dtype), r)
+
+    mean_x = tf_box_filter(x, r) / N
+    mean_y = tf_box_filter(y, r) / N
+    cov_xy = tf_box_filter(x * y, r) / N - mean_x * mean_y
+    var_x  = tf_box_filter(x * x, r) / N - mean_x * mean_x
+
+    A = cov_xy / (var_x + eps)
+    b = mean_y - A * mean_x
+
+    mean_A = tf_box_filter(A, r) / N
+    mean_b = tf_box_filter(b, r) / N
+
+    output = mean_A * x + mean_b
+
+    return output
+
+
+
+def fast_guided_filter(lr_x, lr_y, hr_x, r=1, eps=1e-8):
+    
+    #assert lr_x.shape.ndims == 4 and lr_y.shape.ndims == 4 and hr_x.shape.ndims == 4
+   
+    lr_x_shape = tf.shape(lr_x)
+    #lr_y_shape = tf.shape(lr_y)
+    hr_x_shape = tf.shape(hr_x)
+    
+    N = tf_box_filter(tf.ones((1, lr_x_shape[1], lr_x_shape[2], 1), dtype=lr_x.dtype), r)
+
+    mean_x = tf_box_filter(lr_x, r) / N
+    mean_y = tf_box_filter(lr_y, r) / N
+    cov_xy = tf_box_filter(lr_x * lr_y, r) / N - mean_x * mean_y
+    var_x  = tf_box_filter(lr_x * lr_x, r) / N - mean_x * mean_x
+
+    A = cov_xy / (var_x + eps)
+    b = mean_y - A * mean_x
+
+    mean_A = tf.image.resize_images(A, hr_x_shape[1: 3])
+    mean_b = tf.image.resize_images(b, hr_x_shape[1: 3])
+
+    output = mean_A * hr_x + mean_b
+    
+    return output
+
+
+if __name__ == '__main__':
+    import cv2
+    from tqdm import tqdm
+
+    input_photo = tf.placeholder(tf.float32, [1, None, None, 3])
+    #input_superpixel = tf.placeholder(tf.float32, [16, 256, 256, 3])
+    output = guided_filter(input_photo, input_photo, 5, eps=1)
+    image = cv2.imread('output_figure1/cartoon2.jpg')
+    image = image/127.5 - 1
+    image = np.expand_dims(image, axis=0)
+
+    config = tf.ConfigProto()
+    config.gpu_options.allow_growth = True
+    sess = tf.Session(config=config)
+    sess.run(tf.global_variables_initializer())
+
+    out = sess.run(output, feed_dict={input_photo: image})
+    out = (np.squeeze(out)+1)*127.5
+    out = np.clip(out, 0, 255).astype(np.uint8)
+    cv2.imwrite('output_figure1/cartoon2_filter.jpg', out)

wbc/network.py

@@ -0,0 +1,62 @@
+import tensorflow as tf
+import numpy as np
+import tensorflow.contrib.slim as slim
+
+
+
+def resblock(inputs, out_channel=32, name='resblock'):
+    
+    with tf.variable_scope(name):
+        
+        x = slim.convolution2d(inputs, out_channel, [3, 3], 
+                               activation_fn=None, scope='conv1')
+        x = tf.nn.leaky_relu(x)
+        x = slim.convolution2d(x, out_channel, [3, 3], 
+                               activation_fn=None, scope='conv2')
+        
+        return x + inputs
+
+
+
+
+def unet_generator(inputs, channel=32, num_blocks=4, name='generator', reuse=False):
+    with tf.variable_scope(name, reuse=reuse):
+        
+        x0 = slim.convolution2d(inputs, channel, [7, 7], activation_fn=None)
+        x0 = tf.nn.leaky_relu(x0)
+        
+        x1 = slim.convolution2d(x0, channel, [3, 3], stride=2, activation_fn=None)
+        x1 = tf.nn.leaky_relu(x1)
+        x1 = slim.convolution2d(x1, channel*2, [3, 3], activation_fn=None)
+        x1 = tf.nn.leaky_relu(x1)
+        
+        x2 = slim.convolution2d(x1, channel*2, [3, 3], stride=2, activation_fn=None)
+        x2 = tf.nn.leaky_relu(x2)
+        x2 = slim.convolution2d(x2, channel*4, [3, 3], activation_fn=None)
+        x2 = tf.nn.leaky_relu(x2)
+        
+        for idx in range(num_blocks):
+            x2 = resblock(x2, out_channel=channel*4, name='block_{}'.format(idx))
+            
+        x2 = slim.convolution2d(x2, channel*2, [3, 3], activation_fn=None)
+        x2 = tf.nn.leaky_relu(x2)
+        
+        h1, w1 = tf.shape(x2)[1], tf.shape(x2)[2]
+        x3 = tf.image.resize_bilinear(x2, (h1*2, w1*2))
+        x3 = slim.convolution2d(x3+x1, channel*2, [3, 3], activation_fn=None)
+        x3 = tf.nn.leaky_relu(x3)
+        x3 = slim.convolution2d(x3, channel, [3, 3], activation_fn=None)
+        x3 = tf.nn.leaky_relu(x3)
+
+        h2, w2 = tf.shape(x3)[1], tf.shape(x3)[2]
+        x4 = tf.image.resize_bilinear(x3, (h2*2, w2*2))
+        x4 = slim.convolution2d(x4+x0, channel, [3, 3], activation_fn=None)
+        x4 = tf.nn.leaky_relu(x4)
+        x4 = slim.convolution2d(x4, 3, [7, 7], activation_fn=None)
+        
+        return x4
+
+if __name__ == '__main__':
+    
+
+    pass

wbc/saved_models/checkpoint

@@ -0,0 +1,3 @@
+model_checkpoint_path: "model-33999"
+all_model_checkpoint_paths: "model-33999"
+all_model_checkpoint_paths: "model-37499"

wbc/saved_models/model-33999.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e2df1a5aa86faa4f979720bfc2436f79333a480876f8d6790b7671cf50fe75b
+size 5868300