| 28 | |
| 29 | |
| 30 | def write_worker(args, q_out): |
| 31 | pre_time = time.time() |
| 32 | |
| 33 | if args.input[-1] == '/': |
| 34 | args.input = args.input[:-1] |
| 35 | dirname = os.path.dirname(args.input) |
| 36 | basename = os.path.basename(args.input) |
| 37 | output = os.path.join(dirname, f"shuffled_{basename}") |
| 38 | os.makedirs(output, exist_ok=True) |
| 39 | |
| 40 | path_imgidx = os.path.join(output, "train.idx") |
| 41 | path_imgrec = os.path.join(output, "train.rec") |
| 42 | save_record = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, "w") |
| 43 | more = True |
| 44 | count = 0 |
| 45 | while more: |
| 46 | deq = q_out.get() |
| 47 | if deq is None: |
| 48 | more = False |
| 49 | else: |
| 50 | header, jpeg = mx.recordio.unpack(deq) |
| 51 | # TODO it is currently not fully developed |
| 52 | if isinstance(header.label, float): |
| 53 | label = header.label |
| 54 | else: |
| 55 | label = header.label[0] |
| 56 | |
| 57 | header = mx.recordio.IRHeader(flag=header.flag, label=label, id=header.id, id2=header.id2) |
| 58 | save_record.write_idx(count, mx.recordio.pack(header, jpeg)) |
| 59 | count += 1 |
| 60 | if count % 10000 == 0: |
| 61 | cur_time = time.time() |
| 62 | print('save time:', cur_time - pre_time, ' count:', count) |
| 63 | pre_time = cur_time |
| 64 | print(count) |
| 65 | save_record.close() |
| 66 | |
| 67 | |
| 68 | def main(args): |