# -*- coding: utf-8 -*-
# @Author        : Lyu Kui
# @Email         : 9428.al@gmail.com
# @Create Date   : 2022-07-22 13:10:47
# @Last Modified : 2022-09-08 19:03:24
# @Description   : 

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

import cv2
# from turnsole.ocr_engine import angle_detector
from turnsole.ocr_engine import object_detector
import matplotlib.pyplot as plt


if __name__ == "__main__":

    base_dir = '/home/lk/MyProject/BMW/数据集/文件分类/身份证'

    for (rootDir, dirNames, filenames) in os.walk(base_dir):

        for filename in filenames:
        
            if not filename.endswith('.jpg'):
                continue

            img_path = os.path.join(rootDir, filename)
            print(img_path)

            image = cv2.imread(img_path)
            
            results = object_detector.process(image)

            print(results)

            for item in results:
                xmin = item['location']['xmin']
                ymin = item['location']['ymin']
                xmax = item['location']['xmax']
                ymax = item['location']['ymax']
                cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)

            plt.imshow(image[...,::-1])
            plt.show()