PoseCorrection/backend_refactor/route/upload.py

import logging
import uuid

from PIL import Image
from flask import request, Blueprint

from config import UPLOAD_DIR
from route.util import make_response

bp = Blueprint("upload", __name__, url_prefix='/')


def upload(type):
    # 检查是否正确上传了文件
    if 'file' not in request.files:
        return make_response(400, "no file part.", {})

    file = request.files['file']
    if file.filename == '':
        return make_response(400, "no selected file.", {})

    # 随机 uuid 名称
    random_name = str(uuid.uuid4())

    # 图片格式检查 & 转换
    if type == "photo":
        try:
            img = Image.open(file.stream)
            img.verify()
            img = Image.open(file.stream)
            img = img.convert("RGB")
            img.save(open(UPLOAD_DIR + random_name + ".jpg", "wb"))
        except Exception as e:
            logging.exception(e)
            return make_response(400, "cannot identify image file.", {})
    elif type == "video":
        if not file.filename.endswith(".mp4"):
            return make_response(400, "only mp4 videos are accepted.", {})
        file.save(UPLOAD_DIR + random_name + ".mp4")
    else:
        assert False, "unknown upload format"

    return make_response(201, "success", {"uuid": random_name})


@bp.route("/photos", methods=["POST"])
def upload_photo():
    return upload("photo")


@bp.route("/videos", methods=["POST"])
def upload_video():
    return upload("video")