#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Date    : 2019-06-03 10:12:06
# @Author  : XuYG (xiaoxu04@foxmail.com)
# @Link    : http://www.itecfun.com
# @Version : $V1.0$

import os
from PIL import Image
import pytesseract

class xygORC(object):
	def __init__(self,dir_path,lang,file_name,write_mode):
		'''
		初始化相关参数
		
		Arguments:
			dir_path {[str]} -- [要识别的图片所在目录]
			lang {[str]} -- [tessdata训练数据包的名称，默认采用简体中文识别 chi_sim]
			file_name {[str]} -- [识别的文本保存的文件名称，可自由命名]
			write_mode {[str]} -- [default a+]
		'''
		self.dir_path = dir_path
		self.lang = lang
		self.file_name = file_name
		self.write_mode = write_mode
		self.content = ""

	def recognize(self,image_path,lang):
		'''
		单张图片识别	
		Arguments:
			image_path {[str]} -- [待识别图片路径]
			lang {[str]} -- [tessdata 训练数据包的名称，简体中文用chi_sim进行识别]
		Returns:
			[str] -- [返回识别的本张图片的文字内容]
		'''
		image = Image.open(image_path)
		#使用tessdata chi_sim简体中文训练数据进行识别
		content = pytesseract.image_to_string(image, lang=lang)
		return content
	def save(self,file_name,content,write_mode):
		'''
		保存识别的内容
		Arguments:
			file_name {[str]} -- [保存的文件路径]
			content {[str]} -- [文本内容]
			write_mode {[str]} -- [a,b,r,a+,b+...]
		Returns:
			bool -- [返回保存状态，True or False]
		'''
		#结果保存
		with open(file_name,write_mode) as f:
			f.write(content.replace(' ',''))  #识别后每个字之间有空格，需要去掉空格
			return True
		return False

	def recognize_all(self,dir_path,lang):
		'''
		识别制定目录下所有图片中的文字
		
		Arguments:
			dir_path {[str]} -- [目录路径]
			lang {[str]} -- [tessdata 训练数据包的名称，简体中文用chi_sim进行识别]
		
		Returns:
			[str] -- [目录下所有图片识别后的文字内容]
		'''
		img_list = self.get_file_list(dir_path, [])
		img_list.sort(key=lambda x: os.path.getmtime(x))
		content = ""
		for img in img_list:
			text = self.recognize(img, lang)
			content +=text +"\n "
			print("'{0}'已识别完成！".format(img))
		return content

	def get_file_list(self,dir, file_list):
		'''
		获取指定目录下所有待识别的jpg、png图片列表
		并按时间排序后返回
		
		Arguments:
			dir {[str]} -- [图片存放目录]
			file_list {[list]} -- [图片列表]
		
		Returns:
			[list] -- [目录下所有图片列表]
		'''
		newDir = dir
		if os.path.isfile(dir):
			ext = os.path.splitext(dir)[-1]
			if ext == ".jpg" or ext ==".png":
				file_list.append(dir)
		elif os.path.isdir(dir):
			for file in os.listdir(dir):
				newDir = os.path.join(dir, file)
				self.get_file_list(newDir, file_list)
		return file_list

	def main(self):
		'''
		主函数入口
		'''
		self.content = self.recognize_all(self.dir_path, self.lang)
		self.save(self.file_name, self.content, write_mode)


if __name__ == '__main__':
	dir_path ='E:\\国家政策\\政策学习01：国家大数据战略与社会治理革命' 
	lang  ='chi_sim'  #tessdata chi_sim训练数据包的名称
	file_name = "政策学习.txt" #要保存的识别文字的文件名称，可自由命名
	write_mode ="a+" #以文件末尾追加的形式进行写入
	os.chdir(dir_path) ##将工作目录切换到图片所在目录，主要目的是将文件保存在图片所在目录，方便查阅

	myorc = xygORC(dir_path,lang,file_name,write_mode)

	myorc.main()