'''
Author: liuhy
email: liuhy6@sugon.com
Date: 2023-03-03 10:17:07
LastEditTime: 2023-03-03 11:28:33
FilePath: \lpr\cut_ccpd.py
'''
# -*- coding: utf-8 -*-
# @Author: liuhy
# @Email: 17603873430@163.com
# @Date:   2023-02-28 15:14:13
# @Last Modified time: 2023-03-01 17:52:20
# coding:utf-8
import os
import cv2
import time
import argparse

provinces = ['皖', '沪', '津', '渝', '冀', '晋', '蒙', '辽', '吉', '黑', '苏', '浙', '京', '闽', '赣', '鲁', '豫', '鄂', '湘', '粤', '桂', '琼', '川', '贵', '云', '藏', '陕', '甘', '青', '宁', '新', '警', '学', 'O']
alphabets = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W','X', 'Y', 'Z', 'O']
ads = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X','Y', 'Z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'O']

def cut_img(img_path, save_path):
    images = os.listdir(img_path)
    for i, image in enumerate(images):
        try:
            print('\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\bProcessing:%s'%i, end='',flush=True)
            img = cv2.imread(os.path.join('ccpd_base', image))
            info = image.split('-')
            r1, r2, l1, l2 = info[3].split('_')
            rx1, ry1 = r1.split('&')
            lx1, ly1 = l1.split('&')
            im = img[int(ly1):int(ry1), int(lx1):int(rx1)]

            label_index = info[4].split('_')
            label = provinces[int(label_index[0])]
            label += alphabets[int(label_index[1])]

            if label_index[2].isalnum():
                label += ads[int(label_index[2])]
            else:
                label += alphabets[int(label_index[2])]
            label += ads[int(label_index[3])]
            label += ads[int(label_index[4])]
            label += ads[int(label_index[5])]
            label += ads[int(label_index[6])]

            im = cv2.resize(im, (94, 24))
            if os.path.exists(save_path) is False:
                os.mkdir(save_path)
            cv2.imencode('.jpg', im)[1].tofile(save_path + '/' + label + '.jpg')
        except Exception as e:
            print('wrong image:', image)

if __name__=='__main__':
    parser = argparse.ArgumentParser(description='parameters to vaildate net')
    parser.add_argument('--ccpdpath', default='CCPD/ccpd_base', help='model path to vaildate')
    parser.add_argument('--savepath', default='CCPD/lpr', help='the image path')
    args = parser.parse_args()
    cut_img(args.ccpdpath, args.savepath)