Create LMDB dataset for training and evaluation. ARGS: inputPath : input folder path where starts imagePath outputPath : LMDB output path gtFile : list of image path and label checkValid : if true, check the validity of every image
(data_list, outputPath, checkValid=True)
| 54 | |
| 55 | |
| 56 | def createDataset(data_list, outputPath, checkValid=True): |
| 57 | """ |
| 58 | Create LMDB dataset for training and evaluation. |
| 59 | ARGS: |
| 60 | inputPath : input folder path where starts imagePath |
| 61 | outputPath : LMDB output path |
| 62 | gtFile : list of image path and label |
| 63 | checkValid : if true, check the validity of every image |
| 64 | """ |
| 65 | os.makedirs(outputPath, exist_ok=True) |
| 66 | env = lmdb.open(outputPath, map_size=1099511627776) |
| 67 | cache = {} |
| 68 | cnt = 1 |
| 69 | for imagePath, label in tqdm(data_list, |
| 70 | desc=f'make dataset, save to {outputPath}'): |
| 71 | with open(imagePath, 'rb') as f: |
| 72 | imageBin = f.read() |
| 73 | buf = io.BytesIO(imageBin) |
| 74 | w, h = Image.open(buf).size |
| 75 | if checkValid: |
| 76 | try: |
| 77 | if not checkImageIsValid(imageBin): |
| 78 | print('%s is not a valid image' % imagePath) |
| 79 | continue |
| 80 | except: |
| 81 | continue |
| 82 | |
| 83 | imageKey = 'image-%09d'.encode() % cnt |
| 84 | labelKey = 'label-%09d'.encode() % cnt |
| 85 | whKey = 'wh-%09d'.encode() % cnt |
| 86 | cache[imageKey] = imageBin |
| 87 | cache[labelKey] = label.encode() |
| 88 | cache[whKey] = (str(w) + '_' + str(h)).encode() |
| 89 | |
| 90 | if cnt % 1000 == 0: |
| 91 | writeCache(env, cache) |
| 92 | cache = {} |
| 93 | cnt += 1 |
| 94 | nSamples = cnt - 1 |
| 95 | cache['num-samples'.encode()] = str(nSamples).encode() |
| 96 | writeCache(env, cache) |
| 97 | print('Created dataset with %d samples' % nSamples) |
| 98 | |
| 99 | |
| 100 | if __name__ == '__main__': |
no test coverage detected