def data_loader_for_inference(args=None, isSave=False, isNormalized=True):
depth_dict = _load_anno_mat(args.depth_fname)
info = _load_anno_mat(args.info_fname)
## Load SMPL model (here we load the female model)
## Make sure path is correct
# <========= LOAD SMPL MODEL BASED ON GENDER
if info['gender'][0] == 0: # f
pkl_fname = './models/basicModel_f_lbs_10_207_0_v1.0.0.pkl'
elif info['gender'][0] == 1: # m
pkl_fname = './models/basicModel_m_lbs_10_207_0_v1.0.0.pkl'
else:
pkl_fname = './models/neutral_smpl_with_cocoplus_reg.pkl'
m = load_model(pkl_fname)
print("[***] loading model %s" % pkl_fname)
#print ("[****] initial pose = {}, shape = {}, trans = {}, J_regressor shape = {}, J_regressor[0,:] = {}".format(
# m.pose[:6], m.betas[:], m.trans[:], m.J_regressor, m.J_regressor.shape))
root_pos = m.J_transformed.r[0]
zrot = info['zrot']
print("info['zrot'].shape = {}".format(zrot.shape))
zrot = zrot[0] # body rotation in euler angles
RzBody = np.array(((math.cos(zrot), -math.sin(zrot), 0),
(math.sin(zrot), math.cos(zrot), 0), (0, 0, 1)))
# add gender info: it could be female, male, or neutral;
gender = info['gender'][0]
if gender == 0:
gender_vect = np.array([1.0, .0, .0
]).astype(np.float32) # gender vector - female
elif gender == 1:
gender_vect = np.array([.0, 1.0,
.0]).astype(np.float32) # gender vector - male
else:
gender_vect = np.array([.0, .0, 1.0]).astype(
np.float32) # gender vector - neutral
intrinsic, cam = surreal_util.get_intrinsic()
extrinsic = surreal_util.get_extrinsic(np.reshape(info['camLoc'], [3, -1]))
all_gt_poses = []
all_gt_shapes = []
all_gt_genders = [] # added gender info;
all_images = []
all_depths = []
all_gt_joints2d = []
all_gt_joints3d = []
all_gt_trans = [
] # added for translation between: info['joints3D'][:,:,t].T[0] - root_pos
all_gt_smpl_verts = []
all_gt_smpl_joints3d = []
all_gt_smpl_joints3d_proj = []
gt_cam = cam
if not os.path.exists(args.result_dir):
os.makedirs(args.result_dir)
print("makedirs %s" % args.result_dir)
for t in range(args.t_beg, args.t_end):
# depth: in meters;
dep = depth_dict['depth_%d' % (
t + 1
)] # the depth_1, index is 1-based, defiend by the dataset itselt;
dep = surreal_util.normalizeDepth(dep, isNormalized)
# Read frame of the video
rgb = surreal_util.get_frame(args.info_fname[:-9] + ".mp4", t)
gt_poses = info['pose'][:, t]
gt_shape = info['shape'][:, t]
gt_poses[0:3] = surreal_util.rotateBody(RzBody, gt_poses[0:3])
gt_joints2d = info['joints2D'][:, :, t] # (x,y,vis) x 24;
gt_joints3d = info['joints3D'][:, :, t] # (x,y,z) x 24;
#NOTE: surreal joints are left/right swapped, compared with the smpl joint order;
gt_joints2d_swap = surreal_util.swap_right_left_joints(gt_joints2d)
gt_joints3d_swap = surreal_util.swap_right_left_joints(gt_joints3d)
# read 14 lsp joints;
gt_joints2d, gt_joints3d = surreal_util.read_joints(
gt_joints2d_swap, gt_joints3d_swap, extrinsic) # 3 x 14;
#NOTE:still, surreal dataset has right/left swapped pose compared with official smpl;
gt_poses = surreal_util.swap_right_left_pose(gt_poses)
# all the frames of the same subject should have the same gender, so just repeat it;
all_gt_genders.append(gender_vect)
all_gt_joints2d.append(gt_joints2d[:2, :]) # (x,y,vis) x 24;
all_gt_joints3d.append(gt_joints3d / 1000.0) # change back to meters;
all_gt_poses.append(gt_poses)
all_gt_shapes.append(gt_shape)
all_images.append(rgb)
all_depths.append(dep)
m.pose[:] = gt_poses
m.betas[:] = gt_shape
# Set model translation
""" actually this m.trans[:] can also be added to the smpl vertices/joints
here or after posing/shaping, because of the translation is independent
of shaping/posing!!!
"""
m.trans[:] = info['joints3D'][:, :, t].T[0] - root_pos
all_gt_trans.append(info['joints3D'][:, :, t].T[0] - root_pos)
print("[****] updated smpl pose = {}, shape = {}, trans = {}".format(
m.pose[:6], m.betas[:], m.trans[:]))
print("[****] m.trans[:] = {}".format(info['joints3D'][:, :, t].T[0] -
root_pos))
print("[****] initial, root_joint_3D = {}, root_joint_from_pos = {}".
format(info['joints3D'][:, :, t].T[0], root_pos))
smpl_vertices = m.r
smpl_joints3D = m.J_transformed.r
print("[****] after posing, root_joint_from_pose = {}".format(
smpl_joints3D[0]))
#"""""""""""""""""
#NOTE: debugging
#surreal_util.draw_joints2D(
# rgb,
# # 24 joints
# surreal_util.project_vertices(smpl_joints3D, intrinsic, extrinsic),
# None, m.kintree_table, color = 'b')
#print ("smpl_joints3D shape = ", smpl_joints3D.shape) # in shape [24, 3]
smpl_joints3D = extract_14_joints(
smpl_joints3D, surreal_util.get_lsp_idx_from_smpl_joints(
)) # now in shape [14, 3]
#print ("smpl_joints3D shape = ", smpl_joints3D.shape) # now in shape [14, 3]
all_gt_smpl_verts.append(smpl_vertices)
all_gt_smpl_joints3d.append(smpl_joints3D)
#print ("intrinsic shape = {}, extrinsic shape = {}".format(intrinsic.shape, extrinsic.shape))
#print ("[????] intrinsic in py smpl = {}".format(intrinsic))
#print ("[????] extrinsic in py smpl = {}".format(extrinsic))
proj_smpl_joints3D = surreal_util.project_vertices(
smpl_joints3D, intrinsic, extrinsic)
all_gt_smpl_joints3d_proj.append(proj_smpl_joints3D)
if isSave:
## Write to an .obj file
tmp_name = args.info_fname.split('/')[-1][:-9] + "frame_%03d" % t
outmesh_path = os.path.join(args.result_dir,
tmp_name + '_smpl.obj')
outimg_path = os.path.join(args.result_dir, tmp_name + "_rgb.jpg")
outdep_path = os.path.join(args.result_dir,
tmp_name + "_depth.pfm")
cv2.imwrite(outimg_path, rgb[:, :, [2, 1, 0]])
pfm.save(outdep_path, dep.astype(np.float32))
with open(outmesh_path, 'w') as fp:
for v in m.r:
fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
for f in m.f + 1: # Faces are 1-based, not 0-based in obj files
fp.write('f %d %d %d\n' % (f[0], f[1], f[2]))
## Print message
print('..Output mesh saved to: ', outmesh_path)
print('..Output depth saved to: ', outdep_path)
print('..Output rgb image saved to: ', outimg_path)
return all_images, all_depths, all_gt_poses, all_gt_shapes, \
all_gt_genders, all_gt_trans,\
all_gt_joints2d, all_gt_joints3d, \
all_gt_smpl_verts, all_gt_smpl_joints3d, \
all_gt_smpl_joints3d_proj, gt_cam, intrinsic, extrinsic, m.kintree_table
cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
# offset(i.e., distance between stereo views): B = 0.532725 m = 53.2725 cm
B = 53.2725 # in centimeters;
f = 725.0087 # in pixels
# set zero as a invalid disparity value;
disp_left = np.zeros([height, width], 'float32')
mask = depth_left > 0
disp_left[mask] = f * B / depth_left[mask] # d = fB/z
pfm.show(disp_left, title="disp_left")
disp_pfm_filename = pjoin(
file_path, "vkitti_2.0.3_disparity/" + A[:-26] +
'disparity/Camera_0/disparity_' + A[-9:-4] + ".pfm")
print("disp_pfm_filename = ", disp_pfm_filename)
tmp_len = disp_pfm_filename.rfind("/")
tmp_dir = disp_pfm_filename[0:tmp_len]
print("tmp_dir = ", tmp_dir)
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
pfm.save(disp_pfm_filename, disp_left)
#loadding semantic segmantation labels
seg_filename = pjoin(
file_path, "vkitti_2.0.3_classSegmentation/" + A[:-26] +
'classSegmentation/Camera_0/classgt_' + A[-9:-4] + ".png")
print("semantic label: {}".format(seg_filename))
semantic_rgb_label = cv2.imread(
seg_filename)[:, :, ::-1] # change BRG to RGB via ``::-1`;
pfm.show_uint8(semantic_rgb_label, title="semantic_rgb_label")
semantic_label = encode_color_segmap(semantic_rgb_label)
pfm.show(semantic_label, title="semantic_label")
def test(self, py=None):
self.model.eval()
if not os.path.exists(self.args.resultDir):
os.makedirs(self.args.resultDir)
print('makedirs {}'.format(self.args.resultDir))
dispScale = 1.0
avg_err = .0
avg_rate = .0
img_num = self.test_loader_len
print("[****] To test %d images !!!" % img_num)
for iteration, batch_data in enumerate(self.testing_data_loader):
batch_size = batch_data[0].size()[0]
assert batch_size == 1
features = batch_data[0]
print("[???] features size:", features.size())
height = batch_data[1][0].item()
width = batch_data[2][0].item()
crop_height = batch_data[3][0].item()
crop_width = batch_data[4][0].item()
current_file = batch_data[5][0]
disp_name = batch_data[6][0]
print(height, crop_height, current_file, disp_name)
if os.path.isfile(disp_name):
dispGT = pfm.readPFM(disp_name)
dispGT[dispGT == np.inf] = .0
else:
dispGT = None
if self.kitti2015 or self.kitti2012:
savename = pjoin(self.args.resultDir,
current_file[0:-4] + '.pfm')
elif self.eth3d or self.middlebury:
savename = pjoin(self.args.resultDir, current_file + '.pfm')
else: # scene flow dataset
savename = pjoin(self.args.resultDir, str(iteration) + '.pfm')
if self.cuda:
features = features.cuda()
with torch.no_grad():
if self.model_name == 'MS-GCNet':
disp = self.model(features)
elif self.model_name == 'MS-PSMNet':
disp = self.model(features)
else:
raise Exception(
"No suitable model found ... Wrong name: {}".format(
self.model_name))
#about memory
del features
disp = disp.cpu().detach().numpy() * dispScale
if height <= crop_height and width <= crop_width:
#disp = disp[0, crop_height - height: crop_height, crop_width-width: crop_width]
disp = disp[0, crop_height - height:crop_height, 0:width]
else:
disp = disp[0, :, :]
#save to uint16 png files
#skimage.io.imsave(savename, (disp * 256).astype('uint16'))
if any([
self.kitti2015, self.kitti2012, self.eth3d,
self.middlebury, iteration % 50 == 0
]):
pfm.save(savename, disp)
#print ('saved ', savename)
if dispGT is not None:
if self.eth3d:
self.args.threshold = 1.0
elif self.middlebury:
self.args.threshold = 1.0 # for trainingH;
elif self.kitti2012 or self.kitti2015:
self.args.threshold = 3.0
else: # Scene Flow
self.args.threshold = 1.0
error, rate = get_epe_rate(dispGT, disp, self.max_disp,
self.args.threshold)
avg_err += error
avg_rate += rate
if iteration % 200 == 0:
message_info = "===> Frame {}: ".format(
iteration
) + current_file + " ==> EPE Error: {:.4f}, Bad-{:.1f} Error: {:.4f}".format(
error, self.args.threshold, rate)
""" adapted from Mateo's code """
if py is not None:
gc.collect()
memoryUse = py.memory_info(
)[0] / 2.**20 # memory use in MB...I think
message_info += ', memeory: {:.2f} MB'.format(
memoryUse)
print(message_info)
# save kt15 color
#if self.kitti2015:
if any(
[self.kitti2015, self.kitti2012, self.eth3d, self.middlebury]):
""" disp """
tmp_dir = pjoin(self.args.resultDir, "dispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
tmp_dispname = pjoin(tmp_dir, current_file[0:-4] + '.png')
cv2.imwrite(
tmp_dispname,
KT15FalseClr.writeKT15FalseColor(
np.ascontiguousarray(disp)).astype(
np.uint8)[:, :, ::-1])
if iteration % 50 == 0:
print('saved ', tmp_dispname)
if dispGT is not None: #If KT benchmark submission, then No dispGT;
""" err-disp """
tmp_dir = pjoin(self.args.resultDir, "errDispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
tmp_errdispname = pjoin(tmp_dir,
current_file[0:-4] + '.png')
cv2.imwrite(
tmp_errdispname,
KT15LogClr.writeKT15ErrorDispLogColor(
np.ascontiguousarray(disp),
np.ascontiguousarray(dispGT)).astype(
np.uint8)[:, :, ::-1])
if iteration % 50 == 0:
print('saved ', tmp_errdispname)
if dispGT is not None:
avg_err = avg_err / img_num
avg_rate = avg_rate / img_num
print(
"===> Total {} Frames ==> AVG EPE Error: {:.4f}, AVG Bad-{:.1f} Error: {:.4f}"
.format(img_num, avg_err, self.args.threshold, avg_rate))
print('{} testing finished!'.format(self.model_name))
def test(self):
self.model.eval()
file_path = self.args.data_path
file_list_txt = self.args.test_list
f = open(file_list_txt, 'r')
if self.virtual_kitti2:
filelist = get_virtual_kitti2_filelist(file_list_txt)
else:
filelist = [l.rstrip() for l in f.readlines() if not l.rstrip().startswith('#')]
crop_width = self.args.crop_width
crop_height = self.args.crop_height
if not os.path.exists(self.args.resultDir):
os.makedirs(self.args.resultDir)
print ('makedirs {}'.format(self.args.resultDir))
avg_err = 0
avg_rate1 = 0
avg_rate3 = 0
for index in range(len(filelist)):
#for index in range(1):
current_file = filelist[index]
if self.kitti2015:
leftname = pjoin(file_path, 'image_0/' + current_file)
if index % 20 == 0:
print ("limg: {}".format(leftname))
rightname = pjoin(file_path, 'image_1/' + current_file)
dispname = pjoin(file_path, 'disp_occ_0_pfm/' + current_file[0:-4] + '.pfm')
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT= None
savename = pjoin(self.args.resultDir, current_file[0:-4] + '.pfm')
elif self.kitti2012:
#leftname = pjoin(file_path, 'image_0/' + current_file)
leftname = pjoin(file_path, 'colored_0/' + current_file)
rightname = pjoin(file_path, 'colored_1/' + current_file)
dispname = pjoin(file_path, 'disp_occ_0_pfm/' + current_file[0:-4] + '.pfm')
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT= None
savename = pjoin(self.args.resultDir, current_file[0:-4] + '.pfm')
#disp = Image.open(dispname)
#disp = np.asarray(disp) / 256.0
elif self.virtual_kitti2:
A = current_file
# e.g., /media/ccjData2/datasets/Virtual-KITTI-V2/vkitti_2.0.3_rgb/Scene01/15-deg-left/frames/rgb/Camera_0/rgb_00001.jpg
leftname = pjoin(file_path, "vkitti_2.0.3_rgb/" + A)
rightname = pjoin(file_path, "vkitti_2.0.3_rgb/" + A[:-22] + 'Camera_1/' + A[-13:])
#load depth GT and change it to disparity GT:
depth_png_filename = pjoin(file_path, "vkitti_2.0.3_depth/" + A[:-26] + 'depth/Camera_0/depth_' + A[-9:-4] + ".png")
#print ("imgl = ", leftname, ", imgr = ", rightname, ", depth_left = ", depth_png_filename)
#NOTE: The depth map in centimeters can be directly loaded
depth_left = cv2.imread(depth_png_filename, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
height, width = depth_left.shape[:2]
# Intrinsi: f_x = f_y = 725.0087
# offset(i.e., distance between stereo views): B = 0.532725 m = 53.2725 cm
B = 53.2725 # in centimeters;
f = 725.0087 # in pixels
# set zero as a invalid disparity value;
dispGT = np.zeros([height, width], 'float32')
mask = depth_left > 0
dispGT[mask] = f*B/ depth_left[mask] # d = fB/z
#pfm.show(dispGT, title='dispGT')
savename = pjoin(self.args.resultDir, '%04d.pfm'%(index))
elif self.eth3d or self.middlebury:
leftname = pjoin(file_path, current_file + '/im0.png')
rightname = pjoin(file_path, current_file + '/im1.png')
print ("limg: {}".format(leftname))
H, W = cv2.imread(leftname).shape[:2]
dispname = pjoin(file_path, current_file + '/disp0GT.pfm')
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT= None
savename = pjoin(self.args.resultDir, current_file + '.pfm')
# due to eth3d images have diverse image size;
# so we process each of them seperately;
ds = self.ds # e.g., == 64, encoder sacle;
crop_width = W + (ds - W % ds) % ds
crop_height = H + (ds - H % ds) % ds
print ("input HxW=%dx%d, padding size=%dx%d" %(H,W,crop_height, crop_width))
else:
A = current_file
leftname = pjoin(file_path, A)
rightname = pjoin(file_path, A[:-13] + 'right/' + A[len(A)-8:])
# check disparity GT exists or not!!!
pos = A.find('/')
tmp_len = len('frames_finalpass')
dispname = pjoin(file_path, A[0:pos] + '/disparity' + A[pos+1+tmp_len:-4] + '.pfm')
#print ("[****] ldisp: {}".format(dispname))
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT= None
savename = pjoin(self.args.resultDir, str(index) + '.pfm')
input1, input2, height, width = test_transform(
load_test_data(leftname, rightname, is_data_augment=self.is_data_augment),
crop_height, crop_width)
if self.cuda:
input1 = input1.cuda()
input2 = input2.cuda()
with torch.no_grad():
prediction = self.model(input1, input2)
disp = prediction.cpu().detach().numpy() # in full size
if height <= crop_height and width <= crop_width:
disp = disp[0, crop_height - height: crop_height, crop_width-width: crop_width]
else:
disp = disp[0, :, :]
# for KT benchmark submission;
if dispGT is None:
tmp_savename = savename[:-4] + '.png'
print ("savd ", tmp_savename)
skimage.io.imsave(tmp_savename, (disp * 256).astype('uint16'))
#sys.exit()
if any([self.kitti2015, self.kitti2012, self.eth3d, self.middlebury, index %250 == 0]):
pfm.save(savename, disp)
#print ('saved ', savename)
if 0 and dispGT is not None:
left = np.asarray(Image.open(leftname))[:,:,:3].astype(np.float32)
right = np.asarray(Image.open(rightname))[:,:,:3].astype(np.float32)
#print ("[???]", left.shape)
pfm.save(savename[:-4] + '-iml.pfm', left)
pfm.save(savename[:-4] + '-imr.pfm', right)
pfm.save(savename[:-4] + '-gt.pfm', dispGT.astype(np.float32))
if dispGT is not None:
#error, rate1, rate3 = get_epe_rate2(dispGT, disp, self.args.max_disp, threshold=1.0, threshold2=3.0)
error, rate1, rate3 = get_epe_rate(dispGT, disp, threshold=1.0, threshold2=3.0)
avg_err += error
avg_rate1 += rate1
avg_rate3 += rate3
if index % 250 == 0:
print("===> Frame {}: ".format(index) + leftname + " ==> EPE Error: {:.4f}, Bad-{:.1f} Error: {:.4f}, Bad-{:.1f} Error: {:.4f}".format(
error, 1.0, rate1, 3.0, rate3))
# save kt15 color
#if self.kitti2015:
if any([self.kitti2015, self.kitti2012, self.eth3d, self.middlebury, index%250==0]):
""" disp """
tmp_dir = pjoin(self.args.resultDir, "dispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
if self.eth3d or self.middlebury:
tmp_dispname = pjoin(tmp_dir, current_file + '.png')
elif self.kitti2012 or self.kitti2015:
tmp_dispname = pjoin(tmp_dir, current_file[0:-4] + '.png')
else:
tmp_dispname = pjoin(tmp_dir, '%04d.png'%(index))
cv2.imwrite(tmp_dispname,
KT15FalseClr.writeKT15FalseColor(np.ascontiguousarray(disp)).astype(np.uint8)[:,:,::-1])
if index < 1:
print ('saved ', tmp_dispname)
if dispGT is not None: #If KT benchmark submission, then No dispGT;
""" err-disp """
tmp_dir = pjoin(self.args.resultDir, "errDispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
if self.eth3d or self.middlebury:
tmp_errdispname = pjoin(tmp_dir, current_file + '.png')
elif self.kitti2015 or self.kitti2012:
tmp_errdispname = pjoin(tmp_dir, current_file[0:-4] + '.png')
else:
tmp_errdispname = pjoin(tmp_dir, '%04d.png'%(index))
cv2.imwrite(tmp_errdispname,
KT15LogClr.writeKT15ErrorDispLogColor(np.ascontiguousarray(disp), np.ascontiguousarray(dispGT)).astype(np.uint8)[:,:,::-1])
if index < 1:
print ('saved ', tmp_errdispname)
if dispGT is not None:
avg_err /= len(filelist)
avg_rate1 /= len(filelist)
avg_rate3 /= len(filelist)
print("===> Total {} Frames ==> AVG EPE Error: {:.4f}, AVG Bad-{:.1f} Error: {:.4f}, AVG Bad-{:.1f} Error: {:.4f}".format(
len(filelist), avg_err, 1.0, avg_rate1, 3.0, avg_rate3))
""" save as csv file, Excel file format """
csv_file = os.path.join(self.args.resultDir, 'bad-err.csv')
print ("write ", csv_file, "\n")
timeStamp = datetime.now().strftime('%Y-%m-%d_%H:%M:%S')
messg = timeStamp + ',{},bad-1.0,{:.4f},bad-3.0,{:.4f},epe,{:.4f},fileDir={},for log,{:.3f}(epe); {:.3f}%(bad1); {:.3f}%(bad3)\n'.format(
self.data_type_str, avg_rate1, avg_rate3, avg_err,
self.args.resultDir,
avg_err, avg_rate1*100.0, avg_rate3*100.0)
with open( csv_file, 'w') as fwrite:
fwrite.write(messg)
tmp_len = len('frames_finalpass')
dispname = pjoin(
file_path,
A[0:pos] + '/disparity' + A[pos + 1 + tmp_len:-4] + '.pfm')
#print ("[****] ldisp: {}".format(dispname))
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT = None
savename = pjoin(opt.resultDir, str(index) + '.pfm')
disp = test(leftname, rightname, savename)
if any([opt.kitti2015, opt.kitti2012, index % 250 == 0]):
pfm.save(savename, disp)
""" disp """
tmp_dir = pjoin(opt.resultDir, "dispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
# save kt15 color
if opt.kitti2015 or opt.kitti2012:
tmp_dispname = pjoin(tmp_dir, current_file[0:-4] + '.png')
else:
tmp_dispname = pjoin(tmp_dir, '%04d.png' % (index))
cv2.imwrite(
tmp_dispname,
KT15FalseClr.writeKT15FalseColor(
np.ascontiguousarray(disp)).astype(np.uint8)[:, :, ::-1])
print("save ", tmp_dispname)
def add_to_tfrecord_file(frame_fname, image, depth, joints2d, joints3d, pose, shape,
writer,
cam,
is_save_to_h5files = False,
coder = None,
isSavepfm = False,
out_dir = '',
gender_vect = np.array([.0,.0,1.0]).astype(np.float32) # means 'nuetral'
):
#NOTE: this cv2 i/o makes the tfrecord files slow !!! ???
image_scaled, depth_scaled, joints_scaled, height, width, cam_scaled, \
start_pt, scale_factors, center_scaled = parse_joints_and_scale_image(image,depth,joints2d, cam)
if image_scaled is None:
#print ("[**] probabily no person apper in this frame. with depth shape : ", depth.shape, "image shape : ", image.shape, "depth_scaled shape :", depth_scaled.shape, "image_scaled shape : ", image_scaled.shape, 'frame_fname : ', frame_fname)
return 1, None
#NOTE: saved the values in millimeter as uint16 type;
depth_scaled *= 1000
#depth_scaled[depth_scaled >=65535 ] = 65535 # uint16 max valueose, shape, scale_factors, start_pt, cam_scaled)
depth_scaled[depth_scaled >=65535 ] = 0
depth_scaled = depth_scaled.astype(np.uint16)
#print ("[**] scalded depth shape : ", depth_scaled.shape, "scaled image shape : ", image_scaled.shape)
if isSavepfm:
#e.g., frame_fname = ./datasets/surreal/cmu/val/run0/03_01/03_01_c0008_frame_000;
print ("our_dir = ", out_dir)
tmp_name = frame_fname.split("/")[-1]
tmp_name_img = join(out_dir, "img-scale-{}.pfm".format(tmp_name))
pfm.save(tmp_name_img, image_scaled.astype(np.float32))
print (" save pfm : %s" %tmp_name_img)
tmp_name_dep = join(out_dir, "dep-scale-{}.pfm".format(tmp_name))
pfm.save(tmp_name_dep, depth_scaled.astype(np.float32))
print (" save pfm : %s" % tmp_name_dep)
isSavepfm = False
# Encode image to jpeg format:
image_data_scaled = coder.encode_jpeg(image_scaled)
# NOTE: Encode depth to raw or string or byte format:
depth_data_scaled = depth_scaled.tobytes()
#NOTE: here we have make sure that joints3d is in shape [14, 3];
# which is required by tfrecord example write and parse;
assert (joints3d.shape[1] == 3)
if is_save_to_h5files:
tmp_img_name = frame_fname
example = convert_to_h5_wmosh_wdepth(
depth_scaled.astype(np.float32),
image_scaled.astype(np.float32),
tmp_img_name, height, width, joints_scaled, center_scaled,
joints3d, pose, shape, scale_factors,
start_pt, cam_scaled, gender_vect)
else:
example = convert_to_example_wmosh_wdepth(depth_data_scaled, image_data_scaled, frame_fname,
height, width, joints_scaled, center_scaled, joints3d, pose,
shape, scale_factors, start_pt, cam_scaled, gender_vect )
writer.write(example.SerializeToString())
# Finally return how many were written.
return 1, example # 1 person added
def test(self):
self.model.eval()
file_path = self.args.data_path
file_list = self.args.test_list
f = open(file_list, 'r')
filelist = [l.rstrip() for l in f.readlines()]
crop_width = self.args.crop_width
crop_height = self.args.crop_height
if not os.path.exists(self.args.resultDir):
os.makedirs(self.args.resultDir)
print('makedirs {}'.format(self.args.resultDir))
avg_err = 0
avg_rate = 0
for index in range(len(filelist)):
#for index in range(1):
current_file = filelist[index]
if self.kitti2015:
leftname = pjoin(file_path, 'image_0/' + current_file)
if index % 20 == 0:
print("limg: {}".format(leftname))
rightname = pjoin(file_path, 'image_1/' + current_file)
dispname = pjoin(
file_path, 'disp_occ_0_pfm/' + current_file[0:-4] + '.pfm')
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT = None
savename = pjoin(self.args.resultDir,
current_file[0:-4] + '.pfm')
elif self.kitti2012:
#leftname = pjoin(file_path, 'image_0/' + current_file)
leftname = pjoin(file_path, 'colored_0/' + current_file)
rightname = pjoin(file_path, 'colored_1/' + current_file)
dispname = pjoin(
file_path, 'disp_occ_0_pfm/' + current_file[0:-4] + '.pfm')
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT = None
savename = pjoin(self.args.resultDir,
current_file[0:-4] + '.pfm')
#disp = Image.open(dispname)
#disp = np.asarray(disp) / 256.0
else:
A = current_file
leftname = pjoin(file_path, A)
rightname = pjoin(file_path,
A[:-13] + 'right/' + A[len(A) - 8:])
# check disparity GT exists or not!!!
pos = A.find('/')
tmp_len = len('frames_finalpass')
dispname = pjoin(
file_path,
A[0:pos] + '/disparity' + A[pos + 1 + tmp_len:-4] + '.pfm')
#print ("[****] ldisp: {}".format(dispname))
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT = None
savename = pjoin(self.args.resultDir, str(index) + '.pfm')
input1, input2, height, width = test_transform(
load_test_data(leftname, rightname), crop_height, crop_width)
if self.cuda:
input1 = input1.cuda()
input2 = input2.cuda()
with torch.no_grad():
prediction1, prediction = self.model(input1, input2)
disp = prediction.cpu().detach().numpy() # in full size
disp1 = prediction1.cpu().detach().numpy() # in half size
if height <= crop_height and width <= crop_width:
disp = disp[0, crop_height - height:crop_height,
crop_width - width:crop_width]
disp1 = disp1[0,
crop_height // 2 - height // 2:crop_height // 2,
crop_width // 2 - width // 2:crop_width // 2]
else:
disp = disp[0, :, :]
disp1 = disp1[0, :, :]
#skimage.io.imsave(savename, (disp * 256).astype('uint16'))
if self.kitti2015 or self.kitti2012 or index % 50 == 0:
pfm.save(savename, disp)
#print ('saved ', savename)
if 0 and dispGT is not None:
left = np.asarray(Image.open(leftname))[:, :, :3].astype(
np.float32)
right = np.asarray(Image.open(rightname))[:, :, :3].astype(
np.float32)
#print ("[???]", left.shape)
pfm.save(savename[:-4] + '-iml.pfm', left)
pfm.save(savename[:-4] + '-imr.pfm', right)
pfm.save(savename[:-4] + '-gt.pfm',
dispGT.astype(np.float32))
pfm.save(savename[:-4] + '-half.pfm', disp1)
if dispGT is not None:
error, rate = get_epe_rate(dispGT, disp, self.args.max_disp,
self.args.threshold)
avg_err += error
avg_rate += rate
if index % 20 == 0:
print("===> Frame {}: ".format(index) + leftname +
" ==> EPE Error: {:.4f}, Bad-{:.1f} Error: {:.4f}".
format(error, self.args.threshold, rate))
# save kt15 color
if self.kitti2015:
""" disp """
tmp_dir = pjoin(self.args.resultDir, "dispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
tmp_dispname = pjoin(tmp_dir, current_file[0:-4] + '.png')
cv2.imwrite(
tmp_dispname,
KT15FalseClr.writeKT15FalseColor(
np.ascontiguousarray(disp)).astype(
np.uint8)[:, :, ::-1])
if index % 20 == 0:
print('saved ', tmp_dispname)
if dispGT is not None: #If KT benchmark submission, then No dispGT;
""" err-disp """
tmp_dir = pjoin(self.args.resultDir, "errDispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
tmp_errdispname = pjoin(tmp_dir,
current_file[0:-4] + '.png')
cv2.imwrite(
tmp_errdispname,
KT15LogClr.writeKT15ErrorDispLogColor(
np.ascontiguousarray(disp),
np.ascontiguousarray(dispGT)).astype(
np.uint8)[:, :, ::-1])
if index % 20 == 0:
print('saved ', tmp_errdispname)
if dispGT is not None:
avg_err = avg_err / len(filelist)
avg_rate = avg_rate / len(filelist)
print(
"===> Total {} Frames ==> AVG EPE Error: {:.4f}, AVG Bad-{:.1f} Error: {:.4f}"
.format(len(filelist), avg_err, self.args.threshold, avg_rate))
def infer_eval_1_img(
model_type,
img,
dep,
gender_vect,
joints3d_gt_with_torso_subtracted,
model=None,
result_dir='',
img_name='',
smpl_model_path=None,
isSaveResults=False,
meter_to_mm=1.0 # could be 1000.0 or 1.0, according to the unit you used;
):
#*********************
""" 1) prediction """
#*********************
# Add batch dimension: 1 x D x D x 3
input_img, input_dep, proc_param, img_orig, _, _ = preprocess_image(
img, dep, None, None, None)
input_img = np.expand_dims(input_img, 0)
input_dep = np.expand_dims(input_dep, 0)
gt_gender = np.expand_dims(gender_vect, 0)
if model_type == 'hmr':
joints, verts, cams, joints3d, _, _, _, _, theta, beta, _ = model.predict(
input_img, None, None, None, get_theta=True)
else:
joints, verts, cams, joints3d, _, _, _, _, theta, beta, joints3d_24 = model.predict(
input_img, input_dep, None, None, None, gt_gender, get_theta=True)
#*********************
""" 2) evaluation """
#*********************
#print ("[***] evaluating MPJPE and MPJPE PA !")
joints3d_aligned = align_by_pelvis(joints3d[0], get_pelvis=False)
#joints3d_aligned, pelvis = align_by_pelvis(joints3d[0], get_pelvis = True)
# Convert to mm!
err, err_pa = compute_error_one_sample(
np.reshape(joints3d_gt_with_torso_subtracted * meter_to_mm,
[14, 3]), # gt joints3d;
joints3d_aligned *
meter_to_mm) # regressed joints3d via SMPL regressor;
#print ("[!!!] Processing {}, has mpjpe error = {} mm, mpjpe_pa error = {} mm".format(img_name, err, err_pa))
#*********************
""" 3) save results """
#*********************
if isSaveResults:
if not exists(result_dir):
os.makedirs(result_dir)
print("makedirs %s" % result_dir)
print('[***] input image :', img_name)
save_fig_result = join(
result_dir, 'lsp_14joints3d_%s' % img_name.replace('/', '-'))
draw_lsp_skeleton(joints3d[0], save_fig_result, with_numbers=True)
short_nm = img_name.split('/')[-1]
save_fig_result = join(result_dir, short_nm + '_infer.png')
cam_for_render, joints_orig = visualize(img_orig / 255.0, proc_param,
joints[0], verts[0], cams[0],
save_fig_result)
tmp_name = join(result_dir, short_nm + '_imgOrig.pfm')
pfm.save(tmp_name, img_orig)
tmp_name = join(result_dir, short_nm + '_imgUsed.pfm')
pfm.save(tmp_name, input_img[0])
tmp_name = join(result_dir, short_nm + '_depUsed.pfm')
pfm.save(tmp_name, input_dep[0])
# Write to an .obj file
outmesh_path = join(result_dir,
short_nm + '_mesh_mpjpepa_%.2fmm.obj' % err_pa)
save_to_mesh_file(
smpl_model_path, gender_vect,
np.concatenate(
[np.array([np.pi, .0, .0]),
np.reshape(theta[0, 3:72], [-1])],
axis=0), beta[0], outmesh_path)
return err, err_pa
def simple_infer_1_img(
image_path,
depth_path,
gender='n',
json_path=None,
result_dir='',
datatype='cad120',
has_joints3dgt=False,
smpl_model_path=None,
take_home_message=None, # some note about this experiment;
isSaveResults=False):
sess = tf.Session()
""" new model with depth"""
model = RunModelV2(config, sess=sess)
""" original model w/o depth"""
#model = RunModelV1(config, sess=sess)
# read image and depth; and normalized and rescaled to [-1.0, 1.0];
if '.pfm' in image_path:
img = pfm.load(image_path).astype(np.float32)
else:
img = cv2.cvtColor(cv2.imread(image_path, cv2.IMREAD_COLOR),
cv2.COLOR_BGR2RGB).astype(np.float32)
if depth_path is not None:
dep = pfm.load(depth_path).astype(np.float32)
else:
dep = np.zeros([img.size[0], img.size[1]], dtype=np.float32)
input_img, input_dep, proc_param, img_orig, _, _ = preprocess_image(
img, dep, json_path, None, None)
# Add batch dimension: 1 x D x D x 3
input_img = np.expand_dims(input_img, 0)
input_dep = np.expand_dims(input_dep, 0)
#*********************
""" 1) predicted """
#*********************
if gender == 'f':
gender_vect = np.array([1.0, .0, .0
]).astype(np.float32) # gender vector - female
elif gender == 'm':
gender_vect = np.array([.0, 1.0,
.0]).astype(np.float32) # gender vector - male
elif gender == 'n':
gender_vect = np.array([.0, .0, 1.0]).astype(
np.float32) # gender vector - neutral
gt_gender = np.expand_dims(gender_vect, 0)
joints, verts, cams, joints3d, _, _, _, _, theta, beta, joints3d_24 = model.predict(
input_img, input_dep, None, None, None, gt_gender, get_theta=True)
print("[***] prediction done !")
#*********************
""" 2) evaluation """
#*********************
gender_cad_dict = {
'cad-120/Subject1_rgbd_images': 'f',
'cad-120/Subject3_rgbd_images': 'f',
'cad-120/Subject4_rgbd_images': 'm',
'cad-120/Subject5_rgbd_images': 'm',
'cad-60/Person1': 'm',
'cad-60/Person2': 'f',
'cad-60/Person3': 'f',
'cad-60/Person4': 'm',
}
if has_joints3dgt:
print("[***] evaluating MPJPE and MPJPE PA !")
if datatype == 'cad120':
# e.g., image_path = 'cad-120-small/Subject4_rgbd_images/taking_medicine-1130145737-RGB_114_img_extracted.pfm'
cur_sub = image_path.split('/')[-2]
gender_cad = gender_cad_dict['cad-120/' + cur_sub]
tmp_name = image_path.split('/')[-1]
cad_anno_txt_fname = 'cad-120/' + cur_sub + '/' + tmp_name.split(
'-')[0] + '/' + tmp_name.split('-')[1] + '.txt'
sample_idx = int(tmp_name.split('-')[2].split('_')[1]) - 1
elif datatype == 'cad60':
# e.g., image_path = 'cad-60-small/Person4-0512152943-RGB_1041_dep_extracted.pfm'
tmp_name = image_path.split('/')[-1]
cur_sub = tmp_name.split('-')[0]
gender_cad = gender_cad_dict['cad-60/' + cur_sub]
cad_anno_txt_fname = 'cad-60/' + cur_sub + '/' + tmp_name.split(
'-')[1] + '.txt'
sample_idx = int(tmp_name.split('-')[2].split('_')[1]) - 1
print('gender_cad = %s, sample_idx = %d' % (gender_cad, sample_idx))
joints3d_gt_with_torso_subtracted, _, img_name, joints3d_with_root = get_one_batched_cad_toy_example(
cad_anno_txt_fname='./datasets/cad-60-120/' + cad_anno_txt_fname,
sample_idx=sample_idx,
gender=gender_cad,
swap=True)
#print ("??? joints3d shape = {}".format(joints3d.shape))
joints3d_aligned = align_by_pelvis(joints3d[0], get_pelvis=False)
#print ("??? joints3d_aligned shape = {}, joints3d_gt_with_torso_subtracted shape = {}".format(
# joints3d_aligned.shape,
# joints3d_gt_with_torso_subtracted[0].shape
#))
# Convert to mm!
err, err_pa = compute_error_one_sample(
np.reshape(joints3d_gt_with_torso_subtracted[0] * 1000.0,
[14, 3]), # gt joints3d;
joints3d_aligned *
1000.0) # regressed joints3d via SMPL regressor;
print(
"[!!!] Processing {}, has mpjpe error = {} mm, mpjpe_pa error = {} mm"
.format(img_name, err, err_pa))
#*********************
""" 3) save results """
#*********************
""" save as csv file, Excel file format """
timeStamp = datetime.now().strftime('%Y-%m-%d_%H:%M:%S')
csv_file = './results/experiments_eval_log.csv'
#messg = timeStamp + ',{:>7s},kt15,bad-3.0-noc,{:>6.3f},bad-3.0-all,{:>6.3f},fileDir={}\n'.format(
# eval_type[0], post_score[4]*100.0, post_score[11]*100.0, args.result_dir)
messg = timeStamp + ',image,{},mpjpe,{:.3f},mpjpe_pa,{:.3f},fileDir={},note,{}\n'.format(
img_name, err, err_pa, result_dir, take_home_message)
with open(csv_file, 'a') as fwrite:
fwrite.write(messg)
if isSaveResults:
if not exists(result_dir):
os.makedirs(result_dir)
print("makedirs %s" % result_dir)
print('[***] input image :', img_name)
save_fig_result = join(
result_dir, 'lsp_14joints3d_%s' % img_name.replace('/', '-'))
draw_lsp_skeleton(joints3d_with_root,
save_fig_result,
with_numbers=True)
short_nm = image_path.split('/')[-1][:-4]
save_fig_result = join(result_dir, short_nm + '_infer.png')
cam_for_render, joints_orig = visualize(img_orig / 255.0, proc_param,
joints[0], verts[0], cams[0],
save_fig_result)
tmp_name = join(result_dir, short_nm + '_imgOrig.pfm')
pfm.save(tmp_name, img_orig)
tmp_name = join(result_dir, short_nm + '_imgUsed.pfm')
pfm.save(tmp_name, input_img[0])
tmp_name = join(result_dir, short_nm + '_depUsed.pfm')
pfm.save(tmp_name, input_dep[0])
# Write to an .obj file
outmesh_path = join(result_dir,
short_nm + '_mesh_mpjpepa_%.2fmm.obj' % err_pa)
save_to_mesh_file(
smpl_model_path, gender_vect,
np.concatenate(
[np.array([np.pi, .0, .0]),
np.reshape(theta[0, 3:72], [-1])],
axis=0), beta[0], outmesh_path)