File size: 3,783 Bytes
212d7be |
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 |
from flask import Flask, request, jsonify, abort, redirect, url_for, render_template, send_file
from flask_wtf import FlaskForm
from wtforms import StringField, FileField, BooleanField, DecimalField
from wtforms.validators import DataRequired
from flask import after_this_request
import torch
import os
from model.models import Colorizer, Generator
from model.extractor import get_seresnext_extractor
from utils.xdog import XDoGSketcher
from utils.utils import open_json
from denoising.denoiser import FFDNetDenoiser
from datetime import datetime
from inference import colorize_single_image, colorize_images, colorize_cbr
if torch.cuda.is_available():
device = 'cuda'
else:
device = 'cpu'
generator = Generator()
generator.load_state_dict(torch.load('model/generator.pth'))
extractor = get_seresnext_extractor()
extractor.load_state_dict(torch.load('model/extractor.pth'))
colorizer = Colorizer(generator, extractor)
colorizer = colorizer.eval().to(device)
sketcher = XDoGSketcher()
xdog_config = open_json('configs/xdog_config.json')
for key in xdog_config.keys():
if key in sketcher.params:
sketcher.params[key] = xdog_config[key]
denoiser = FFDNetDenoiser(device)
app = Flask(__name__)
app.config.update(dict(
SECRET_KEY="lol kek",
WTF_CSRF_SECRET_KEY="cheburek"
))
color_args = {'colorizer':colorizer, 'sketcher':sketcher, 'device':device, 'dfm' : True}
class SubmitForm(FlaskForm):
file = FileField(validators=[DataRequired()])
denoise = BooleanField(default = 'checked')
denoise_sigma = DecimalField(label = 'Denoise sigma', validators=[DataRequired()], default = 25, places = None)
autohint = BooleanField(default = None)
autohint_sigma = DecimalField(label = 'Autohint sigma', validators=[DataRequired()], default= 0.0003, places = None)
ignore_gray = BooleanField(label = 'Ignore gray autohint', default = None)
@app.route('/img/<path>')
def show_image(path):
return f'<img src="/static/{path}">'
@app.route('/', methods=('GET', 'POST'))
def submit_data():
form = SubmitForm()
if form.validate_on_submit():
input_data = form.file.data
_, ext = os.path.splitext(input_data.filename)
filename = str(datetime.now()) + ext
input_data.save(filename)
color_args['auto_hint'] = form.autohint.data
color_args['auto_hint_sigma'] = float(form.autohint_sigma.data)
color_args['ignore_gray'] = form.ignore_gray.data
color_args['denoiser'] = None
if form.denoise.data:
color_args['denoiser'] = denoiser
color_args['denoiser_sigma'] = float(form.denoise_sigma.data)
if ext.lower() in ('.cbr', '.cbz', '.rar', '.zip'):
result_name = colorize_cbr(filename, color_args)
os.remove(filename)
@after_this_request
def remove_file(response):
try:
os.remove(result_name)
except Exception as error:
app.logger.error("Error removing or closing downloaded file handle", error)
return response
return send_file(result_name, mimetype='application/vnd.comicbook-rar', attachment_filename=result_name, as_attachment=True)
elif ext.lower() in ('.jpg', '.png', ',jpeg'):
random_name = str(datetime.now()) + '.png'
new_image_path = os.path.join('static', random_name)
colorize_single_image(filename, new_image_path, color_args)
os.remove(filename)
return redirect(f'/img/{random_name}')
else:
return 'Wrong format'
return render_template('submit.html', form=form) |