Matting/api/image.py

import os
from flask import request, logging, send_file, jsonify
from app import app
from .resp import success_resp, error_resp
import numpy as np
from .req import ReplaceForm, is_empty
import tools
import cv2
from .file import get_upload_file_path, get_output_dir, file_url, get_output_file_path, get_file_id

log = logging.create_logger(app)


@app.route("/image/seg", methods=['POST'])
def seg():
    if 'fileId' not in request.values:
        return "请先上传图片！", 500
    file_id = request.values.get('fileId')

    file_path = get_upload_file_path(file_id)
    _, fg, _, path = tools.seg(file_path, get_output_dir())
    # filename = os.path.relpath(path, OUTPUT_DIR)
    return jsonify(success_resp({
        "fileId": file_id,
        "url": file_url(path, True)
    }))


@app.route("/image/replace", methods=["POST"])
def replace():
    form_data = request.form
    form = ReplaceForm(form_data)
    if form.validate() is False:
        return jsonify(error_resp("参数错误！"))

    if is_empty(form.bg_file_id.data) and is_empty(form.background.data):
        return jsonify(error_resp("请选择需要替换的背景！"))

    img_path = get_upload_file_path(form.file_id.data)
    bg_path = form.background.data

    if is_empty(form.bg_file_id.data) is False:
        bg_path = get_upload_file_path(form.bg_file_id.data)
        if os.path.exists(bg_path):
            bg_path = form.background.data

    result = tools.replace(img_path=img_path, background=bg_path, save_dir=get_output_dir())

    return jsonify({
        "fileId": get_file_id(result, True),
        "url": file_url(result, True)
    })


@app.route("/image/resize", methods=["POST", 'GET'])
def resize():
    # flip: 翻转， 0 为沿X轴翻转，正数为沿Y轴翻转，负数为同时沿X轴和Y轴翻转
    # reset: 是否重头开始，否则从上一次的处理开始,默认为重头开始
    # resize: 重新设定尺寸
    # rect: 裁剪，left,top,right,bottom 四个参数
    reset = request.values.get("reset")
    file_id = request.values.get("fileId")
    width = request.values.get("width")
    height = request.values.get("height")
    rotate = request.values.get("rotate")
    flip = request.values.get("flip")
    rect = request.values.get('rect')

    if is_empty(reset):
        reset = True
    else:
        reset = False

    file_path = get_upload_file_path(file_id)

    is_get = request.method.upper() == "GET"

    if os.path.exists(file_path) is not True:
        return "文件上传失败，请重新上传！" if is_get else jsonify(error_resp("文件上传失败，请重新上传")), 400

    out_file = get_output_file_path(file_id, "resize")
    if reset is False and os.path.exists(out_file):
        img = cv2.imread(out_file)
    else:
        img = cv2.imread(file_path)
    origin_h = img.shape[0]
    origin_w = img.shape[1]

    img = crop(img, rect, origin_w, origin_h)

    change_size = True
    if is_empty(width) and is_empty(height):
        change_size = False
        w = origin_w
        h = origin_h
    elif is_empty(width):
        h = int(height)
        w = round(h * origin_w / origin_h)
    elif is_empty(height):
        w = int(width)
        h = round(w * origin_h / origin_w)
    else:
        w = int(width)
        h = int(height)

    dst = img

    if change_size:
        dst = cv2.resize(img, (w, h))

    if is_empty(flip) is not True:
        flip = int(flip)
        dst = cv2.flip(dst, flip)

    if is_empty(rotate) is False:
        dst = rot_degree(dst, float(rotate), w=w, h=h)

    if dst is not None:
        cv2.imwrite(out_file, dst)
        if is_get:
            return send_file(out_file, mimetype="image/png")
        else:
            return jsonify(success_resp({
                "fileId": get_file_id(out_file),
                "url": file_url(out_file)
            }))

    if is_get:
        return send_file(file_path, mimetype="image/png")
    return jsonify(success_resp({
        "fileId": get_file_id(file_path),
        "url": file_url(file_path)
    }))


def crop(img, rect: str, w, h):
    # 裁剪
    if is_empty(rect):
        return img
    r = rect.split(',')
    if len(r) != 4:
        return img

    left = int(r[0])
    top = int(r[1])
    right = int(r[2])
    bottom = int(r[3])

    if left < 0:
        left = 0
    if right > w:
        right = w
    if top < 0:
        top = 0
    if bottom > h:
        bottom = h

    if left == 0 and top == 0 and right == w and bottom == h:
        return img

    return img[top:bottom, left:right]


def rot_degree(img, degree, w, h):
    center = (w / 2, h / 2)

    M = cv2.getRotationMatrix2D(center, degree, 1)
    top_right = np.array((w - 1, 0)) - np.array(center)
    bottom_right = np.array((w - 1, h - 1)) - np.array(center)
    top_right_after_rot = M[0:2, 0:2].dot(top_right)
    bottom_right_after_rot = M[0:2, 0:2].dot(bottom_right)
    new_width = max(int(abs(bottom_right_after_rot[0] * 2) + 0.5), int(abs(top_right_after_rot[0] * 2) + 0.5))
    new_height = max(int(abs(top_right_after_rot[1] * 2) + 0.5), int(abs(bottom_right_after_rot[1] * 2) + 0.5))
    offset_x = (new_width - w) / 2
    offset_y = (new_height - h) / 2
    M[0, 2] += offset_x
    M[1, 2] += offset_y
    dst = cv2.warpAffine(img, M, (new_width, new_height))
    return dst