jsphoto-api / app.py
jsaplication's picture
Update app.py
3c28963
raw
history blame
No virus
1.97 kB
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__)
@app.route('/reconstruir', methods=['POST'])
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)