Spaces:
Running
Running
File size: 4,753 Bytes
95f8bbc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
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
|