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, get_file_id_path
from utils import color
import time

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("参数错误！"))

    img_path = get_file_id_path(form.file_id.data)

    exist, alpha_path, path = tools.has_seg(img_path, get_output_dir())

    if exist:
        alpha = cv2.imread(alpha_path)
        fg = cv2.imread(path)
    else:
        alpha, fg, _, path = tools.seg(img_path, get_output_dir())

    if is_empty(form.bg_file_id.data) and is_empty(form.background.data):
        return jsonify(success_resp({
            "fileId": form.file_id.data,
            "url": file_url(path, True)
        }))

    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) is False:
            bg_path = form.background.data

    bg = color.get_bg(bg_path, fg.shape)
    if bg is None:
        return jsonify(success_resp({
            "fileId": form.file_id.data,
            "url": file_url(path, True)
        }))

    alpha = alpha / 255.0
    alpha = alpha[:, :, np.newaxis]
    com = alpha * fg + (1 - alpha) * bg
    com = com.astype('uint8')
    filename = os.path.basename(img_path)
    names = os.path.splitext(filename)
    save_name = "{}_{}{}".format(names[0], time.time(), names[1])
    com_save_path = os.path.join(get_output_dir(), save_name)
    cv2.imwrite(com_save_path, com)

    return jsonify(success_resp({
        "fileId": get_file_id(com_save_path, True),
        "url": file_url(com_save_path, 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')
    save = request.values.get('onSave')
    download = request.values.get('download')
    save = is_empty(save) is False
    download = is_empty(download) is False

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

    file_path = get_file_id_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]

    dst, r_w, r_h = crop(img, rect, origin_w, origin_h)

    change_size = True
    if is_empty(width) and is_empty(height):
        change_size = False

    if change_size:
        if 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)
        origin = img if dst is None else dst
        dst = cv2.resize(origin, (w, h))
        r_w = w
        r_h = h

    if r_w is None:
        r_w = origin_w
    if r_h is None:
        r_h = origin_h

    if is_empty(flip) is not True:
        flip = int(flip)
        origin = img if dst is None else dst
        dst = cv2.flip(origin, flip)

    if is_empty(rotate) is False and int(rotate) != 0:
        origin = img if dst is None else dst
        dst, new_w, new_h = rot_degree(origin, float(rotate), w=r_w, h=r_h)
        if new_w is not None:
            r_w = new_w
            r_h = new_h

    result_path = file_path

    if dst is not None:
        cv2.imwrite(out_file, dst)
        result_path = out_file
    else:
        dst = img

    if save:
        result_path = get_output_file_path(file_id)
        cv2.imwrite(result_path, dst)

    if is_get:
        mimetype = 'application/octet-stream' if download else 'image/png'
        return send_file(result_path, mimetype=mimetype)

    return jsonify(success_resp({
        "fileId": get_file_id(result_path, True),
        "url": file_url(result_path, True),
        "width": r_w,
        "height": r_h,
    }))


def crop(img, rect: str, w, h):
    # 裁剪,
    # 如果没有改变，返回None
    if is_empty(rect):
        return None, None, None
    r = rect.split(',')
    if len(r) != 4:
        return None, None, None

    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 None, None, None

    return img[top:bottom, left:right], right-left, bottom-top


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, new_width, new_height