文字檢測模型EAST應用詳解 ckpt pb的tf載入，opencv載入

geek_codinggirl發表於2020-04-24
模型ASTOpenCV

參考連結：https://github.com/argman/EAST （專案來源）

                  https://github.com/opencv/opencv/issues/12491  （遇到的問題）

　　　　  https://www.pyimagesearch.com/2018/08/20/opencv-text-detection-east-text-detector/   （opencv載入）

 

文字檢測有很多比較好的現成的模型比如yolov3，pesnet，pennet，east。不一一贅述，講一下自己跑通east的過程。

在https://github.com/argman/EAST連結中下載專案，windows下，各種包的版本要正確否則會出一些亂七八糟的錯誤。

執行EAST/eval.py。沒有什麼特別的問題要說，我在cpu下單張640*480的圖能夠達到每張0.4秒左右，還是非常優秀的。中英文數字都可。

 

但是原始碼是ckpt，非常大,轉成pb會稍微小點。新增：

##生成pb模型，但需要修改model.py
output_graph_def = tf.graph_util.convert_variables_to_constants(self.sess,  # The session is used to retrieve the weights
    tf.get_default_graph().as_graph_def(),  # The graph_def is used to retrieve the nodes
    ["feature_fusion/Conv_7/Sigmoid", "feature_fusion/concat_3"]
)
output_graph='D:\\work\\video\\hand_tracking_no_op\\hand_tracking\\EAST\\east_icdar2015_resnet_v1_50_rbox\\out.pb'
with tf.gfile.GFile(output_graph, "wb") as f:
    f.write(output_graph_def.SerializeToString())
print("%d ops in the final graph." % len(output_graph_def.node))

位置在eval.py中的

saver.restore(self.sess, model_path)後面。注意如果你想要opencv載入pb還要修改model.py中的內容，這個在後面一篇文章中會講到。
生成後用tf載入，方法跟載入ckpt相似：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list

try:
    os.makedirs(FLAGS.output_dir)
except OSError as e:
    if e.errno != 17:
        raise

print("load_graph")
graph = load_graph(FLAGS.checkpoint_path)

input_images = graph.get_tensor_by_name(
    'import/input_images:0')

f_score = graph.get_tensor_by_name('import/feature_fusion/Conv_7/Sigmoid:0')
f_geometry = graph.get_tensor_by_name(
    'import/feature_fusion/concat_3:0')

with tf.Session(graph=graph) as sess:

    im_fn_list = get_images()
    for im_fn in im_fn_list:
        im = cv2.imread(im_fn)[:, :, ::-1]
        start_time = time.time()
        im_resized, (ratio_h, ratio_w) = resize_image(im)

        timer = {'net': 0, 'restore': 0, 'nms': 0}
        start = time.time()

        #file_writer = tf.summary.FileWriter('tmp/log', sess.graph)

        score, geometry = sess.run([f_score, f_geometry], feed_dict={
                                    input_images: [im_resized]})
        timer['net'] = time.time() - start

        boxes, timer = detect(score_map=score, geo_map=geometry, timer=timer)
        print('{} : net {:.0f}ms, restore {:.0f}ms, nms {:.0f}ms'.format(
            im_fn, timer['net']*1000, timer['restore']*1000, timer['nms']*1000))

        if boxes is not None:
            boxes = boxes[:, :8].reshape((-1, 4, 2))
            boxes[:, :, 0] /= ratio_w
            boxes[:, :, 1] /= ratio_h

        duration = time.time() - start_time
        print('[timing] {}'.format(duration))

        # save to file
        if boxes is not None:
            res_file = os.path.join(
                FLAGS.output_dir,
                '{}.txt'.format(
                    os.path.basename(im_fn).split('.')[0]))

            with open(res_file, 'w') as f:
                for box in boxes:
                    # to avoid submitting errors
                    box = sort_poly(box.astype(np.int32))
                    if np.linalg.norm(box[0] - box[1]) < 5 or np.linalg.norm(box[3]-box[0]) < 5:
                        continue
                    f.write('{},{},{},{},{},{},{},{}\r\n'.format(
                        box[0, 0], box[0, 1], box[1, 0], box[1, 1], box[2, 0], box[2, 1], box[3, 0], box[3, 1],
                    ))
                    cv2.polylines(im[:, :, ::-1], [box.astype(np.int32).reshape((-1, 1, 2))], True, color=(255, 255, 0), thickness=1)
        if not FLAGS.no_write_images:
            img_path = os.path.join(FLAGS.output_dir, os.path.basename(im_fn))
            cv2.imwrite(img_path, im[:, :, ::-1])

以上就是EAST的ckpt轉pb用tf載入啦。
下一篇講opencv載入east的pb。



相關文章