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)