<commit_before>from __future__ import print_function

import argparse
import logging
import os

import mxnet as mx
import numpy as np
from PIL import Image

from crnn import crnn


class SimpleBatch(object):
    def __init__(self, data_names, data, label_names, label):
        self.data = data
        self.label = label
        self.data_names = data_names
        self.label_names = label_names

        self.pad = 0
        self.index = None

    @property
    def provide_data(self):
        return [(n, x.shape) for n, x in zip(self.data_names, self.data)]

    @property
    def provide_label(self):
        return [(n, x.shape) for n, x in zip(self.label_names, self.label)]


class OCRIter(mx.io.DataIter):
    def __init__(self, batch_size, classes, data_shape, num_label, init_states,dataset_lst):
        super(OCRIter, self).__init__()
        self.batch_size = batch_size
        self.data_shape = data_shape
        self.num_label = num_label
        self.init_states = init_states
        self.init_state_arrays = [mx.nd.zeros(x[1]) for x in init_states]
        self.classes = classes
        self.dataset_lst_file = open(dataset_lst)
        self.provide_data = [('data', (batch_size, 1, data_shape[1], data_shape[0]))] + init_states
        self.provide_label = [('label', (self.batch_size, num_label))]

    def __iter__(self):
        init_state_names = [x[0] for x in self.init_states]
        data = []
        label = []
        cnt = 0
        for m_line in self.dataset_lst_file:
            img_path,img_label = m_line.strip().split('\t')
            cnt += 1
            img = Image.open(img_path).convert('L').resize(self.data_shape)
            img = np.array(img).reshape((1, data_shape[1], data_shape[0]))
            data.append(img)
            plate_str = img_label
            ret = np.zeros(self.num_label, int)
            for number in range(len(plate_str)):
                ret[number] = self.classes.index(plate_str[number]) + 1
            label.append(ret)
            if cnt % self.batch_size == 0:
                data_all = [mx.nd.array(data)] + self.init_state_arrays
                label_all = [mx.nd.array(label)]
                data_names = ['data'] + init_state_names
                label_names = ['label']
                data.clear()
                label.clear()
                yield SimpleBatch(data_names, data_all, label_names, label_all)
                continue

    def reset(self):
        if self.dataset_lst_file.seekable():
            self.dataset_lst_file.seek(0)


def ctc_label(p):
    ret = []
    p1 = [0] + p
    for i in range(len(p)):
        c1 = p1[i]
        c2 = p1[i + 1]
        if c2 == 0 or c2 == c1:
            continue
        ret.append(c2)
    return ret


def remove_blank(l):
    ret = []
    for i in range(len(l)):
        if l[i] == 0:
            break

        ret.append(l[i])
    return ret


def Accuracy(label, pred):
    global BATCH_SIZE
    global SEQ_LENGTH
    hit = 0.
    total = 0.
    for i in range(BATCH_SIZE):
        l = remove_blank(label[i])
        p = []
        for k in range(SEQ_LENGTH):
            p.append(np.argmax(pred[k * BATCH_SIZE + i]))
        p = ctc_label(p)
        if len(p) == len(l):
            match = True
            for k in range(len(p)):
                if p[k] != int(l[k]):
                    match = False
                    break
            if match:
                hit += 1.0
        total += 1.0
    return hit / total


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--name',required=True,help='the model name')
    parser.add_argument('--charset',required=True,help='the charset file')
    parser.add_argument('--train_lst',required=True,help='the csv which contains all train list')
    parser.add_argument('--val_lst',required=True,help='the csv which contains all val list')
    parser.add_argument('--batch_size',required=True,type=int,default=64,help='train/val batch size,default is 64')
    parser.add_argument('--seq_len',required=True,type=int,default=17,help='the sequence length effected by image width')
    parser.add_argument('--num_label',required=True,type=int,default=9,help='output label length,must less than seq_len')
    parser.add_argument('--imgH',required=True,type=int,default=32,help='image height,must divided by 16')
    parser.add_argument('--imgW',required=True,type=int,default=200,help='image width')
    parser.add_argument('--num_hidden',type=int,default=256,help='num of parameters in lstm hidden layer')
    parser.add_argument('--num_lstm',type=int,default=2,help='num of lstm layers')
    parser.add_argument('--gpu',type=bool,action='store_true',help='enable train with gpu(0)')
    parser.add_argument('--from_epoch',type=int,action='store_true',help='continue train from specific epoch file')
    opt = parser.parse_args()
    model_name = opt.name
    log_file_name = model_name+'.log'
    log_file = open(log_file_name, 'w')
    log_file.close()
    logging.basicConfig()
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    fh = logging.FileHandler(log_file_name)
    logger.addHandler(fh)
    model_dir_path = os.path.join(os.getcwd(), 'model')
    if not os.path.exists(model_dir_path):
        os.mkdir(model_dir_path)<commit_msg>fix train script argument parser bug.<commit_after>from __future__ import print_function

import argparse
import logging
import os

