Hugging Face's logo

Spaces:
davertor
/
colorizing_images
Runtime error

App Files Community
1
colorizing_images
File size: 5,736 Bytes 
0cb9530
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
 
#importing the libraries
import os, sys, re
import streamlit as st
import PIL
from PIL import Image
import cv2
import numpy as np
import uuid

import ssl
ssl._create_default_https_context = ssl._create_unverified_context

# Import torch libraries
import fastai
import torch

# Import util functions from app_utils
from app_utils import download
from app_utils import generate_random_filename
from app_utils import clean_me
from app_utils import clean_all
from app_utils import get_model_bin
from app_utils import convertToJPG

# Import util functions from deoldify
# NOTE:  This must be the first call in order to work properly!
from deoldify import device
from deoldify.device_id import DeviceId
#choices:  CPU, GPU0...GPU7
device.set(device=DeviceId.CPU)
from deoldify.visualize import *


####### INPUT PARAMS ###########
model_folder = 'models/'
max_img_size = 800
################################

@st.cache(allow_output_mutation=True)
def load_model(model_dir, option):
    if option.lower() == 'artistic':
        model_url = 'https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth'
        get_model_bin(model_url, os.path.join(model_dir, "ColorizeArtistic_gen.pth"))
        colorizer = get_image_colorizer(artistic=True)
    elif option.lower() == 'stable':
        model_url = "https://www.dropbox.com/s/usf7uifrctqw9rl/ColorizeStable_gen.pth?dl=0"
        get_model_bin(model_url, os.path.join(model_dir, "ColorizeStable_gen.pth"))
        colorizer = get_image_colorizer(artistic=False)

    return colorizer

def resize_img(input_img, max_size):
    img = input_img.copy()
    img_height, img_width = img.shape[0],img.shape[1]

    if max(img_height, img_width) > max_size:
        if img_height > img_width:
            new_width = img_width*(max_size/img_height)
            new_height = max_size
            resized_img = cv2.resize(img,(int(new_width), int(new_height)))
            return resized_img

        elif img_height <= img_width:
            new_width = img_height*(max_size/img_width)
            new_height = max_size
            resized_img = cv2.resize(img,(int(new_width), int(new_height)))
            return resized_img

    return img

def get_image_download_link(img,filename,text):
    button_uuid = str(uuid.uuid4()).replace('-', '')
    button_id = re.sub('\d+', '', button_uuid)
    
    custom_css = f""" 
        <style>
            #{button_id} {{
                background-color: rgb(255, 255, 255);
                color: rgb(38, 39, 48);
                padding: 0.25em 0.38em;
                position: relative;
                text-decoration: none;
                border-radius: 4px;
                border-width: 1px;
                border-style: solid;
                border-color: rgb(230, 234, 241);
                border-image: initial;

            }} 
            #{button_id}:hover {{
                border-color: rgb(246, 51, 102);
                color: rgb(246, 51, 102);
            }}
            #{button_id}:active {{
                box-shadow: none;
                background-color: rgb(246, 51, 102);
                color: white;
                }}
        </style> """

    buffered = BytesIO()
    img.save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    href =  custom_css + f'<a href="data:file/txt;base64,{img_str}" id="{button_id}" download="{filename}">{text}</a>'
    return href


# General configuration
# st.set_page_config(layout="centered")
st.set_page_config(layout="wide")
st.set_option('deprecation.showfileUploaderEncoding', False)
st.markdown('''
<style>
    .uploadedFile {display: none}
<style>''',
unsafe_allow_html=True)

# Main window configuration
st.title("Black and white colorizer")
st.markdown("This app puts color into your black and white pictures")
title_message = st.empty()

title_message.markdown("**Model loading, please wait** ⌛")

# # Sidebar
color_option = st.sidebar.selectbox('Select colorizer mode',
                                    ('Artistic', 'Stable'))
                                    
# st.sidebar.title('Model parameters')
# det_conf_thres = st.sidebar.slider("Detector confidence threshold", 0.1, 0.9, value=0.5, step=0.1)
# det_nms_thres = st.sidebar.slider("Non-maximum supression IoU", 0.1, 0.9, value=0.4, step=0.1)

# Load models
try:
    colorizer = load_model(model_folder, color_option)
except Exception as e: 
    print(e)
    colorizer = None
    print('Error while loading the model. Please refresh the page')

if colorizer is not None:
    print('Running colorizer')
    title_message.markdown("**To begin, please upload an image** 👇")

    #Choose your own image
    uploaded_file = st.file_uploader("Upload a black and white photo", type=['png', 'jpg', 'jpeg'])

    # show = st.image(use_column_width='auto')
    input_img_pos = st.empty()
    output_img_pos = st.empty()

    if uploaded_file is not None:
        img_name = uploaded_file.name

        pil_img = PIL.Image.open(uploaded_file)
        img_rgb = np.array(pil_img)

        resized_img_rgb = resize_img(img_rgb, max_img_size)
        resized_pil_img = PIL.Image.fromarray(resized_img_rgb)

        title_message.markdown("**Processing your image, please wait** ⌛")

        output_pil_img = colorizer.plot_transformed_pil_image(resized_pil_img, render_factor=35, compare=False)
        
        title_message.markdown("**To begin, please upload an image** 👇")

        # Plot images
        input_img_pos.image(resized_pil_img, 'Input image', use_column_width=True)
        output_img_pos.image(output_pil_img, 'Output image', use_column_width=True)

        st.markdown(get_image_download_link(output_pil_img, img_name, 'Download '+img_name), unsafe_allow_html=True)