FORMAT_CONVERSION_MAXCOMPUTE/isr/isr_interface.py

import base64
import json
import os
import time
import sys
import traceback
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
from format_convert import _global
import cv2
import numpy as np
from PIL import Image
from format_convert.utils import log, get_md5_from_bytes, request_post, np2pil, bytes2np, pil2np
from isr.post_process import get_seal_part, replace_seal_part
from isr.model import get_tiny_inference_model, seal_model, seal_model_se
from isr.pre_process import count_red_pixel, get_anchors, get_classes, get_colors
from isr.utils import get_best_predict_size, pil_resize, letterbox_image, draw_boxes, adjust_boxes
import tensorflow as tf
from flask import Flask, request

tf.compat.v1.disable_eager_execution()
sess1 = tf.compat.v1.Session(graph=tf.Graph())
sess2 = tf.compat.v1.Session(graph=tf.Graph())


def remove_seal(image_np, model):
    # inference data
    image_seal = image_np
    h, w = image_seal.shape[:2]
    best_h, best_w = get_best_predict_size(image_seal)
    X = np.zeros((1, best_h, best_w, 3))

    # resize
    image_seal = pil_resize(image_seal, best_h, best_w)
    # cv2.imshow("resize", image_seal)
    X[0] = image_seal / 255

    # predict
    with sess2.as_default():
        with sess2.graph.as_default():
            pred = model.predict(X)
            pred = pred[0]*255.
            pred = pred.astype(np.uint8)
            pred = pil_resize(pred, h, w)
    # cv2.imshow("pred", pred)
    # cv2.waitKey(0)
    return pred


def detect_seal(image_np, model):
    image_pil = np2pil(image_np)

    # 首先判断红色像素
    # if not count_red_pixel(image_np):
    #     return image_np, [], []

    # create image input
    h, w = image_np.shape[:2]
    # best_h, best_w = get_best_predict_size(image_np, times=32, max_size=1280)
    best_h, best_w = 1024, 1024
    image_resize = letterbox_image(image_pil, tuple(reversed([best_h, best_w])))
    # cv2.imshow("letterbox_image", pil2np(image_resize))
    # cv2.waitKey(0)
    # image_resize = pil_resize(image_np, best_h, best_w)

    # image_resize = image_pil.resize((int(416), int(416)), Image.BICUBIC)
    image_resize = np.array(image_resize, dtype='float32')
    image_resize = image_resize.astype('float32') / 255.
    image_resize = np.expand_dims(image_resize, 0)

    # create image shape input
    image_shape = np.array([image_pil.size[1], image_pil.size[0]])
    image_shape = np.expand_dims(image_shape, 0)

    # inference data
    with sess1.as_default():
        with sess1.graph.as_default():
            try:
                out_boxes, out_scores, out_classes = model.predict([image_resize, image_shape],steps=1)
            except ValueError as e:
                if '(0, 4)' in traceback.format_exc():
                    log("there is no seal!")
                    out_boxes, out_scores, out_classes = [], [], []
                else:
                    traceback.print_exc()
                    raise e

    if not isinstance(out_boxes, list) and not isinstance(out_classes, list):
        out_boxes = out_boxes.astype(np.int32)
        out_classes = out_classes.astype(np.int32)

        boxes = adjust_boxes(image_pil, out_boxes)
    else:
        boxes = out_boxes

    # # draw
    # class_names = get_classes(os.path.abspath(os.path.dirname(__file__))+"/yolo_data/my_classes.txt")
    # colors = get_colors(len(class_names))
    # image_draw = draw_boxes(image_pil, out_boxes, out_classes, out_scores, class_names, colors)
    # image_draw = cv2.cvtColor(np.array(image_draw), cv2.COLOR_RGB2BGR)
    # cv2.namedWindow('detect', cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)
    # cv2.imshow("detect", image_draw)
    # cv2.waitKey(0)
    return image_np, boxes, out_classes


def isr(data, isr_yolo_model, isr_model):
    log("into isr_interface isr")
    try:
        img_data = base64.b64decode(data)
        img_np = bytes2np(img_data)
        _img, boxes, classes = detect_seal(img_np, isr_yolo_model)
        if not boxes and not classes:
            return {"image": img_np}
        part_list = get_seal_part(_img, boxes, classes)
        new_part_list = []
        for part in part_list:
            part_remove = remove_seal(part, isr_model)
            new_part_list.append(part_remove)
        img_replace = replace_seal_part(img_np, new_part_list, boxes)
        return {"image": img_replace}
    except TimeoutError:
        return {"image": [-5]}
    except:
        traceback.print_exc()
        return {"image": [-1]}


