in eft/datasets/video_dataset.py [0:0]
def __init__(self, options, dataset, ignore_3d=False, use_augmentation=True, is_train=True, bMiniTest=False, bEnforceUpperOnly= False):
super(BaseDataset, self).__init__()
self.datasetName = dataset
self.is_train = is_train
self.options = options
self.img_dir = config.DATASET_FOLDERS[dataset]
self.normalize_img = Normalize(mean=constants.IMG_NORM_MEAN, std=constants.IMG_NORM_STD)
self.bEnforceUpperOnly = bEnforceUpperOnly
if options is not None:
if dataset=='coco3d' :
config.SetDBName(dataset, options.db_coco3d_name)
if dataset=='mpii3d':
config.SetDBName(dataset, options.db_mpii3d_name)
if dataset=='lspet3d':
config.SetDBName(dataset, options.db_lspet3d_name)
if dataset=='posetrack3d':
config.SetDBName(dataset, options.db_posetrack3d_name)
if dataset=='cocoplus3d':
config.SetDBName(dataset, options.db_cocoplus3d_name)
self.data = np.load(config.DATASET_FILES[is_train][dataset])
"""
The folling should have the same length
self.imgname
self.annotIds
self.scale
self.center
self.pose
self.keypoints
self.gender
self.has_smpl
self.maskname
self.partname
This should reset: self.length
"""
self.imgname = self.data['imgname']
#COCO only
self.annotIds = None
if 'annotIds' in self.data.files:
self.annotIds = self.data['annotIds']
assert len(self.imgname)==len(self.annotIds)
print(">>> BaseDataset:: Loading DB {}: {} samples".format(self.datasetName, len(self.imgname) ))
# Get paths to gt masks, if available
try:
self.maskname = self.data['maskname']
except KeyError:
pass
try:
self.partname = self.data['partname']
except KeyError:
pass
# Bounding boxes are assumed to be in the center and scale format
self.scale = self.data['scale']
self.center = self.data['center']
# If False, do not do augmentation
self.use_augmentation = use_augmentation
# Get gt SMPL parameters, if available
try:
self.pose = self.data['pose'].astype(np.float)
self.betas = self.data['shape'].astype(np.float)
if 'has_smpl' in self.data:
self.has_smpl = self.data['has_smpl']
else:
self.has_smpl = np.ones(len(self.imgname))
except KeyError:
self.has_smpl = np.zeros(len(self.imgname))
if ignore_3d:
self.has_smpl = np.zeros(len(self.imgname))
# Get gt 3D pose, if available
try:
self.pose_3d = self.data['S']
self.has_pose_3d = 1
except KeyError:
self.has_pose_3d = 0
if ignore_3d:
self.has_pose_3d = 0
# Get 2D keypoints
try:
keypoints_gt = self.data['part']
except KeyError:
keypoints_gt = np.zeros((len(self.imgname), 24, 3))
try:
keypoints_openpose = self.data['openpose']
except KeyError:
keypoints_openpose = np.zeros((len(self.imgname), 25, 3))
self.keypoints = np.concatenate([keypoints_openpose, keypoints_gt], axis=1)
# Get hand data, if available
if self.options.bUseHand3D:
try:
self.rhand_2d = self.data['rhand_2d']
self.lhand_2d = self.data['lhand_2d']
self.rhand_3d = self.data['rhand_3d']
self.lhand_3d = self.data['lhand_3d']
except KeyError:
pass
# Get gender data, if available
try:
gender = self.data['gender']
self.gender = np.array([0 if str(g) == 'm' else 1 for g in gender]).astype(np.int32)
except KeyError:
self.gender = -1*np.ones(len(self.imgname)).astype(np.int32)
self.length = self.scale.shape[0]
self.bLoadImage = True
if self.bEnforceUpperOnly or (self.options is not None and self.options.bUpperBodyTest):
scaleFactor = 1.2
self.scale_upperbody = np.zeros(self.scale.shape)
self.center_upperbody = np.zeros(self.center.shape)
for i in range(self.keypoints.shape[0]):
validity = self.keypoints[i,25:,2]==1 #(N, 24)
# validity[ [0,1,2,3,4,5]] = False #Disable lower bodies
validity[ [0,1,4,5]] = False #Disable lower bodies
if np.sum(validity)>=3:
min_pt = np.min(self.keypoints[i,25:,:2][validity], axis=0)
max_pt = np.max(self.keypoints[i,25:,:2][validity], axis=0)
# bbox= [ min_pt[0], min_pt[1], max_pt[0], max_pt[1] ]
bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]
center = [bbox[0] + bbox[2]/2, bbox[1] + bbox[3]/2]
scale = scaleFactor*max(bbox[2], bbox[3])/200
self.center_upperbody[i]= center
self.scale_upperbody[i]= scale
else:
# self.center_upperbody[i]= [0,0]
self.scale_upperbody[i]= -1 #Garbage should be ignored
self.center_facePart = np.zeros(self.center.shape)
self.scale_facePart = np.zeros(self.scale.shape) #More extreme case
for i in range(self.keypoints.shape[0]):
validity = self.keypoints[i,25:,2]==1 #(N, 24)
validity[ [0,1,2,3,4,5,1,14, 6,11 ]] = False #Disable, even hand, arms. Only consider shoulders
if np.sum(validity)>=3:
min_pt = np.min(self.keypoints[i,25:,:2][validity], axis=0)
max_pt = np.max(self.keypoints[i,25:,:2][validity], axis=0)
# bbox= [ min_pt[0], min_pt[1], max_pt[0], max_pt[1] ]
bbox= [ min_pt[0], min_pt[1], max_pt[0] - min_pt[0], max_pt[1] - min_pt[1]]
center = [bbox[0] + bbox[2]/2, bbox[1] + bbox[3]/2]
scale = scaleFactor*max(bbox[2], bbox[3])/200
self.center_facePart[i]= center
self.scale_facePart[i]= scale
else:
# self.center_upperbody[i]= [0,0]
self.scale_facePart[i]= -1 #Garbage should be ignored
if bMiniTest: #Sampling to 1/10 size
# validIdx = range(0,len(self.imgname),10)
if self.options.bExemplar_badsample_finder:
validIdx = range(0,len(self.imgname),100)
else:
validIdx = range(0,len(self.imgname),10)
print(">>> Debug: sampling originalSize {} -> miniSize {}".format(len(self.imgname), len(validIdx)))
self.imgname = self.imgname[validIdx]
try:
self.pose = self.pose[validIdx]
except:
pass
try:
self.betas = self.betas[validIdx]
except:
pass
try:
self.pose_3d = self.pose_3d[validIdx]
except:
pass
try:
self.keypoints = self.keypoints [validIdx]
except:
pass
try:
self.keypoints = self.keypoints [validIdx]
except:
pass
try:
self.scale_upperbody = self.scale_upperbody[validIdx]
except:
pass
try:
self.center_upperbody = self.center_upperbody[validIdx]
except:
pass
# try:
# self.scale_facePart = self.scale_facePart[validIdx]
# except:
# pass
# try:
# self.center_facePart = self.center_facePart[validIdx]
# except:
# pass
self.has_smpl = self.has_smpl[validIdx]
# self.has_pose_3d = self.has_pose_3d[validIdx]
self.scale =self.scale[validIdx]
self.center= self.center[validIdx]
self.gender= self.gender[validIdx]
if self.annotIds is not None:
self.annotIds = self.annotIds[validIdx]
#Debug: choose specific range to run paralle (training code only)
if self.options is not None and\
is_train == True and\
self.options.bExemplar_dataLoaderStart >=0:# and self.options.bExemplar_dataLoaderEnd >=0:
startIdx = self.options.bExemplar_dataLoaderStart
print(">> {}".format( len(self.imgname)))
if self.options.bExemplar_dataLoaderEnd>=0:
endIdx = min( self.options.bExemplar_dataLoaderEnd, len(self.imgname))
else:
endIdx = len(self.imgname)
#Run only from bExemplar_dataLoaderStart ~ bExemplar_dataLoaderEnd
validIdx = range(startIdx,endIdx)
print(validIdx)
# for i in range(len(self.imgname)):
# if seqName in self.imgname[i]:
# validIdx.append(i)
self.imgname = self.imgname[validIdx]
try:
self.pose = self.pose[validIdx]
except:
pass
try:
self.betas = self.betas[validIdx]
except:
pass
try:
self.pose_3d = self.pose_3d[validIdx]
except:
pass
#Hand part
try:
self.rhand_2d = self.rhand_2d[validIdx]
except:
pass
try:
self.rhand_3d = self.rhand_3d[validIdx]
except:
pass
try:
self.lhand_2d = self.lhand_2d[validIdx]
except:
pass
try:
self.lhand_3d = self.lhand_3d[validIdx]
except:
pass
# self.has_pose_3d = self.has_pose_3d[validIdx]
self.keypoints = self.keypoints [validIdx]
self.has_smpl = self.has_smpl[validIdx]
self.scale =self.scale[validIdx]
self.center= self.center[validIdx]
self.gender= self.gender[validIdx]
if self.annotIds is not None: #bug fixed
self.annotIds = self.annotIds[validIdx]
# self.dataset= self.dataset[validIdx]
# item['maskname'] = self.maskname[index]
# item['partname'] = ''
print("Selected IDx: {} - {}. Total number {}".format(startIdx, endIdx, self.__len__()))
#Debug: choose specific sequence
if False:
# seqName ='office_phoneCall_00'
seqName ='downtown'
validIdx =[]
for i in range(len(self.imgname)):
if seqName in self.imgname[i]:
validIdx.append(i)
self.imgname = self.imgname[validIdx]
self.pose = self.pose[validIdx]
self.betas = self.betas[validIdx]
# self.pose_3d = self.pose_3d[validIdx]
self.keypoints = self.keypoints [validIdx]
self.has_smpl = self.has_smpl[validIdx]
# self.has_pose_3d = self.has_pose_3d[validIdx]
self.scale =self.scale[validIdx]
self.center= self.center[validIdx]
self.gender= self.gender[validIdx]
# self.dataset= self.dataset[validIdx]
# item['maskname'] = self.maskname[index]
# item['partname'] = ''
#Check
dbLeng = len(self.imgname)
self.length = dbLeng
assert(dbLeng==len(self.scale)==len(self.center)==len(self.keypoints)==len(self.gender)==len(self.has_smpl))
if self.annotIds is not None: #Bug fixed
dbLeng = len(self.annotIds)
if hasattr(self,'pose'):
assert(dbLeng ==len(self.pose))
if hasattr(self,'betas'):
assert(dbLeng ==len(self.betas))
if hasattr(self,'pose_3d'):
assert(dbLeng ==len(self.pose_3d))
if hasattr(self,'maskname'):
assert(dbLeng ==len(self.maskname))
if hasattr(self,'partname'):
assert(dbLeng ==len(self.partname))
if hasattr(self,'rhand_2d'):
assert(dbLeng ==len(self.rhand_2d))
if hasattr(self,'lhand_2d'):
assert(dbLeng ==len(self.lhand_2d))