文本轮廓检测

from openvino.inference_engine import IECore
import numpy as np
import time
import cv2 as cv


def text_detection_demo():
    ie = IECore()
    for device in ie.available_devices:
        print(device)

    model_xml = "/home/bhc/BHC/model/intel/text-detection-0004/FP16/text-detection-0004.xml"
    model_bin = "/home/bhc/BHC/model/intel/text-detection-0004/FP16/text-detection-0004.bin"

    net = ie.read_network(model=model_xml, weights=model_bin)
    input_blob = next(iter(net.input_info))
    text_it = iter(net.outputs)
    out_blob1 = next(text_it)  # model/link_logits_/add              #输出(1, 16, 192, 320)
    out_blob2 = next(text_it)  # model/segm_logits/add               #（1, 2, 192, 320）2：text\no-text
    print(out_blob1, out_blob2)

    n, c, h, w = net.input_info[input_blob].input_data.shape
    print(n, c, h, w)

    src = cv.imread("002.png")
    image = cv.resize(src, (w, h))
    image = image.transpose(2, 0, 1)

    exec_net = ie.load_network(network=net, device_name="CPU")
    res = exec_net.infer(inputs={input_blob:[image]})

    res = res[out_blob2]                                            
    res = np.squeeze(res, 0)
    res = res.transpose(1, 2, 0)
    res = np.argmax(res, 2)

    hh, ww = res.shape
    mask = np.zeros((hh, ww), dtype=np.uint8)                                               #灰度图像
    mask[np.where(res > 0)] = 255                                                           #text则白色
    mask = cv.resize(mask, (src.shape[1], src.shape[0]))                            
    ret, binary = cv.threshold(mask, 127, 255, cv.THRESH_BINARY)                            #阈值二值化，大于127则255
    cv.imshow("mask", binary)
    # result = cv.addWeighted(src, 0.5, mask, 0.5, 0)

    contours, hiearchy = cv.findContours(binary, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)  #轮廓检测
    for cnt in range(len(contours)):
        x, y, w, h = cv.boundingRect(contours[cnt])                                         #轮廓矩形
        cv.rectangle(src, (x, y), (x+w, y+h), (244, 255, 0), 2, 8, 0)

    cv.imshow("Text Detection", src)
    cv.waitKey(0)


if __name__ == "__main__":
    text_detection_demo()