列出了三段代码，在输出文件夹中创建4个空 tfrecord文件后，该过程挂起。 该代码假定读取一个元组列表（video_path，label），对于每个元组，我们读取该视频和标签数据的帧，进行一些处理，然后写入tfrecord。我想使用多处理来加速它，因为元组实例彼此独立。
def add_to_tfrecord(video_path, info_path, save_path, coder, bad_file_logger):
frames = video2frames(video_path)
pose, shape, gender, joint3d, joint2d, zrot, camLoc = load_mat(info_path)
# video path : */run0/*/104_30_c0003.mp4
# print('length of frames', len(frames))
projFunc = proj3d2(camLoc)
image_shape = frames[0].shape
if gender ==0:
(model, pelvis_loc) = (f_model, f_pelvis_pos)
else:
(model, pelvis_loc) = (m_model, m_pelvis_pos)
with tf.python_io.TFRecordWriter(save_path) as writer:
for i in range(len(frames)):
trans = joint3d[i, 0] - pelvis_loc
pose[i,:3] = rotateBody(zrot, pose[i,:3])
model.pose[:] = pose[i]
model.betas[:] = shape[i]
model.trans[:] = trans
verts = model.r
face_verts_3d = np.matmul(face_regressor, verts) # (5x6890) (6890x3)
face_verts_2d = projFunc(face_verts_3d)
# img_data = coder.encode_jpeg(frames[i])
#NOTICE, to make consistent with dataloader, there should be only 14 gt 3d joints
# joint3d_i = np.concatenate( (joint3d[i][ [8, 5, 2, 1, 4, 7, 21, 19, 17, 16, 18, 20, 12, 15] ], face_verts_3d), axis=0 )
joint3d_i = joint3d[i][ [8, 5, 2, 1, 4, 7, 21, 19, 17, 16, 18, 20, 12, 15] ]
joint2d_i = np.concatenate( (joint2d[i][ [8, 5, 2, 1, 4, 7, 21, 19, 17, 16, 18, 20, 12, 15] ], face_verts_2d), axis=0 )
joint3d_i = np.matmul(R_world2cv, joint3d_i.T).T
label2d_i, vis = vis_annot(joint2d_i, image_shape)
label2d_i[13,2] = 0
vis[13] = 0
if np.sum(vis)<10:
bad_file = '%s %d / %d'%(video_path, i, len(frames))
print('Bad! ' + bad_file+'\r')
bad_file_logger.write( bad_file + '\n' )
else:
min_pt = np.min(joint2d_i[vis], axis=0)
max_pt = np.max(joint2d_i[vis], axis=0)
center = (min_pt + max_pt) / 2.
person_height = np.linalg.norm(max_pt - min_pt)
scale = 150. / person_height
image_scaled, scale_factors = resize_img(frames[i], scale)
height, width = image_scaled.shape[:2]
joints_scaled = np.copy(label2d_i)
joints_scaled[:, 0] *= scale_factors[0]
joints_scaled[:, 1] *= scale_factors[1]
center_scaled = np.round(center * scale_factors).astype(np.int)
# # scale camera: Flength, px, py
cam_scaled = [600, 160, 120]
cam_scaled[0] = scale # TODO: think about it and return sth more useful
cam_scaled[1] *= scale_factors[0]
cam_scaled[2] *= scale_factors[1]
# Crop 300x300 around the center
margin = 150
start_pt = np.maximum(center_scaled - margin, 0).astype(int)
end_pt = (center_scaled + margin).astype(int)
end_pt[0] = min(end_pt[0], width)
end_pt[1] = min(end_pt[1], height)
image_scaled = image_scaled[start_pt[1]:end_pt[1], start_pt[0]:end_pt[
0], :]
# Update others too.
joints_scaled[:, 0] -= start_pt[0]
joints_scaled[:, 1] -= start_pt[1]
center_scaled -= start_pt
# Update principal point:
cam_scaled[1] -= start_pt[0]
cam_scaled[2] -= start_pt[1]
height, width = image_scaled.shape[:2]
# disp2dp(image_scaled, joints_scaled)
joints_scaled_vis, vis_u = vis_annot(joints_scaled, (height, width))
if np.sum(vis_u)<10:
bad_file = '%s %d / %d'%(video_path, i, len(frames))
print('Bad! ' + bad_file+'\r')
bad_file_logger.write( bad_file + '\n' )
continue
img_data = coder.encode_jpeg(image_scaled)
example = convert_to_surreal_example(img_data, video_path + '_%d'%i, height=height,
width=width, label=joints_scaled, center=center_scaled, gt3d=joint3d_i, pose=pose[i], shape=shape[i],
trans=trans, gender=gender, cam=np.array(cam_scaled))
writer.write(example.SerializeToString())
bad_file_logger.flush()
def worker(task, save_dir):
# # video path : */run0/*/104_30_c0003.mp4
video_path, info_path = task
name = video_path.split('/')[-3] + '_' + video_path.split('/')[-1][:-4] + '.tfrecord'
save_path = os.path.join(save_dir, name)
try:
add_to_tfrecord_mp(video_path, info_path, save_path, coder)
except ValueError:
print(video_path)
pass
return 0
def process_surreal_mp(in_path, out_path):
import multiprocessing as mp
sets = ['train', 'test', 'val']
for set in sets:
save_dir = os.path.join(out_path, set)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
data_path = os.path.join(in_path, set+'.txt')
video_clip_infos = read_file(data_path) # return list of tuples (video_path, gt_mat_path)
pool = mp.Pool(4)
result = [pool.apply_async(worker, args=(video_clip_info, save_dir) ) for video_clip_info in video_clip_infos]# , chunksize=100)
pool.close()
pool.join()