Spaces:
Sleeping
Sleeping
import os | |
import cv2 | |
import torch | |
from flask import Flask, request, jsonify, send_file | |
# Importe as classes e funções necessárias para seus modelos aqui | |
# Carregue os modelos | |
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' | |
model_gfpgan_1_2 = GFPGANer(model_path='GFPGANv1.2.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler) | |
model_gfpgan_1_3 = GFPGANer(model_path='GFPGANv1.3.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler) | |
model_gfpgan_1_4 = GFPGANer(model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler) | |
# Defina o modelo RestoreFormer se necessário | |
# model_restoreformer = ... | |
# Defina o modelo CodeFormer se necessário | |
# model_codeformer = ... | |
# Defina o modelo RealESR-General-x4v3 se necessário | |
# model_realesr_general = ... | |
app = Flask(__name__) | |
def reconstruir_imagem(): | |
try: | |
version = request.form.get('version', 'v1.4') | |
scale = int(request.form.get('scale', 2)) | |
img_file = request.files['imagem'] | |
temp_filename = 'temp.jpg' | |
img_file.save(temp_filename) | |
if version == 'v1.2': | |
face_enhancer = model_gfpgan_1_2 | |
elif version == 'v1.3': | |
face_enhancer = model_gfpgan_1_3 | |
elif version == 'v1.4': | |
face_enhancer = model_gfpgan_1_4 | |
# Adicione mais condições para outros modelos, se necessário | |
output, save_path = inference(temp_filename, version, scale) | |
if output is not None: | |
return send_file(save_path, mimetype='image/jpeg') | |
else: | |
return jsonify({'error': 'Falha na reconstrução da imagem'}) | |
except Exception as e: | |
return jsonify({'error': str(e)}) | |
if __name__ == '__main__': | |
app.run(host='0.0.0.0', port=80) |