| | import pickle |
| | import numpy as np |
| | import os |
| | import os.path as osp |
| | import sys |
| | import mxnet as mx |
| |
|
| |
|
| | class RecBuilder(): |
| | def __init__(self, path, image_size=(112, 112)): |
| | self.path = path |
| | self.image_size = image_size |
| | self.widx = 0 |
| | self.wlabel = 0 |
| | self.max_label = -1 |
| | assert not osp.exists(path), '%s exists' % path |
| | os.makedirs(path) |
| | self.writer = mx.recordio.MXIndexedRecordIO(os.path.join(path, 'train.idx'), |
| | os.path.join(path, 'train.rec'), |
| | 'w') |
| | self.meta = [] |
| |
|
| | def add(self, imgs): |
| | |
| | |
| | |
| | assert len(imgs) > 0 |
| | label = self.wlabel |
| | for img in imgs: |
| | idx = self.widx |
| | image_meta = {'image_index': idx, 'image_classes': [label]} |
| | header = mx.recordio.IRHeader(0, label, idx, 0) |
| | if isinstance(img, np.ndarray): |
| | s = mx.recordio.pack_img(header,img,quality=95,img_fmt='.jpg') |
| | else: |
| | s = mx.recordio.pack(header, img) |
| | self.writer.write_idx(idx, s) |
| | self.meta.append(image_meta) |
| | self.widx += 1 |
| | self.max_label = label |
| | self.wlabel += 1 |
| |
|
| |
|
| | def add_image(self, img, label): |
| | |
| | |
| | |
| | idx = self.widx |
| | header = mx.recordio.IRHeader(0, label, idx, 0) |
| | if isinstance(label, list): |
| | idlabel = label[0] |
| | else: |
| | idlabel = label |
| | image_meta = {'image_index': idx, 'image_classes': [idlabel]} |
| | if isinstance(img, np.ndarray): |
| | s = mx.recordio.pack_img(header,img,quality=95,img_fmt='.jpg') |
| | else: |
| | s = mx.recordio.pack(header, img) |
| | self.writer.write_idx(idx, s) |
| | self.meta.append(image_meta) |
| | self.widx += 1 |
| | self.max_label = max(self.max_label, idlabel) |
| |
|
| | def close(self): |
| | with open(osp.join(self.path, 'train.meta'), 'wb') as pfile: |
| | pickle.dump(self.meta, pfile, protocol=pickle.HIGHEST_PROTOCOL) |
| | print('stat:', self.widx, self.wlabel) |
| | with open(os.path.join(self.path, 'property'), 'w') as f: |
| | f.write("%d,%d,%d\n" % (self.max_label+1, self.image_size[0], self.image_size[1])) |
| | f.write("%d\n" % (self.widx)) |
| |
|
| |
|
