Spaces:
Running
on
Zero
Running
on
Zero
| 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): | |
| #!!! img should be BGR!!!! | |
| #assert label >= 0 | |
| #assert label > self.last_label | |
| 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): | |
| #!!! img should be BGR!!!! | |
| #assert label >= 0 | |
| #assert label > self.last_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)) | |