File size: 4,211 Bytes
c080fe4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d9451cf
 
c080fe4
b326d68
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c080fe4
 
 
 
 
 
 
d9451cf
 
c080fe4
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
import os
os.system("hub install deoldify==1.0.1")
import gradio as gr
import paddlehub as hub
from pathlib import Path
from datetime import datetime
from typing import Optional


model = hub.Module(name='deoldify')
# NOTE:  Max is 45 with 11GB video cards. 35 is a good default
render_factor=35


def colorize_image(image):
    # now = datetime.now().strftime("%Y%m%d-%H%M%S-%f")
    if not os.path.exists("./output"):
        os.makedirs("./output")
    # if image is not None:
    #     image.save(f"./output/{now}-input.jpg")    
    model.predict(image.name)
    
    return f'./output/DeOldify/'+Path(image.name).stem+".png"

# def inference(img, version, scale, weight):
def inference(img, version, scale):
    # weight /= 100
    print(img, version, scale)
    try:
        extension = os.path.splitext(os.path.basename(str(img)))[1]
        img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
        if len(img.shape) == 3 and img.shape[2] == 4:
            img_mode = 'RGBA'
        elif len(img.shape) == 2:  # for gray inputs
            img_mode = None
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        else:
            img_mode = None
        h, w = img.shape[0:2]
        if h < 300:
            img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)

        try:
            # _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True, weight=weight)
            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
        except RuntimeError as error:
            print('Error', error)
        try:
            if scale != 2:
                interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
                h, w = img.shape[0:2]
                output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
        except Exception as error:
            print('wrong scale input.', error)
        if img_mode == 'RGBA':  # RGBA images should be saved in png format
            extension = 'png'
        else:
            extension = 'jpg'
        save_path = f'output/out.{extension}'
        cv2.imwrite(save_path, output)
output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
        return output, save_path
    except Exception as error:
        print('global exception', error)
        return None, None



demo = gr.Interface(
    inference, [
        gr.inputs.Image(type="filepath", label="Input"),
        # gr.inputs.Radio(['v1.2', 'v1.3', 'v1.4', 'RestoreFormer', 'CodeFormer'], type="value", default='v1.4', label='version'),
        gr.inputs.Radio(['v1.2', 'v1.3', 'v1.4', 'RestoreFormer','CodeFormer','RealESR-General-x4v3'], type="value", default='v1.4', label='version'),
        gr.inputs.Number(label="Rescaling factor", default=2),
        # gr.Slider(0, 100, label='Weight, only for CodeFormer. 0 for better quality, 100 for better identity', default=50)
    ], [
        gr.outputs.Image(type="numpy", label="Output (The whole image)"),
        gr.outputs.File(label="Download the output image")
    ],
    title=title,
    description=description,
    article=article,
    # examples=[['AI-generate.jpg', 'v1.4', 2, 50], ['lincoln.jpg', 'v1.4', 2, 50], ['Blake_Lively.jpg', 'v1.4', 2, 50],
    #           ['10045.png', 'v1.4', 2, 50]]).launch()
    examples=[['a1.jpg', 'v1.4', 2], ['a2.jpg', 'v1.4', 2], ['a3.jpg', 'v1.4', 2],['a4.jpg', 'v1.4', 2]])
    
demo.queue(concurrency_count=4)
demo.launch()



def create_interface():
    with gr.Blocks() as enhancer:
        gr.Markdown("Colorize old black & white photos")
        with gr.Column(scale=1, label = "Colorize photo", visible=True) as colorize_column:
            colorize_input = gr.Image(type="file")
            colorize_button = gr.Button("Colorize!")
            colorize_output = gr.outputs.Image(type="numpy", label="Output (The Colorized image)")
            download_colorize_button = gr.outputs.File(label="Download the colorized image!")
            colorize_button.click(colorize_image, inputs=colorize_input, outputs=colorize_output)
    enhancer.launch()

def run_code():
    create_interface()

# The main function
run_code()