in visualization/POF/data/DomeReaderTempConst.py [0:0]
def __init__(self, mode='training', objtype=0, shuffle=False, batch_size=1, crop_noise=False, full_only=True, head_top=True):
super(DomeReaderTempConst, self).__init__(objtype, shuffle, batch_size, crop_noise)
assert mode in ('training', 'evaluation')
assert objtype in (0, 1)
self.image_root = '/media/posefs0c/panopticdb/'
# read data from a4plus with consecutive frames
path_to_db = './data/a4plus_collected.pkl'
path_to_calib = './data/camera_data_a4.pkl'
with open(path_to_db, 'rb') as f:
db_data = pickle.load(f)
with open('./data/a4_hands_annotated.txt') as f:
hand_annots = {}
for line in f:
strs = line.split()
hand_annots[tuple(strs[:3])] = eval(strs[3])
if mode == 'training':
mode_data = db_data['training_data']
else:
mode_data = db_data['testing_data']
with open(path_to_calib, 'rb') as f:
calib_data = pickle.load(f)
human3d = {'1_body': [], '1_left_hand': [], '1_right_hand': [], '1_body_valid': [], 'left_hand_valid': [], 'right_hand_valid': [],
'2_body': [], '2_left_hand': [], '2_right_hand': [], '2_body_valid': []}
calib = {'1_K': [], '1_R': [], '1_t': [], '1_distCoef': [],
'2_K': [], '2_R': [], '2_t': [], '2_distCoef': []}
img_dirs_1 = []
img_dirs_2 = []
map_next = {}
for data3d in mode_data:
seqName = data3d['seqName']
frame_str = data3d['frame_str']
frame = int(frame_str)
if frame % 5: # a4plus is sampled 1 out of 5, frame number *0 and *5 is the first frame, *1 and *6 is the second frame
continue
map_next[(seqName, frame_str)] = None
for data3d in mode_data:
seqName = data3d['seqName']
frame_str = data3d['frame_str']
frame = int(frame_str)
if frame % 5 != 1:
continue
prev_key = (seqName, '{:08d}'.format(frame - 1))
if prev_key not in map_next:
continue
map_next[prev_key] = data3d
for data3d in mode_data:
seqName = data3d['seqName']
frame_str = data3d['frame_str']
frame = int(frame_str)
if frame % 5:
continue
# check for manual annotation, remove the annotation if the hand is annotated as incorrect.
if 'left_hand' in data3d and not hand_annots[(seqName, frame_str, 'left')]:
del data3d['left_hand']
if 'right_hand' in data3d and not hand_annots[(seqName, frame_str, 'righ')]:
del data3d['right_hand']
next_data = map_next[(seqName, frame_str)]
if next_data is None:
continue
if objtype == 0:
body3d = np.array(data3d['body']['landmarks'], dtype=np.float32).reshape(-1, 3)
nose_lear = body3d[16] - body3d[1]
nose_rear = body3d[18] - body3d[1]
neck_nose = body3d[1] - body3d[0]
n = np.cross(nose_lear, nose_rear)
n = n / nl.norm(n)
d = np.dot(neck_nose, n)
assert d > 0
head_top_kp = body3d[0] + 1.5 * d * n
if head_top:
body3d = np.concatenate((body3d, head_top_kp[np.newaxis, :]), axis=0)
chest = 0.5 * body3d[0] + 0.25 * (body3d[6] + body3d[12])
body3d_1 = np.concatenate((body3d, chest[np.newaxis, :]), axis=0)
body3d = np.array(next_data['body']['landmarks'], dtype=np.float32).reshape(-1, 3)
nose_lear = body3d[16] - body3d[1]
nose_rear = body3d[18] - body3d[1]
neck_nose = body3d[1] - body3d[0]
n = np.cross(nose_lear, nose_rear)
n = n / nl.norm(n)
d = np.dot(neck_nose, n)
assert d > 0
head_top_kp = body3d[0] + 1.5 * d * n
if head_top:
body3d = np.concatenate((body3d, head_top_kp[np.newaxis, :]), axis=0)
chest = 0.5 * body3d[0] + 0.25 * (body3d[6] + body3d[12])
body3d_2 = np.concatenate((body3d, chest[np.newaxis, :]), axis=0)
elif objtype == 1:
# left hand or right hand must be valid
if 'left_hand' in data3d:
left_hand3d_1 = np.array(data3d['left_hand']['landmarks'], dtype=np.float32).reshape(-1, 3)
if 'left_hand' in next_data:
left_hand3d_2 = np.array(next_data['left_hand']['landmarks'], dtype=np.float32).reshape(-1, 3)
if 'right_hand' in data3d:
right_hand3d_1 = np.array(data3d['right_hand']['landmarks'], dtype=np.float32).reshape(-1, 3)
if 'right_hand' in next_data:
right_hand3d_2 = np.array(next_data['right_hand']['landmarks'], dtype=np.float32).reshape(-1, 3)
if ('left_hand' not in data3d or 'left_hand' not in next_data) and ('right_hand' not in data3d or 'right_hand' not in next_data):
# one hand must be valid for both frames
continue
if objtype == 0:
for camIdx, camDict in data3d['body']['2D'].items():
if camIdx not in next_data['body']['2D']: # no data from this camera in the next frame
continue
if full_only:
cond_inside_1 = all(camDict['insideImg'])
cond_inside_2 = all(next_data['body'][camIdx]['insideImg'])
else: # if not full_only, use the image if at least half keypoints are visible
inside_ratio_1 = np.float(np.sum(camDict['insideImg'])) / len(camDict['insideImg'])
inside_ratio_2 = np.float(np.sum(next_data['body']['2D'][camIdx]['insideImg'])) / len(next_data['body']['2D'][camIdx]['insideImg'])
cond_inside_1 = (inside_ratio_1 > 0.1)
cond_inside_2 = (inside_ratio_2 > 0.1)
if any(camDict['occluded']) or any(next_data['body']['2D'][camIdx]['occluded']) or not cond_inside_1 or not cond_inside_2:
continue
human3d['1_body'].append(body3d_1)
human3d['2_body'].append(body3d_2)
human3d['1_body_valid'].append(np.ones((20 if head_top else 19,), dtype=bool))
human3d['2_body_valid'].append(np.ones((20 if head_top else 19,), dtype=bool))
calib['1_K'].append(calib_data[seqName][camIdx]['K'].astype(np.float32))
calib['2_K'].append(calib_data[seqName][camIdx]['K'].astype(np.float32))
calib['1_R'].append(calib_data[seqName][camIdx]['R'].astype(np.float32))
calib['2_R'].append(calib_data[seqName][camIdx]['R'].astype(np.float32))
calib['1_t'].append(calib_data[seqName][camIdx]['t'][:, 0].astype(np.float32))
calib['2_t'].append(calib_data[seqName][camIdx]['t'][:, 0].astype(np.float32))
calib['1_distCoef'].append(calib_data[seqName][camIdx]['distCoef'].astype(np.float32))
calib['2_distCoef'].append(calib_data[seqName][camIdx]['distCoef'].astype(np.float32))
img_dirs_1.append('{}/{}/hdImgs/{}/{}/00_{:02d}_{}.jpg'.format(self.image_root, 'a4', seqName, frame_str, camIdx, frame_str))
img_dirs_2.append('{}/{}/hdImgs/{}/{}/00_{:02d}_{}.jpg'.format(self.image_root, 'a4', seqName, next_data['frame_str'], camIdx, next_data['frame_str']))
elif objtype == 1:
if 'left_hand' in data3d and 'left_hand' in next_data:
for camIdx, camDict in data3d['left_hand']['2D'].items():
if camIdx not in next_data['left_hand']['2D']:
continue
if any(data3d['left_hand']['2D'][camIdx]['occluded']) or not all(data3d['left_hand']['2D'][camIdx]['insideImg']) or data3d['left_hand']['2D'][camIdx]['overlap']:
continue
if any(next_data['left_hand']['2D'][camIdx]['occluded']) or not all(next_data['left_hand']['2D'][camIdx]['insideImg']) or next_data['left_hand']['2D'][camIdx]['overlap']:
continue
human3d['1_left_hand'].append(left_hand3d_1)
human3d['2_left_hand'].append(left_hand3d_2)
human3d['1_right_hand'].append(np.zeros((21, 3), dtype=np.float32))
human3d['2_right_hand'].append(np.zeros((21, 3), dtype=np.float32))
human3d['left_hand_valid'].append(np.ones((21,), dtype=bool))
human3d['right_hand_valid'].append(np.zeros((21,), dtype=bool))
calib['1_K'].append(calib_data[seqName][camIdx]['K'].astype(np.float32))
calib['2_K'].append(calib_data[seqName][camIdx]['K'].astype(np.float32))
calib['1_R'].append(calib_data[seqName][camIdx]['R'].astype(np.float32))
calib['2_R'].append(calib_data[seqName][camIdx]['R'].astype(np.float32))
calib['1_t'].append(calib_data[seqName][camIdx]['t'][:, 0].astype(np.float32))
calib['2_t'].append(calib_data[seqName][camIdx]['t'][:, 0].astype(np.float32))
calib['1_distCoef'].append(calib_data[seqName][camIdx]['distCoef'].astype(np.float32))
calib['2_distCoef'].append(calib_data[seqName][camIdx]['distCoef'].astype(np.float32))
img_dirs_1.append('{}/{}/hdImgs/{}/{}/00_{:02d}_{}.jpg'.format(self.image_root, 'a4', seqName, frame_str, camIdx, frame_str))
img_dirs_2.append('{}/{}/hdImgs/{}/{}/00_{:02d}_{}.jpg'.format(self.image_root, 'a4', seqName, next_data['frame_str'], camIdx, next_data['frame_str']))
if 'right_hand' in data3d and 'right_hand' in next_data:
for camIdx, camDict in data3d['right_hand']['2D'].items():
if camIdx not in next_data['right_hand']['2D']:
continue
if any(data3d['right_hand']['2D'][camIdx]['occluded']) or not all(data3d['right_hand']['2D'][camIdx]['insideImg']) or data3d['right_hand']['2D'][camIdx]['overlap']:
continue
if any(next_data['right_hand']['2D'][camIdx]['occluded']) or not all(next_data['right_hand']['2D'][camIdx]['insideImg']) or next_data['right_hand']['2D'][camIdx]['overlap']:
continue
human3d['1_right_hand'].append(right_hand3d_1)
human3d['2_right_hand'].append(right_hand3d_2)
human3d['1_left_hand'].append(np.zeros((21, 3), dtype=np.float32))
human3d['2_left_hand'].append(np.zeros((21, 3), dtype=np.float32))
human3d['left_hand_valid'].append(np.zeros((21,), dtype=bool))
human3d['right_hand_valid'].append(np.ones((21,), dtype=bool))
calib['1_K'].append(calib_data[seqName][camIdx]['K'].astype(np.float32))
calib['2_K'].append(calib_data[seqName][camIdx]['K'].astype(np.float32))
calib['1_R'].append(calib_data[seqName][camIdx]['R'].astype(np.float32))
calib['2_R'].append(calib_data[seqName][camIdx]['R'].astype(np.float32))
calib['1_t'].append(calib_data[seqName][camIdx]['t'][:, 0].astype(np.float32))
calib['2_t'].append(calib_data[seqName][camIdx]['t'][:, 0].astype(np.float32))
calib['1_distCoef'].append(calib_data[seqName][camIdx]['distCoef'].astype(np.float32))
calib['2_distCoef'].append(calib_data[seqName][camIdx]['distCoef'].astype(np.float32))
img_dirs_1.append('{}/{}/hdImgs/{}/{}/00_{:02d}_{}.jpg'.format(self.image_root, 'a4', seqName, frame_str, camIdx, frame_str))
img_dirs_2.append('{}/{}/hdImgs/{}/{}/00_{:02d}_{}.jpg'.format(self.image_root, 'a4', seqName, next_data['frame_str'], camIdx, next_data['frame_str']))
human3d.update(calib)
human3d['1_img_dirs'] = img_dirs_1
human3d['2_img_dirs'] = img_dirs_2
self.register_tensor(human3d, order_dict)
self.num_samples = len(self.tensor_dict['1_img_dirs'])