Spaces:
Sleeping
Sleeping
File size: 4,127 Bytes
bf18f18 d6ac96e bf18f18 3c28963 bf18f18 3c28963 d6ac96e bf18f18 3c28963 d6ac96e 7686b0e bf18f18 d6ac96e bf18f18 7686b0e bf18f18 d6ac96e 3c28963 d6ac96e 3c28963 d6ac96e 602b069 bf18f18 3c28963 d6ac96e 602b069 d6ac96e 602b069 d6ac96e 602b069 d6ac96e 602b069 3c28963 602b069 a4981cb 3c28963 f8fca07 |
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 |
import os
import cv2
import torch
from gfpgan.utils import GFPGANer
from flask import Flask, request, jsonify, send_file
from basicsr.archs.srvgg_arch import SRVGGNetCompact
from realesrgan.utils import RealESRGANer
import base64
model_realesr = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
model_path_realesr = 'realesr-general-x4v3.pth'
# Background enhancer with RealESRGAN
model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
model_path = 'realesr-general-x4v3.pth'
half = True if torch.cuda.is_available() else False
upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
model_gfpgan_1_4 = GFPGANer(model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
#os.makedirs('output', exist_ok=True)
# 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)
if version == 'v1.4':
face_enhancer = GFPGANer(
model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
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'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
app = Flask(__name__)
@app.route('/reconstruir', methods=['POST'])
def reconstruir_imagem():
try:
token = request.form.get('token')
version = request.form.get('version',"v1.4")
scale = int(request.form.get('scale',2))
img_file = request.files['imagem']
if token == "api_key_abcd":
temp_filename = 'temp.jpg'
img_file.save(temp_filename)
output, save_path = inference(temp_filename, version, scale)
if output is not None:
# return send_file(save_path, mimetype='image/jpeg')
with open(save_path, 'rb') as image_file:
encoded_image = base64.b64encode(image_file.read()).decode('utf-8')
return jsonify({'status': 'success', 'message':'Imagem restaurada com sucesso', 'image_base64': encoded_image})
else:
return jsonify({'status': 'error', 'message':'Falha na reconstrução da imagem'})
else:
return jsonify({'status': 'error', 'message': 'Token invalido'})
except Exception as e:
return jsonify({'status': 'error', 'message': str(e)})
if __name__ == '__main__':
app.run(host='0.0.0.0', port=80)
|