hamg22 commited on
Commit
67d2e4d
1 Parent(s): 3c9d987

Upload 2 files

Browse files
Files changed (2) hide show
  1. pixelup_demo.py +246 -0
  2. requirements (1).txt +11 -0
pixelup_demo.py ADDED
@@ -0,0 +1,246 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # -*- coding: utf-8 -*-
2
+ """PixelUP DEMO.ipynb
3
+
4
+ Automatically generated by Colaboratory.
5
+
6
+ Original file is located at
7
+ https://colab.research.google.com/drive/1RufGXUNy-7zx5biALX7zEgXb0rM9cWnb
8
+ """
9
+
10
+ import os
11
+ os.sys.path
12
+
13
+ !pip install opencv-python
14
+
15
+ !pip install basicsr
16
+ !pip install facexlib
17
+ !pip install gfpgan
18
+ !pip install tqdm
19
+ !pip install -U gradio
20
+
21
+ # Commented out IPython magic to ensure Python compatibility.
22
+ # %pip install realesrgan
23
+
24
+ import gradio as gr
25
+ import cv2
26
+ import numpy
27
+ import os
28
+ import random
29
+ from basicsr.archs.rrdbnet_arch import RRDBNet
30
+ from basicsr.utils.download_util import load_file_from_url
31
+
32
+ from realesrgan import RealESRGANer
33
+ from realesrgan.archs.srvgg_arch import SRVGGNetCompact
34
+
35
+ last_file = None
36
+ img_mode = "RGBA"
37
+
38
+
39
+ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
40
+ """Real-ESRGAN function to restore (and upscale) images.
41
+ """
42
+ if not img:
43
+ return
44
+
45
+ # Define model parameters
46
+ if model_name == 'RealESRGAN_x4plus': # x4 RRDBNet model
47
+ model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
48
+ netscale = 4
49
+ file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
50
+ elif model_name == 'RealESRNet_x4plus': # x4 RRDBNet model
51
+ model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
52
+ netscale = 4
53
+ file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']
54
+ elif model_name == 'RealESRGAN_x4plus_anime_6B': # x4 RRDBNet model with 6 blocks
55
+ model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
56
+ netscale = 4
57
+ file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
58
+ elif model_name == 'RealESRGAN_x2plus': # x2 RRDBNet model
59
+ model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
60
+ netscale = 2
61
+ file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
62
+ elif model_name == 'realesr-general-x4v3': # x4 VGG-style model (S size)
63
+ model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
64
+ netscale = 4
65
+ file_url = [
66
+ 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',
67
+ 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
68
+ ]
69
+
70
+ # Determine model paths
71
+ model_path = os.path.join('weights', model_name + '.pth')
72
+ if not os.path.isfile(model_path):
73
+ ROOT_DIR = os.path.dirname(os.path.abspath("."))
74
+ for url in file_url:
75
+ # model_path will be updated
76
+ model_path = load_file_from_url(
77
+ url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
78
+
79
+ # Use dni to control the denoise strength
80
+ dni_weight = None
81
+ if model_name == 'realesr-general-x4v3' and denoise_strength != 1:
82
+ wdn_model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')
83
+ model_path = [model_path, wdn_model_path]
84
+ dni_weight = [denoise_strength, 1 - denoise_strength]
85
+
86
+ # Restorer Class
87
+ upsampler = RealESRGANer(
88
+ scale=netscale,
89
+ model_path=model_path,
90
+ dni_weight=dni_weight,
91
+ model=model,
92
+ tile=0,
93
+ tile_pad=10,
94
+ pre_pad=10,
95
+ half=False,
96
+ gpu_id=None
97
+ )
98
+
99
+ # Use GFPGAN for face enhancement
100
+ if face_enhance:
101
+ from gfpgan import GFPGANer
102
+ face_enhancer = GFPGANer(
103
+ model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
104
+ upscale=outscale,
105
+ arch='clean',
106
+ channel_multiplier=2,
107
+ bg_upsampler=upsampler)
108
+
109
+ # Convert the input PIL image to cv2 image, so that it can be processed by realesrgan
110
+ cv_img = numpy.array(img)
111
+ img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)
112
+
113
+ # Apply restoration
114
+ try:
115
+ if face_enhance:
116
+ _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
117
+ else:
118
+ output, _ = upsampler.enhance(img, outscale=outscale)
119
+ except RuntimeError as error:
120
+ print('Error', error)
121
+ print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
122
+ else:
123
+ # Save restored image and return it to the output Image component
124
+ if img_mode == 'RGBA': # RGBA images should be saved in png format
125
+ extension = 'png'
126
+ else:
127
+ extension = 'jpg'
128
+
129
+ out_filename = f"output_{rnd_string(8)}.{extension}"
130
+ cv2.imwrite(out_filename, output)
131
+ global last_file
132
+ last_file = out_filename
133
+ return out_filename
134
+
135
+
136
+ def rnd_string(x):
137
+ """Returns a string of 'x' random characters
138
+ """
139
+ characters = "abcdefghijklmnopqrstuvwxyz_0123456789"
140
+ result = "".join((random.choice(characters)) for i in range(x))
141
+ return result
142
+
143
+
144
+ def reset():
145
+ """Resets the Image components of the Gradio interface and deletes
146
+ the last processed image
147
+ """
148
+ global last_file
149
+ if last_file:
150
+ print(f"Deleting {last_file} ...")
151
+ os.remove(last_file)
152
+ last_file = None
153
+ return gr.update(value=None), gr.update(value=None)
154
+
155
+
156
+ def has_transparency(img):
157
+ """This function works by first checking to see if a "transparency" property is defined
158
+ in the image's info -- if so, we return "True". Then, if the image is using indexed colors
159
+ (such as in GIFs), it gets the index of the transparent color in the palette
160
+ (img.info.get("transparency", -1)) and checks if it's used anywhere in the canvas
161
+ (img.getcolors()). If the image is in RGBA mode, then presumably it has transparency in
162
+ it, but it double-checks by getting the minimum and maximum values of every color channel
163
+ (img.getextrema()), and checks if the alpha channel's smallest value falls below 255.
164
+ https://stackoverflow.com/questions/43864101/python-pil-check-if-image-is-transparent
165
+ """
166
+ if img.info.get("transparency", None) is not None:
167
+ return True
168
+ if img.mode == "P":
169
+ transparent = img.info.get("transparency", -1)
170
+ for _, index in img.getcolors():
171
+ if index == transparent:
172
+ return True
173
+ elif img.mode == "RGBA":
174
+ extrema = img.getextrema()
175
+ if extrema[3][0] < 255:
176
+ return True
177
+ return False
178
+
179
+
180
+ def image_properties(img):
181
+ """Returns the dimensions (width and height) and color mode of the input image and
182
+ also sets the global img_mode variable to be used by the realesrgan function
183
+ """
184
+ global img_mode
185
+ if img:
186
+ try:
187
+ img_mode = img.mode # Get the color mode directly
188
+ properties = f"Width: {img.size[0]}, Height: {img.size[1]} | Color Mode: {img_mode}"
189
+ return properties
190
+ except Exception as e:
191
+ print(f"Error processing image: {e}")
192
+
193
+ return "Invalid image"
194
+
195
+ def main():
196
+ # Gradio Interface
197
+ with gr.Blocks(title="PixelUP Demo", theme="dark") as demo:
198
+
199
+ gr.Markdown(
200
+ """# <div align="center"> </div>
201
+ <div align="center"><img width="200" height="74" src="https://i.imgur.com/UzFFDdL.png"></div>
202
+ """
203
+ )
204
+
205
+ with gr.Accordion("Options/Parameters"):
206
+ with gr.Row():
207
+ model_name = gr.Dropdown(label="Choose Model",
208
+ choices=["RealESRGAN_x4plus", "RealESRNet_x4plus", "RealESRGAN_x4plus_anime_6B",
209
+ "RealESRGAN_x2plus", "realesr-general-x4v3"],
210
+ value="realesr-general-x4v3", show_label=True)
211
+ denoise_strength = gr.Slider(label="Denoise Strength",
212
+ minimum=0, maximum=1, step=0.1, value=0.5)
213
+ outscale = gr.Slider(label="Image Upscaling Factor",
214
+ minimum=1, maximum=10, step=1, value=2, show_label=True)
215
+ face_enhance = gr.Checkbox(label="Face Enhancement using GFPGAN ",
216
+ value=False, show_label=True)
217
+
218
+ with gr.Row():
219
+ with gr.Group():
220
+ input_image = gr.Image(label="Source Image", type="pil", image_mode="RGBA")
221
+ input_image_properties = gr.Textbox(label="Image Properties", max_lines=1)
222
+ output_image = gr.Image(label="Restored Image", image_mode="RGBA")
223
+ with gr.Row():
224
+ restore_btn = gr.Button("Restore Image")
225
+ reset_btn = gr.Button("Reset")
226
+
227
+ # Event listeners:
228
+ input_image.change(fn=image_properties, inputs=input_image, outputs=input_image_properties)
229
+ restore_btn.click(fn=realesrgan,
230
+ inputs=[input_image, model_name, denoise_strength, face_enhance, outscale],
231
+ outputs=output_image)
232
+ reset_btn.click(fn=reset, inputs=[], outputs=[output_image, input_image])
233
+ # reset_btn.click(None, inputs=[], outputs=[input_image], _js="() => (null)\n")
234
+ # Undocumented method to clear a component's value using Javascript
235
+
236
+ gr.Markdown(
237
+ """*Please note that support for animated GIFs is not yet implemented. Should an animated GIF is chosen for restoration,
238
+ the demo will output only the first frame saved in PNG format (to preserve probable transparency).*
239
+ """
240
+ )
241
+
242
+ demo.launch(share=True)
243
+
244
+
245
+ if __name__ == "__main__":
246
+ main()
requirements (1).txt ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch
2
+ torchvision
3
+ numpy
4
+ opencv-python
5
+ Pillow
6
+ basicsr
7
+ facexlib
8
+ gfpgan
9
+ tqdm
10
+ gradio
11
+ realesrgan