Spaces:
Sleeping
Sleeping
jsaplication
commited on
Commit
•
bf18f18
1
Parent(s):
b089cbd
Update app.py
Browse files
app.py
CHANGED
@@ -1,51 +1,96 @@
|
|
|
|
1 |
import os
|
2 |
import cv2
|
3 |
import torch
|
|
|
4 |
from flask import Flask, request, jsonify, send_file
|
|
|
|
|
|
|
5 |
|
6 |
-
# Importe as classes e funções necessárias para seus modelos aqui
|
7 |
-
|
8 |
-
# Carregue os modelos
|
9 |
model_realesr = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
|
10 |
model_path_realesr = 'realesr-general-x4v3.pth'
|
11 |
|
12 |
-
|
13 |
-
|
|
|
|
|
|
|
|
|
14 |
model_gfpgan_1_4 = GFPGANer(model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
|
15 |
|
16 |
-
# Defina o modelo RestoreFormer se necessário
|
17 |
-
# model_restoreformer = ...
|
18 |
|
19 |
-
|
20 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
21 |
|
22 |
-
|
23 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
24 |
|
25 |
app = Flask(__name__)
|
26 |
|
27 |
@app.route('/reconstruir', methods=['POST'])
|
28 |
def reconstruir_imagem():
|
29 |
try:
|
30 |
-
version = request.form.get('version',
|
31 |
-
scale = int(request.form.get('scale',
|
32 |
img_file = request.files['imagem']
|
33 |
|
34 |
temp_filename = 'temp.jpg'
|
35 |
img_file.save(temp_filename)
|
36 |
|
37 |
-
if version == 'v1.2':
|
38 |
-
face_enhancer = model_gfpgan_1_2
|
39 |
-
elif version == 'v1.3':
|
40 |
-
face_enhancer = model_gfpgan_1_3
|
41 |
-
elif version == 'v1.4':
|
42 |
-
face_enhancer = model_gfpgan_1_4
|
43 |
-
# Adicione mais condições para outros modelos, se necessário
|
44 |
-
|
45 |
output, save_path = inference(temp_filename, version, scale)
|
46 |
|
47 |
if output is not None:
|
48 |
-
return send_file(save_path, mimetype='image/jpeg')
|
|
|
|
|
|
|
49 |
else:
|
50 |
return jsonify({'error': 'Falha na reconstrução da imagem'})
|
51 |
|
@@ -53,4 +98,4 @@ def reconstruir_imagem():
|
|
53 |
return jsonify({'error': str(e)})
|
54 |
|
55 |
if __name__ == '__main__':
|
56 |
-
app.run(host='0.0.0.0', port=80)
|
|
|
1 |
+
|
2 |
import os
|
3 |
import cv2
|
4 |
import torch
|
5 |
+
from gfpgan.utils import GFPGANer
|
6 |
from flask import Flask, request, jsonify, send_file
|
7 |
+
from basicsr.archs.srvgg_arch import SRVGGNetCompact
|
8 |
+
from realesrgan.utils import RealESRGANer
|
9 |
+
import base64
|
10 |
|
|
|
|
|
|
|
11 |
model_realesr = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
|
12 |
model_path_realesr = 'realesr-general-x4v3.pth'
|
13 |
|
14 |
+
# Background enhancer with RealESRGAN
|
15 |
+
model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
|
16 |
+
model_path = 'realesr-general-x4v3.pth'
|
17 |
+
half = True if torch.cuda.is_available() else False
|
18 |
+
upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
|
19 |
+
|
20 |
model_gfpgan_1_4 = GFPGANer(model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
|
21 |
|
|
|
|
|
22 |
|
23 |
+
os.makedirs('output', exist_ok=True)
|
24 |
+
|
25 |
+
|
26 |
+
# def inference(img, version, scale, weight):
|
27 |
+
def inference(img, version, scale):
|
28 |
+
# weight /= 100
|
29 |
+
print(img, version, scale)
|
30 |
+
try:
|
31 |
+
extension = os.path.splitext(os.path.basename(str(img)))[1]
|
32 |
+
img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
|
33 |
+
if len(img.shape) == 3 and img.shape[2] == 4:
|
34 |
+
img_mode = 'RGBA'
|
35 |
+
elif len(img.shape) == 2: # for gray inputs
|
36 |
+
img_mode = None
|
37 |
+
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
38 |
+
else:
|
39 |
+
img_mode = None
|
40 |
+
|
41 |
+
h, w = img.shape[0:2]
|
42 |
+
if h < 300:
|
43 |
+
img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
|
44 |
+
|
45 |
+
if version == 'v1.4':
|
46 |
+
face_enhancer = GFPGANer(
|
47 |
+
model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
|
48 |
+
|
49 |
+
try:
|
50 |
+
# _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True, weight=weight)
|
51 |
+
_, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
|
52 |
+
except RuntimeError as error:
|
53 |
+
print('Error', error)
|
54 |
|
55 |
+
try:
|
56 |
+
if scale != 2:
|
57 |
+
interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
|
58 |
+
h, w = img.shape[0:2]
|
59 |
+
output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
|
60 |
+
except Exception as error:
|
61 |
+
print('wrong scale input.', error)
|
62 |
+
if img_mode == 'RGBA': # RGBA images should be saved in png format
|
63 |
+
extension = 'png'
|
64 |
+
else:
|
65 |
+
extension = 'jpg'
|
66 |
+
save_path = f'output/out.{extension}'
|
67 |
+
cv2.imwrite(save_path, output)
|
68 |
+
|
69 |
+
output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
|
70 |
+
return output, save_path
|
71 |
+
except Exception as error:
|
72 |
+
print('global exception', error)
|
73 |
+
return None, None
|
74 |
|
75 |
app = Flask(__name__)
|
76 |
|
77 |
@app.route('/reconstruir', methods=['POST'])
|
78 |
def reconstruir_imagem():
|
79 |
try:
|
80 |
+
version = request.form.get('version',"v1.4")
|
81 |
+
scale = int(request.form.get('scale',2))
|
82 |
img_file = request.files['imagem']
|
83 |
|
84 |
temp_filename = 'temp.jpg'
|
85 |
img_file.save(temp_filename)
|
86 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
87 |
output, save_path = inference(temp_filename, version, scale)
|
88 |
|
89 |
if output is not None:
|
90 |
+
# return send_file(save_path, mimetype='image/jpeg')
|
91 |
+
with open(save_path, 'rb') as image_file:
|
92 |
+
encoded_image = base64.b64encode(image_file.read()).decode('utf-8')
|
93 |
+
return jsonify({'image_base64': encoded_image})
|
94 |
else:
|
95 |
return jsonify({'error': 'Falha na reconstrução da imagem'})
|
96 |
|
|
|
98 |
return jsonify({'error': str(e)})
|
99 |
|
100 |
if __name__ == '__main__':
|
101 |
+
app.run(host='0.0.0.0', port=80)
|