# 接口配置
app = Flask(__name__)


@app.route('/isr', methods=['POST'])
def _isr():
    _global._init()
    _global.update({"port": globals().get("port")})
    start_time = time.time()

    log("into isr_interface _isr")
    try:
        if not request.form:
            log("isr no data!")
            return json.dumps({"text": str([-9]), "bbox": str([-9])})
        data = request.form.get("data")
        log("isr_interface get data time" + str(time.time()-start_time))

        img_data = base64.b64decode(data)
        img_np = bytes2np(img_data)
        _md5 = request.form.get("md5")
        _global.update({"md5": _md5})

        # 初始化模型
        isr_yolo_model = globals().get("global_isr_yolo_model")
        isr_model = globals().get("global_isr_model")
        if isr_model is None or isr_yolo_model is None:
            print("=========== init isr model ===========")
            isr_yolo_model, isr_model = IsrModels().get_model()
            globals().update({"global_isr_yolo_model": isr_yolo_model})
            globals().update({"global_isr_model": isr_model})

        # 检测印章
        _img, boxes, classes = detect_seal(img_np, isr_yolo_model)
        # 检测不到，直接返回
        if not boxes and not classes:
            log("no seal detected!")
            return json.dumps({"image": [1]})
        else:
            log("there is " + str(len(boxes)) + " seals")

        # 截取
        part_list = get_seal_part(_img, boxes, classes)

        # 去除印章
        new_part_list = []
        for part in part_list:
            part_remove = remove_seal(part, isr_model)
            new_part_list.append(part_remove)

        # 替换
        img_replace = replace_seal_part(img_np, new_part_list, boxes)

        # numpy转为可序列化的string
        success, img_encode = cv2.imencode(".jpg", img_replace)
        # numpy -> bytes
        img_bytes = img_encode.tobytes()
        # bytes -> base64 bytes
        img_base64 = base64.b64encode(img_bytes)
        # base64 bytes -> string (utf-8)
        base64_string = img_base64.decode('utf-8')
        return json.dumps({"image": base64_string})
    except TimeoutError:
        return json.dumps({"image": [-5]})
    except:
        traceback.print_exc()
        return json.dumps({"image": [-1]})
    finally:
        log("isr interface finish time " + str(time.time()-start_time))


class IsrModels:
    def __init__(self):
        # python文件所在目录
        _dir = os.path.abspath(os.path.dirname(__file__))

        # detect
        model_path = _dir + "/models/seal_detect_yolo.h5"
        anchors = get_anchors(_dir + "/yolo_data/my_anchors.txt")
        class_names = get_classes(_dir + "/yolo_data/my_classes.txt")
        colors = get_colors(len(class_names))
        with sess1.as_default():
            with sess1.graph.as_default():
                self.isr_yolo_model = get_tiny_inference_model(anchors, len(class_names), weights_path=model_path)
                self.isr_yolo_model.load_weights(model_path)

        # remove
        model_path = _dir + "/models/seal_remove_unet.h5"
        with sess2.as_default():
            with sess2.graph.as_default():
                self.isr_model = seal_model_se(input_shape=(None, None, 3),
                                               output_shape=(None, None, 3))
                self.isr_model.load_weights(model_path)

    def get_model(self):
        return [self.isr_yolo_model, self.isr_model]


def test_isr_model(from_remote=False):
    file_path = "C:/Users/Administrator/Desktop/test_image/error10.jpg"
    with open(file_path, "rb") as f:
        file_bytes = f.read()
    file_base64 = base64.b64encode(file_bytes)
    _md5 = get_md5_from_bytes(file_bytes)[0]

    _global._init()
    _global.update({"port": 15010, "md5": _md5})

    if from_remote:
        file_json = {"data": file_base64, "md5": _md5}
        # _url = "http://192.168.2.102:17000/ocr"
        _url = "http://127.0.0.1:17000/ocr"
        print(json.loads(request_post(_url, file_json)))
    else:
        isr_yolo_model, isr_model = IsrModels().get_model()
        result = isr(file_base64, isr_yolo_model, isr_model)
        # print(result)
        if type(result.get("image")) == list:
            print(result)
        else:
            img = result.get("image")
            print(img.shape)
            cv2.namedWindow('img', cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)
            cv2.imshow("img", img)
            cv2.waitKey(0)
        # print(result)


if __name__ == "__main__":
    test_isr_model()