ocr 文字区域检测及识别

ocr 文字区域检测及识别

# coding=utf-8

from PIL import Image, ImageFilter, ImageEnhance
from skimage.filters import threshold_otsu
import skimage.morphology as sm
from skimage.measure import regionprops
import matplotlib.pyplot as plt
import numpy as np
import pytesseract
import re
import os
import time
import logging
logging.basicConfig(level=logging.INFO, format="%(message)s", filename=\'train_output.log\')

class ShopCert(object):
    def cut_region(self, img):
        """
        先按规则缩小搜索范围
        """
        w, h = img.size
        if h<1200:
            factor = max(1, 1600.0/h)
            newsize = int(w*factor), int(h*factor)
            img = img.resize(newsize, Image.ANTIALIAS)
        if w<h:
            box = (w*0.4, h*0.18, w*0.96, h*0.6)
        else:
            box = (w*0.1, h*0.18, w*0.96, h*0.9)
        return img.crop(box)

    
    def detect_text(self, img):
        """
        检测字符区域
        """
        imgM = np.array(img.convert(\'L\'))
        imgM = 1 * (imgM < threshold_otsu(imgM)) 
        imgM = sm.binary_closing(imgM, np.ones((5, 10)))
        imgM = sm.remove_small_objects(imgM, 600)
        label_img = sm.label(imgM)
        imgList = []
        for region in regionprops(label_img):
            minr, minc, maxr, maxc = region.bbox
            w, h = (maxc-minc), (maxr-minr)
            if h > w * 0.2:
                continue
            box = minc-5, minr-3, maxc+5, maxr+3
            imgList.append(img.crop(box))
        return imgList

    
    
    def clear_noise(self, box):
        """
        降噪处理
        """
        box = box.convert(\'L\')
 #       box = box.point(lambda x: 0 if x<50 else x)
        box = box.point(lambda x: 200 if x>200 else x)
        box = ImageEnhance.Contrast(box).enhance(2.5)
        return box
    
    def predict(self, fname, lang=\'eng\'):
        """
        ocr 识别
        """
        img = Image.open(fname)
        # 先大致缩小范围
        region = self.cut_region(img)
        # 候选字符区域
#        region = self.clear_noise(region)
        boxList = self.detect_text(region)
        # 遍历识别
        for box in boxList:
            box = self.clear_noise(box)
            w, h = box.size
            if float(w)/h > 12.5:
                res = pytesseract.image_to_string(box, lang=\'chi_sim\', config=\'-psm 7\')
            else:
                res = pytesseract.image_to_string(box, lang=\'eng\', config=\'-psm 7\')
            res = re.sub(\'\s\', \'\', res)  # 去除中间空白
            res = re.findall(r\'[0-9][A-Z0-9]{13,20}\', res)  # 13-20位
            for line in res:
                line = line.strip()
                if line.find(u\'年\')>1:
                    continue
                print \'line\', line
                if len(line)> 14:
                    box.save(\'img/clearNoise/%s_%s.jpg\' % (fname.split(\'/\')[-1].split(\'.\')[0], line))
                    return line
                else:
                    print \'error line\', line
        return \'error\'


def show_pic(path=\'img/origin2/\'):
    fnames = [os.path.join(path, fname) for fname in os.listdir(path)]
    for i, fname in enumerate(fnames, 0):
        print fname
        img = Image.open(fname)
      #  img.save(\'./tesseract-train/cert.normal.exp%d.ttf\' % i)
        img = ImageEnhance.Contrast(img).enhance(2.0)
        img = img.filter(ImageFilter.MedianFilter).convert(\'L\')
        plt.figure(figsize=(10, 12), dpi=300)
        plt.imshow(img, plt.cm.gray)
        plt.title(fname.split(\'/\')[-1]+\'_%d\' % i)
        plt.show()if __name__ == \'__main__\':
    test = ShopCert()
    path = \'img/origin2/\'
    fnames = [os.path.join(path, fname) for fname in os.listdir(path) if fname.endswith(\'jpg\')]
    fnames.sort()
    
    arguments = \'mode: L; enhance:2.0; h:0.5; dh:0.15\'
    logging.info(\'%s\' % arguments)
    logging.info("%s: %s" % (\'imgname\', \'result\'))
    start_time = time.time()
    cnt = 0
    for idx, fname in enumerate(fnames, 1):
        print idx, fname
        y_true = fname.split(\'/\')[-1].split(\'_\')[0]
        y_pred = test.predict(fname)
        if y_true == y_pred:
            cnt +=1
            print fname
        else: 
            print \'***\'*20
            print \'error\'
        logging.info("%s: %s" % (fname, y_pred))
        print \'y_true\', y_true
        print \'y_pred\', y_pred
        acc =  float(cnt)/idx
        print acc, cnt
        print \'==\'*20, idx
        logging.info(\'%.3f %d/%d\' % (acc, cnt, idx))
    print \'cost time: \', time.time()-start_time
    logging.info(\'accuracy: %.2f\' % acc)