mirror of
https://github.com/no1xuan/colourize.git
synced 2026-08-28 14:22:07 +08:00
贫困的蚊子
67 lines
1.9 KiB
Python
67 lines
1.9 KiB
Python
import time
|
|
import os
|
|
import base64
|
|
from flask import Flask, jsonify
|
|
from utils.form_models import ImageForm
|
|
from datetime import timedelta
|
|
from color_transformer import img_transformer
|
|
|
|
app = Flask(__name__)
|
|
app.config['SECRET_KEY'] = 'PAPERCLUBASDFASDF'
|
|
app.config['WTF_CSRF_ENABLED'] = False # 关闭CSRF保护
|
|
app.send_file_max_age_default = timedelta(seconds=1)
|
|
app.config['UPLOAD_FOLDER'] = './tempImgs/'
|
|
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
|
|
|
|
def img2byte(img_path):
|
|
# 读取文件并返回字节流，文件过大时使用生成器逐块读取
|
|
with open(img_path, 'rb') as f:
|
|
return f.read()
|
|
|
|
def process_and_cleanup(img_filepath):
|
|
try:
|
|
# 图像处理
|
|
ori_bytes = img2byte(img_filepath)
|
|
img_bytes = img_transformer(ori_bytes, method=1)
|
|
|
|
return img_bytes
|
|
finally:
|
|
# 删除临时图像文件
|
|
if os.path.exists(img_filepath):
|
|
os.remove(img_filepath)
|
|
|
|
@app.route('/colourizeImg', methods=['POST'])
|
|
def main():
|
|
form = ImageForm()
|
|
|
|
if form.validate_on_submit():
|
|
base64_data = form.base64_image.data
|
|
img_data = base64.b64decode(base64_data.split(',')[1])
|
|
filename = f"uploaded_{time.time()}.png"
|
|
img_filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
|
|
|
|
# 将图像保存为临时文件
|
|
with open(img_filepath, 'wb') as f:
|
|
f.write(img_data)
|
|
|
|
# 处理图像并清理临时文件
|
|
img_bytes = process_and_cleanup(img_filepath)
|
|
|
|
response_data = {
|
|
"status": 2, # 2成功，1失败
|
|
"processed_image": img_bytes
|
|
}
|
|
|
|
return jsonify(response_data)
|
|
else:
|
|
errors = [error for error in form.errors.values()]
|
|
response_data = {
|
|
"status": 1,
|
|
"message": errors
|
|
}
|
|
|
|
return jsonify(response_data)
|
|
|
|
if __name__ == '__main__':
|
|
app.run(host='127.0.0.1', port=3007, debug=True)
|