weimc/White-box-Cartoonization

1#!/usr/bin/env python2 3from __future__ import annotations4import argparse5import functools6import os7import pathlib8import sys9from typing import Callable10import uuid11 12import gradio as gr13import huggingface_hub14import numpy as np15import PIL.Image16 17from io import BytesIO18from wbc.cartoonize import Cartoonize19 20ORIGINAL_REPO_URL = 'https://github.com/SystemErrorWang/White-box-Cartoonization'21TITLE = 'SystemErrorWang/White-box-Cartoonization'22DESCRIPTION = f"""This is a demo for {ORIGINAL_REPO_URL}.23 24"""25ARTICLE = """26 27"""28 29SAFEHASH = [x for x in "0123456789-abcdefghijklmnopqrstuvwxyz_ABCDEFGHIJKLMNOPQRSTUVWXYZ"]30def compress_UUID():31    '''32    根据http://www.ietf.org/rfc/rfc1738.txt,由uuid编码扩bai大字符域生成du串33    包括:[0-9a-zA-Z\-_]共64个34    长度:(32-2)/3*2=2035    备注:可在地球上人zhi人都用,使用100年不重复(2^120)36    :return:String37    '''38    row = str(uuid.uuid4()).replace('-', '')39    safe_code = ''40    for i in range(10):41        enbin = "%012d" % int(bin(int(row[i * 3] + row[i * 3 + 1] + row[i * 3 + 2], 16))[2:], 10)42        safe_code += (SAFEHASH[int(enbin[0:6], 2)] + SAFEHASH[int(enbin[6:12], 2)])43    safe_code = safe_code.replace('-', '')44    return safe_code45 46 47def parse_args() -> argparse.Namespace:48    parser = argparse.ArgumentParser()49    parser.add_argument('--device', type=str, default='cpu')50    parser.add_argument('--theme', type=str)51    parser.add_argument('--live', action='store_true')52    parser.add_argument('--share', action='store_true')53    parser.add_argument('--port', type=int)54    parser.add_argument('--disable-queue',55                        dest='enable_queue',56                        action='store_false')57    parser.add_argument('--allow-flagging', type=str, default='never')58    parser.add_argument('--allow-screenshot', action='store_true')59    return parser.parse_args()60 61def run(62        image,63        cartoonize : Cartoonize64) -> tuple[PIL.Image.Image]:65 66    out_path = compress_UUID()+'.png'67    cartoonize.run_sigle(image.name, out_path)68 69    return PIL.Image.open(out_path)70 71 72def main():73    gr.close_all()74 75    args = parse_args()76 77    cartoonize = Cartoonize(os.path.join(os.path.dirname(os.path.abspath(__file__)),'wbc/saved_models/'))78 79    func = functools.partial(run, cartoonize=cartoonize)80    func = functools.update_wrapper(func, run)81 82    gr.Interface(83        func,84        [85            gr.inputs.Image(type='file', label='Input Image'),86        ],87        [88            gr.outputs.Image(89                type='pil',90                label='Result'),91        ],92        # examples=examples,93        theme=args.theme,94        title=TITLE,95        description=DESCRIPTION,96        article=ARTICLE,97        allow_screenshot=args.allow_screenshot,98        allow_flagging=args.allow_flagging,99        live=args.live,100    ).launch(101        enable_queue=args.enable_queue,102        server_port=args.port,103        share=args.share,104    )105 106 107if __name__ == '__main__':108    main()109