import mxnet as mx
import numpy as np
from PIL import Image

from crnn import crnn


class SimpleBatch(object):
    def __init__(self, data_names, data, label_names, label):
        self.data = data
        self.label = label
        self.data_names = data_names
        self.label_names = label_names

        self.pad = 0
        self.index = None

    @property
    def provide_data(self):
        return [(n, x.shape) for n, x in zip(self.data_names, self.data)]

    @property
    def provide_label(self):
        return [(n, x.shape) for n, x in zip(self.label_names, self.label)]


class OCRIter(mx.io.DataIter):
    def __init__(self, batch_size, classes, data_shape, num_label, init_states,dataset_lst):
        super(OCRIter, self).__init__()
        self.batch_size = batch_size
        self.data_shape = data_shape
        self.num_label = num_label
        self.init_states = init_states
        self.init_state_arrays = [mx.nd.zeros(x[1]) for x in init_states]
        self.classes = classes
        self.dataset_lst_file = open(dataset_lst)
        self.provide_data = [('data', (batch_size, 1, data_shape[1], data_shape[0]))] + init_states
        self.provide_label = [('label', (self.batch_size, num_label))]

    def __iter__(self):
        init_state_names = [x[0] for x in self.init_states]
        data = []
        label = []
        cnt = 0
        for m_line in self.dataset_lst_file:
            img_path,img_label = m_line.strip().split('\t')
            cnt += 1
            img = Image.open(img_path).convert('L').resize(self.data_shape)
            img = np.array(img).reshape((1, data_shape[1], data_shape[0]))
            data.append(img)
            plate_str = img_label
            ret = np.zeros(self.num_label, int)
            for number in range(len(plate_str)):
                ret[number] = self.classes.index(plate_str[number]) + 1
            label.append(ret)
            if cnt % self.batch_size == 0:
                data_all = [mx.nd.array(data)] + self.init_state_arrays
                label_all = [mx.nd.array(label)]
                data_names = ['data'] + init_state_names
                label_names = ['label']
                data.clear()
                label.clear()
                yield SimpleBatch(data_names, data_all, label_names, label_all)
                continue

    def reset(self):
        if self.dataset_lst_file.seekable():
            self.dataset_lst_file.seek(0)


def ctc_label(p):
    ret = []
    p1 = [0] + p
    for i in range(len(p)):
        c1 = p1[i]
        c2 = p1[i + 1]
        if c2 == 0 or c2 == c1:
            continue
        ret.append(c2)
    return ret


def remove_blank(l):
    ret = []
    for i in range(len(l)):
        if l[i] == 0:
            break

        ret.append(l[i])
    return ret


def Accuracy(label, pred):
    global BATCH_SIZE
    global SEQ_LENGTH
    hit = 0.
    total = 0.
    for i in range(BATCH_SIZE):
        l = remove_blank(label[i])
        p = []
        for k in range(SEQ_LENGTH):
            p.append(np.argmax(pred[k * BATCH_SIZE + i]))
        p = ctc_label(p)
        if len(p) == len(l):
            match = True
            for k in range(len(p)):
                if p[k] != int(l[k]):
                    match = False
                    break
            if match:
                hit += 1.0
        total += 1.0
    return hit / total


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--name',required=True,help='the model name')
    parser.add_argument('--charset',required=True,help='the charset file')
    parser.add_argument('--train_lst',required=True,help='the csv which contains all train list')
    parser.add_argument('--val_lst',required=True,help='the csv which contains all val list')
    parser.add_argument('--batch_size',type=int,default=64,help='train/val batch size,default is 64')
    parser.add_argument('--seq_len',type=int,default=17,help='the sequence length effected by image width')
    parser.add_argument('--num_label',type=int,default=9,help='output label length,must less than seq_len')
    parser.add_argument('--imgH',type=int,default=32,help='image height,must divided by 16')
    parser.add_argument('--imgW',type=int,default=200,help='image width')
    parser.add_argument('--num_hidden',type=int,default=256,help='num of parameters in lstm hidden layer')
    parser.add_argument('--num_lstm',type=int,default=2,help='num of lstm layers')
    parser.add_argument('--gpu',action='store_true',help='enable train with gpu(0)')
    parser.add_argument('--from_epoch',type=int,help='continue train from specific epoch file')
    opt = parser.parse_args()
    model_name = opt.name
    log_file_name = model_name+'.log'
    log_file = open(log_file_name, 'w')
    log_file.close()
    logging.basicConfig()
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    fh = logging.FileHandler(log_file_name)
    logger.addHandler(fh)
    model_dir_path = os.path.join(os.getcwd(), 'model')
    if not os.path.exists(model_dir_path):
        os.mkdir(model_dir_path)
{'fraction_non_alphanumeric': 0.12507133346014837, 'fraction_numerical': 0.009891573140574473, 'mean_word_length': 2.5392123864018847, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '993686', 'n_tokens_mistral': 3488, 'n_tokens_neox': 3360, 'n_words': 728}