| 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)) |
|
|
|
|
