| |
| |
| |
| |
|
|
|
|
| import sys |
| import pathlib |
| import itertools |
| import re |
|
|
| from joblib import Parallel, delayed |
|
|
| import cv2 |
| import numpy as np |
| from PIL import Image |
| from torch.utils.data import IterableDataset, ChainDataset |
|
|
| import xml.etree.ElementTree as ET |
|
|
| pt = re.compile(r'(\{(.*)\})?(.*)') |
|
|
|
|
| def find_image_root(image_name, leaf_dir, root_dir): |
| image_root = leaf_dir.parent |
| while image_root >= root_dir: |
| image_paths = list(image_root.rglob(image_name)) |
| if len(image_paths) > 1: |
| print(ValueError(f'{image_name} is ambiguous\n{image_paths}'), sys.stderr) |
| break |
| elif len(image_paths) == 1: |
| return pathlib.Path(str(image_paths[0])[:-len(image_name)]) |
| image_root = image_root.parent |
| return None |
|
|
|
|
| class XMLRawDataset(IterableDataset): |
| IMAGE_TYPE_NONE = 0 |
| IMAGE_TYPE_BYTE = 1 |
| IMAGE_TYPE_ENCODED = 2 |
| IMAGE_TYPE_IMAGE = 3 |
| IMAGE_TYPE_GRAY_IMAGE = 4 |
|
|
| @staticmethod |
| def from_list(input_paths, |
| image_type=IMAGE_TYPE_IMAGE, |
| accept_empty=False, accept_only_en=False, |
| keep_remain=False, |
| n_jobs=4): |
|
|
| if type(input_paths) == str: |
| input_paths = [input_paths] |
| else: |
| assert type(input_paths) == list |
| xml_files = [] |
| for p in input_paths: |
| print(p) |
| p = pathlib.Path(p) |
| assert p.exists() |
| xmls = p.rglob('*.xml') |
| xml_files.extend([(f, p) for f in xmls]) |
| xml_files = sorted(xml_files) |
| return ChainDataset(Parallel(n_jobs)([ |
| delayed(XMLRawDataset)(xml_file, input_path, image_type, accept_empty, accept_only_en, keep_remain) |
| for xml_file, input_path in xml_files])) |
|
|
| def __init__(self, xml_file, input_path, |
| image_type=IMAGE_TYPE_IMAGE, |
| accept_empty=False, accept_only_en=False, |
| keep_remain=False, |
| sensitive=False): |
| self.xml_file = xml_file |
| self.input_path = input_path |
| self.image_type = image_type |
| self.accept_empty = accept_empty |
| self.accept_only_en = accept_only_en |
| self.keep_remain = keep_remain |
| self.sensitive = sensitive |
| with open(xml_file, 'r') as f: |
| try: |
| tree = ET.parse(f) |
| self.root = tree.getroot() |
| except ValueError: |
| with open(xml_file, encoding="utf-8") as file: |
| self.root = ET.fromstring(file.read()) |
| except ET.ParseError as e: |
| print(f) |
| raise e |
| groups = pt.match(self.root.tag).groups() |
| if groups[1]: |
| self.namespace = f"{{{groups[1]}}}" |
| else: |
| self.namespace = '' |
| if groups[2] != 'OCRDATASET': |
| return [] |
| self.remains = None |
|
|
| def __iter__(self): |
| if self.remains is not None: |
| for v in self.remains: |
| yield v |
| self.remains = None |
| return |
| if self.keep_remain: |
| self.remains = [] |
| image_root = None |
|
|
| for xpage in self.root: |
| page = xpage.attrib |
|
|
| image_name = page['IMAGENAME'] |
| if image_root is None: |
| image_root = find_image_root(image_name, self.xml_file, self.input_path) |
| if image_root is None: |
| print(FileNotFoundError(f'{image_name} is not Found'), file=sys.stderr) |
| continue |
| image_path = image_root / image_name |
| if not image_path.exists(): |
| print(FileNotFoundError(f'{image_name} is not Found'), file=sys.stderr) |
| continue |
| ext = image_path.suffix |
|
|
| if self.image_type is not XMLRawDataset.IMAGE_TYPE_NONE: |
| try: |
| with open(image_path, 'rb') as f: |
| imageBin = f.read() |
| imageBuf = np.frombuffer(imageBin, dtype=np.uint8) |
| img = cv2.imdecode(imageBuf, cv2.IMREAD_ANYCOLOR) |
| if img is None: |
| raise Exception('image cant not decode') |
| except Exception as e: |
| print(f'error occured at {image_path}\n{e}') |
| continue |
|
|
| if self.accept_only_en: |
| ilines = itertools.chain.from_iterable([ |
| |
| xpage.iterfind(f'.//{self.namespace}INLINE[@TYPE="欧文"]'), |
| xpage.iterfind(f'.//{self.namespace}INLINE[@TYPE="回転欧文"]'), |
| ]) |
| else: |
| data_list = [ |
| xpage.iterfind(f'.//{self.namespace}LINE'), |
| ] |
| if self.accept_empty: |
| data_list.extend([ |
| xpage.iterfind(f'.//{self.namespace}BLOCK[@TYPE="回転欧文"]'), |
| xpage.iterfind(f'.//{self.namespace}BLOCK[@TYPE="欧文"]'), |
| xpage.iterfind(f'.//{self.namespace}BLOCK[@TYPE="ノンブル"]'), |
| xpage.iterfind(f'.//{self.namespace}BLOCK[@TYPE="柱"]'), |
| ]) |
| ilines = itertools.chain.from_iterable(data_list) |
| for xline in ilines: |
| line = xline.attrib |
| line['done'] = True |
|
|
| if 'ERROR' in line: |
| continue |
|
|
| line['X'] = int(line['X']) |
| line['Y'] = int(line['Y']) |
| line['WIDTH'] = int(line['WIDTH']) |
| line['HEIGHT'] = int(line['HEIGHT']) |
| line['path'] = image_path |
| line['tag'] = xline.tag |
|
|
| if self.accept_empty: |
| line['STRING'] = line.get('STRING', None) |
| if 'DIRECTION' not in line: |
| if line['WIDTH'] > line['HEIGHT']: |
| line['DIRECTION'] = '横' |
| else: |
| line['DIRECTION'] = '縦' |
| elif len(line.get('STRING', '')) == 0: |
| continue |
| line['label'] = line['STRING'] |
| line['direction'] = line['DIRECTION'] |
|
|
| chars = list() |
| line['cattrs'] = chars |
|
|
| for char in xline: |
| char = char.attrib |
| char['X'] = int(char['X']) - line['X'] |
| char['Y'] = int(char['Y']) - line['Y'] |
| char['WIDTH'] = int(char['WIDTH']) |
| char['HEIGHT'] = int(char['HEIGHT']) |
| chars.append(char) |
|
|
| if self.image_type is not self.IMAGE_TYPE_NONE: |
| y1, y2 = line['Y'], line['Y'] + line['HEIGHT'] |
| x1, x2 = line['X'], line['X'] + line['WIDTH'] |
| g = img[y1:y2, x1:x2] |
| if self.image_type == self.IMAGE_TYPE_ENCODED: |
| retval, g = cv2.imencode(ext, g) |
| assert retval |
| elif self.image_type == self.IMAGE_TYPE_IMAGE: |
| g = Image.fromarray(g) |
| elif self.image_type == self.IMAGE_TYPE_GRAY_IMAGE: |
| g = Image.fromarray(g).convert('L') |
| else: |
| g = None |
|
|
| if not self.sensitive and line['label']: |
| line['label'] = line['label'].lower() |
|
|
| yield g, line |
| if self.keep_remain: |
| for xline in xpage.iterfind("./*"): |
| line = xline.attrib |
| if 'done' in line: |
| continue |
| line['path'] = image_path |
| line['tag'] = xline.tag |
| self.remains.append(line) |
|
|
|
|
| class SyntheticDataset(IterableDataset): |
| def __init__(self, chars, fontpath): |
| from PIL import ImageFont, ImageDraw |
| import functools |
| self.chars = chars |
| dtmp = ImageDraw.Draw(Image.new('L', (400, 200))) |
| self._font = ImageFont.truetype(fontpath, 32) |
| self._textsize = functools.partial(dtmp.multiline_textsize, font=self._font) |
|
|
| def __iter__(self): |
| import random |
| from more_itertools import chunked |
| from PIL import ImageDraw |
| chars = [] |
| chars.extend(self.chars) |
| random.shuffle(chars) |
|
|
| for s in chunked(chars, 40): |
| s = ''.join(s) |
| w, h = self._textsize(s) |
| g = Image.new('RGB', (w, h), (255, 255, 255)) |
| d = ImageDraw.Draw(g) |
| d.text((0, 0), s, font=self._font, fill=(0, 0, 0)) |
| line = { |
| 'path': None, |
| 'direction': '横', |
| 'label': s, |
| 'WIDTH': w, |
| 'HEIGHT': h, |
| 'STRING': s, |
| 'cattrs': None, |
| } |
| yield g.convert('L'), line |
|
|
|
|
| class ListRawDataset(IterableDataset): |
| IMAGE_TYPE_BYTE = 0 |
| IMAGE_TYPE_ENCODED = 1 |
| IMAGE_TYPE_IMAGE = 2 |
| IMAGE_TYPE_GRAY_IMAGE = 3 |
|
|
| def __init__(self, input_paths, read_image=True, image_type=IMAGE_TYPE_IMAGE, n_jobs=4): |
| self.read_image = read_image |
| self.image_type = image_type |
|
|
| if type(input_paths) == str: |
| input_paths = [input_paths] |
| else: |
| assert type(input_paths) == list |
|
|
| self.images = list() |
| self.labels = list() |
|
|
| for input_path in input_paths: |
| p = pathlib.Path(input_path) |
| assert p.exists() |
| with open(p, 'r') as f: |
| lines = f.readlines() |
|
|
| for line in lines: |
| image_path, label = line.strip('\n').split() |
| image_path = p.parent / image_path |
| self.images.append(image_path) |
| self.labels.append(label) |
|
|
| def __iter__(self): |
| for image_path, label in zip(self.images, self.labels): |
| ext = image_path.suffix |
|
|
| try: |
| with open(image_path, 'rb') as f: |
| imageBin = f.read() |
| imageBuf = np.frombuffer(imageBin, dtype=np.uint8) |
| img = cv2.imdecode(imageBuf, cv2.IMREAD_ANYCOLOR) |
| h, w, c = img.shape |
| if img is None: |
| raise Exception('image cant not decode') |
| except Exception as e: |
| print(f'error occured at {image_path}\n{e}') |
| continue |
|
|
| if self.read_image: |
| g = img |
| if self.image_type == self.IMAGE_TYPE_ENCODED: |
| retval, g = cv2.imencode(ext, g) |
| assert retval |
| elif self.image_type == self.IMAGE_TYPE_IMAGE: |
| g = Image.fromarray(g) |
| elif self.image_type == self.IMAGE_TYPE_GRAY_IMAGE: |
| g = Image.fromarray(g).convert('L') |
| else: |
| g = None |
| line = { |
| 'path': image_path, |
| 'direction': '横', |
| 'label': label, |
| 'WIDTH': w, |
| 'HEIGHT': h, |
| 'STRING': label, |
| } |
| yield g, line |
|
|
|
|
| class XMLRawDatasetWithCli(IterableDataset): |
| def __init__(self, img_data, xml_data, |
| accept_empty=False, accept_only_en=False, |
| sensitive=False, yield_block_pillar=True, yield_block_page_num=True): |
|
|
| self.xml_data = xml_data |
| self.orig_img = img_data |
| self.accept_empty = accept_empty |
| self.accept_only_en = accept_only_en |
| self.sensitive = sensitive |
| self.root = xml_data.getroot() |
| self.remains = None |
| self.yield_block_pillar = yield_block_pillar |
| self.yield_block_page_num = yield_block_page_num |
|
|
| groups = pt.match(self.root.tag).groups() |
| if groups[1]: |
| self.namespace = f"{{{groups[1]}}}" |
| else: |
| self.namespace = '' |
| if groups[2] != 'OCRDATASET': |
| return [] |
|
|
| def __iter__(self): |
| from PIL import Image |
| root = self.xml_data.getroot() |
| for xml_line in root.find('PAGE'): |
| is_block_page_num = (xml_line.tag == 'BLOCK') and (xml_line.attrib['TYPE'] == 'ノンブル') |
| is_block_pillar = (xml_line.tag == 'BLOCK') and (xml_line.attrib['TYPE'] == '柱') |
| do_yield = (xml_line.tag == 'LINE') or (is_block_page_num and self.yield_block_page_num) or (is_block_pillar and self.yield_block_pillar) |
| if not do_yield: |
| print("This {0} elemetn will be skipped.".format(xml_line.tag)) |
| print(xml_line.attrib) |
| continue |
| xml_attrib = xml_line.attrib |
|
|
| w = xml_attrib['WIDTH'] |
| h = xml_attrib['HEIGHT'] |
| y1, y2 = int(xml_attrib['Y']), int(xml_attrib['Y']) + int(xml_attrib['HEIGHT']) |
| x1, x2 = int(xml_attrib['X']), int(xml_attrib['X']) + int(xml_attrib['WIDTH']) |
| g = Image.fromarray(self.orig_img[y1:y2, x1:x2]).convert('L') |
|
|
| direction = '横' |
| if 'DIRECTION' not in xml_attrib: |
| if int(w) > int(h): |
| direction = '横' |
| else: |
| direction = '縦' |
| else: |
| direction = xml_attrib['DIRECTION'] |
|
|
| line = { |
| 'path': None, |
| 'direction': direction, |
| 'label': None, |
| 'WIDTH': w, |
| 'HEIGHT': h, |
| 'STRING': None, |
| 'cattrs': None, |
| 'X' : xml_attrib['X'], |
| 'Y' : xml_attrib['Y'], |
| 'TYPE' :xml_attrib['TYPE'], |
| 'tag' : xml_line.tag, |
| 'STRING': None, |
| 'done': True |
| } |
|
|
| yield g, line |
|
|
