| 28 | return ''.join(re.findall("[0-9a-zA-Z]",key)) |
| 29 | |
| 30 | class LoadDatasetLmdb(Dataset): |
| 31 | def __init__(self,config,lmdb_file): |
| 32 | num_workers = config['train']['num_workers'] |
| 33 | self.fixKey = config['train']['fixKeyON'] |
| 34 | self.fixKeyType = config['train']['fixKeytype'] |
| 35 | assert self.fixKeyType in ['En','Ch'] |
| 36 | self.env = lmdb.open(lmdb_file, max_readers=num_workers, readonly=True, lock=False, readahead=False, meminit=False) |
| 37 | if not self.env: |
| 38 | print('cannot creat lmdb from %s' % (lmdb_file)) |
| 39 | sys.exit(0) |
| 40 | |
| 41 | with self.env.begin(write=False) as txn: |
| 42 | nSamples = int(txn.get('num-samples'.encode('utf-8'))) |
| 43 | self.nSamples = nSamples |
| 44 | |
| 45 | |
| 46 | def __len__(self): |
| 47 | return self.nSamples |
| 48 | |
| 49 | def __getitem__(self, index): |
| 50 | assert index <= len(self), 'index range error' |
| 51 | index += 1 |
| 52 | with self.env.begin(write=False) as txn: |
| 53 | img_key = 'image-%09d' % index |
| 54 | imgbuf = txn.get(img_key.encode('utf-8')) |
| 55 | buf = six.BytesIO() |
| 56 | buf.write(imgbuf) |
| 57 | buf.seek(0) |
| 58 | try: |
| 59 | img = Image.open(buf).convert('RGB') |
| 60 | except IOError: |
| 61 | print('Corrupted image for %d' % index) |
| 62 | return self[index + 1] |
| 63 | |
| 64 | label_key = 'label-%09d' % index |
| 65 | label = txn.get(label_key.encode('utf-8')).decode().replace('\ufeff', '').replace('\u3000', '').strip() |
| 66 | if self.fixKey: |
| 67 | if self.fixKeyType == 'En': |
| 68 | label = fixkeyEn(label) |
| 69 | label = label.lower() |
| 70 | elif self.fixKeyType == 'Ch': |
| 71 | label = fixkeyCh(label) |
| 72 | return (img, label) |
| 73 | |
| 74 | |
| 75 | class resizeNormalize(object): |