Spaces:
Running
Running
| import os | |
| import h5py | |
| from functools import reduce | |
| import torch.utils.data as data | |
| from ..pose import generateSampleBox | |
| from opt import opt | |
| class Mscoco(data.Dataset): | |
| def __init__(self, train=True, sigma=1, | |
| scale_factor=0.25, rot_factor=30, label_type='Gaussian'): | |
| self.img_folder = '../data/' # root image folders | |
| self.is_train = train # training set or test set | |
| self.inputResH = 320 | |
| self.inputResW = 256 | |
| self.outputResH = 80 | |
| self.outputResW = 64 | |
| self.sigma = sigma | |
| self.scale_factor = (0.2, 0.3) | |
| self.rot_factor = rot_factor | |
| self.label_type = label_type | |
| self.nJoints_coco = 17 | |
| self.nJoints_mpii = 16 | |
| self.nJoints = 33 | |
| self.accIdxs = (1, 2, 3, 4, 5, 6, 7, 8, # COCO | |
| 9, 10, 11, 12, 13, 14, 15, 16, 17, | |
| 18, 19, 20, 21, 22, 23, # MPII | |
| 28, 29, 32, 33) | |
| self.flipRef = ((2, 3), (4, 5), (6, 7), # COCO | |
| (8, 9), (10, 11), (12, 13), | |
| (14, 15), (16, 17), | |
| (18, 23), (19, 22), (20, 21), # MPII | |
| (28, 33), (29, 32), (30, 31)) | |
| ''' | |
| Create train/val split | |
| ''' | |
| # COCO | |
| with h5py.File('../data/coco/annot_clean.h5', 'r') as annot: | |
| # train | |
| self.imgname_coco_train = annot['imgname'][:-5887] | |
| self.bndbox_coco_train = annot['bndbox'][:-5887] | |
| self.part_coco_train = annot['part'][:-5887] | |
| # val | |
| self.imgname_coco_val = annot['imgname'][-5887:] | |
| self.bndbox_coco_val = annot['bndbox'][-5887:] | |
| self.part_coco_val = annot['part'][-5887:] | |
| # MPII | |
| with h5py.File('../data/mpii/annot_mpii.h5', 'r') as annot: | |
| # train | |
| self.imgname_mpii_train = annot['imgname'][:-1358] | |
| self.bndbox_mpii_train = annot['bndbox'][:-1358] | |
| self.part_mpii_train = annot['part'][:-1358] | |
| # val | |
| self.imgname_mpii_val = annot['imgname'][-1358:] | |
| self.bndbox_mpii_val = annot['bndbox'][-1358:] | |
| self.part_mpii_val = annot['part'][-1358:] | |
| self.size_coco_train = self.imgname_coco_train.shape[0] | |
| self.size_coco_val = self.imgname_coco_val.shape[0] | |
| self.size_train = self.imgname_coco_train.shape[0] + self.imgname_mpii_train.shape[0] | |
| self.size_val = self.imgname_coco_val.shape[0] + self.imgname_mpii_val.shape[0] | |
| self.train, self.valid = [], [] | |
| def __getitem__(self, index): | |
| sf = self.scale_factor | |
| if self.is_train and index < self.size_coco_train: # COCO | |
| part = self.part_coco_train[index] | |
| bndbox = self.bndbox_coco_train[index] | |
| imgname = self.imgname_coco_train[index] | |
| imgset = 'coco' | |
| elif self.is_train: # MPII | |
| part = self.part_mpii_train[index - self.size_coco_train] | |
| bndbox = self.bndbox_mpii_train[index - self.size_coco_train] | |
| imgname = self.imgname_mpii_train[index - self.size_coco_train] | |
| imgset = 'mpii' | |
| elif index < self.size_coco_val: | |
| part = self.part_coco_val[index] | |
| bndbox = self.bndbox_coco_val[index] | |
| imgname = self.imgname_coco_val[index] | |
| imgset = 'coco' | |
| else: | |
| part = self.part_mpii_val[index - self.size_coco_val] | |
| bndbox = self.bndbox_mpii_val[index - self.size_coco_val] | |
| imgname = self.imgname_mpii_val[index - self.size_coco_val] | |
| imgset = 'mpii' | |
| if imgset == 'coco': | |
| imgname = reduce(lambda x, y: x + y, map(lambda x: chr(int(x)), imgname)) | |
| else: | |
| imgname = reduce(lambda x, y: x + y, map(lambda x: chr(int(x)), imgname))[:13] | |
| img_path = os.path.join(self.img_folder, imgset, 'images', imgname) | |
| metaData = generateSampleBox(img_path, bndbox, part, self.nJoints, | |
| imgset, sf, self, train=self.is_train) | |
| inp, out_bigcircle, out_smallcircle, out, setMask = metaData | |
| label = [] | |
| for i in range(opt.nStack): | |
| if i < 2: | |
| # label.append(out_bigcircle.clone()) | |
| label.append(out.clone()) | |
| elif i < 4: | |
| # label.append(out_smallcircle.clone()) | |
| label.append(out.clone()) | |
| else: | |
| label.append(out.clone()) | |
| return inp, label, setMask, imgset | |
| def __len__(self): | |
| if self.is_train: | |
| return self.size_train | |
| else: | |
| return self.size_val | |