Matting/api/file.py

import os
from flask import request, jsonify, send_file
from werkzeug.utils import secure_filename
from werkzeug.datastructures import FileStorage
import uuid
from app import app
from .resp import success_resp, error_resp
import time

UPLOAD_PATH = ""
OUTPUT_DIR = ""


def get_upload_dir():
    return UPLOAD_PATH


def get_output_dir():
    return OUTPUT_DIR


def set_options(upload_path, save_dir):
    global UPLOAD_PATH, OUTPUT_DIR
    UPLOAD_PATH = upload_path
    OUTPUT_DIR = save_dir
    print("dirs", UPLOAD_PATH, OUTPUT_DIR)


@app.route("/outputs/<path:filename>")
def get_outputs_files(filename):
    # 处理后的文件
    return send_file(os.path.join(OUTPUT_DIR, filename), mimetype="image/png")


@app.route("/uploads/<path:filename>")
def get_upload_files(filename):
    # 用户上传的原始文件
    return send_file(os.path.join(UPLOAD_PATH, filename), mimetype="image/png")


@app.route("/upload", methods=['POST'])
def upload_file():
    # 上传图片
    if 'file' not in request.files:
        return jsonify(error_resp("请选择图片！"))
    file = request.files['file']
    file_path, file_id = save_upload_file(file)
    return jsonify(success_resp({
        "fileId": file_id,
        "url": file_url(file_path)
    }))


def file_url(path, output=False):
    if output is True:
        file_id = os.path.relpath(path, OUTPUT_DIR)
        return "/outputs/{}".format(file_id)
    file_id = os.path.relpath(path, UPLOAD_PATH)
    return "/uploads/{}".format(file_id)


def save_upload_file(file: FileStorage):
    # 如果区分用户，加文件前缀或者文件夹名称
    if file.filename == '':
        filename = "unknown.png"
    else:
        filename = secure_filename(file.filename)
    paths = os.path.splitext(filename)
    t = time.strftime("%Y%m%d%H%M%S", time.localtime())
    filename = "{}_{}{}_{}".format(paths[0], t, uuid.uuid4().hex, paths[1])

    file_path = get_upload_file_path(filename)
    file.save(file_path)
    return file_path, filename


def get_output_file_path(filename: str, suffix: str = None, fun: str = None):
    if filename == '':
        filename = "unknown.png"
    else:
        filename = secure_filename(filename)
    paths = os.path.splitext(filename)
    if suffix is None or len(suffix) == 0:
        suffix = time.strftime("%Y%m%d%H%M%S", time.localtime())
    if fun is not None:
        suffix = "{}_{}".format(fun, suffix)
    filename = "{}_{}{}".format(paths[0], suffix, paths[1])
    return os.path.join(OUTPUT_DIR, filename)


def get_upload_file_path(name):
    return os.path.join(UPLOAD_PATH, name)