def scale_translation(log_dir, scale):
log_files = os.listdir(log_dir)
# log_files = sorted(log_files, key=lambda x: os.path.splitext())
for log_file in log_files:
cam = Cam(os.path.join(log_dir, log_file), max_d=192)
intrinsic = cam.intrinsic_mat
extrinsic = cam.extrinsic_mat
depth_min = cam.depth_min
depth_interval = cam.depth_interval
print(intrinsic)
print(extrinsic)
print(depth_min, depth_interval)
# fx, fy = intrinsic[0, 0], intrinsic[1, 1]
# intrinsic[0, 0], intrinsic[1, 1] = fx / 3.0, fy / 3.0
extrinsic[:3, 3] = extrinsic[:3, 3] * scale
with open(os.path.join(log_dir, log_file), 'w') as outfile:
outfile.write('extrinsic\n')
for row in extrinsic:
row = [str(item) for item in row]
row = ' '.join(row)
outfile.write(row + '\n')
outfile.write('\n')
outfile.write('intrinsic\n')
for row in intrinsic:
row = [str(item) for item in row]
row = ' '.join(row)
outfile.write(row + '\n')
outfile.write('\n')
outfile.write(str(depth_min) + ' ' + str(depth_interval))
def __iter__(self):
if self.shuffle:
self.rng.shuffle(self.sample_list)
for data in self.sample_list:
imgs = []
cams = []
for view in range(self.view_num):
# read_image
# // [fixedTODO]: center image is left to augmentor or tf Graph
# // I have done it here
img = cv2.imread(data[2 * view])
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# load cam and do basic interval_scale
cam = Cam(data[2 * view + 1], max_d=self.max_d)
cam.depth_interval = cam.depth_interval * self.interval_scale
imgs.append(img)
cams.append(cam.get_mat_form())
# load depth image of ref view
depth_image = PFMReader(data[2 * self.view_num]).data
# depth_image = np.zeros((10, 10))
# mask invalid depth_image
ref_cam = cams[0]
depth_min, depth_interval = Cam.get_depth_meta(
ref_cam, 'depth_min', 'depth_interval')
depth_start = depth_min + depth_interval
depth_end = depth_min + (self.max_d - 2) * depth_interval
# depth_image's shape: (h, w, 1)
depth_image = mask_depth_image(depth_image, depth_start, depth_end)
# view_num, h, w, 3
imgs = np.array(imgs)
# (view_num, )
cams = np.array(cams)
if self.test and self.count % 10 == 0:
print('Forward pass: d_min = %f, d_max = %f.' %
(depth_min, depth_min +
(self.max_d - 1) * depth_interval))
assert cams.shape == (self.view_num, 2, 4, 4)
self.count += 1
yield [imgs, cams, depth_image]
def generate_3d_point_cloud(rgb_path, depth_path, cam_path):
# depth_id = os.path.splitext(depth_path)[0]
img = cv2.imread(rgb_path)
rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
cam_id = os.path.splitext(cam_path)[0]
# assert depth_id == cam_id, (depth_id, cam_id)
cam = Cam(cam_path)
intrinsic = cam.intrinsic_mat
print('intrinsic: ')
print(intrinsic)
pfm_reader = PFMReader(depth_path)
depth_map = pfm_reader.data
ma = np.ma.masked_equal(depth_map, 0.0, copy=False)
print('value range: ', ma.min(), ma.max())
point_list = PointCloudGenerator.gen_3d_point_with_rgb(
depth_map, rgb, intrinsic)
PointCloudGenerator.write_as_obj(point_list, '%s.obj' % cam_id)
def post_process(base_dir):
"""
generate *.obj, generate rainbow depth map and rainbow prob map
:param base_dir:
:return:
"""
all_files = os.listdir(base_dir)
img_files = list(filter(is_img_file, all_files))
out_dir = path.join(base_dir, 'post_process')
if path.exists(out_dir):
shutil.rmtree(out_dir)
os.makedirs(out_dir)
for img_file in img_files:
file_id, _ = path.splitext(img_file)
""" depth_map and prob_map """
depth_path = path.join(base_dir, file_id + '_init.pfm')
prob_path = path.join(base_dir, file_id + '_prob.pfm')
# cmap = plt.cm.rainbow
depth_reader, prob_reader = PFMReader(depth_path), PFMReader(prob_path)
depth_map, prob_map = depth_reader.data, prob_reader.data
plt.imsave(path.join(out_dir, file_id + '_depth.png'),
depth_map * 255,
cmap='rainbow')
plt.imsave(path.join(out_dir, file_id + '_prob.png'),
prob_map * 255,
cmap='rainbow')
""" .obj file """
img = cv2.imread(path.join(base_dir, img_file))
rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
cam = Cam(path.join(base_dir, file_id + '.txt'))
intrinsic = cam.intrinsic_mat
ma = np.ma.masked_equal(depth_map, 0.0, copy=False)
print('value range: ', ma.min(), ma.max())
depth_point_list = PointCloudGenerator.gen_3d_point_with_rgb(
depth_map, rgb, intrinsic)
PointCloudGenerator.write_as_obj(
depth_point_list, path.join(out_dir, '%s_depth.obj' % file_id))
prob_point_list = PointCloudGenerator.gen_3d_point_with_rgb(
prob_map, rgb, intrinsic)
PointCloudGenerator.write_as_obj(
prob_point_list, path.join(out_dir, '%s_prob.obj' % file_id))
def test(model, sess_init, args):
"""
outputs prob_map, depth_map, rgb, and meshlab .obj file
:param model:
:param sess_init:
:param args:
:return:
"""
data_dir = args.data
out_base = args.out
view_num = args.view_num
max_h = args.max_h
max_w = args.max_w
max_d = args.max_d
interval_scale = args.interval_scale
logger.info('data_dir: %s, out_dir: %s' % (data_dir, out_base))
pred_conf = PredictConfig(
model=model,
session_init=sess_init,
input_names=['imgs', 'cams'],
output_names=['prob_map', 'coarse_depth', 'refine_depth', 'cost_volume_regularization/regularized_cost_volume']
)
# create imgs and cams data
# data_points = list(DTU.make_test_data(data_dir, view_num, max_h, max_w, max_d, interval_scale))
data_points = DTU.make_test_dataset(data_dir, view_num, max_h, max_w, max_d, interval_scale)
# model.batch_size = len(data_points)
pred_func = OfflinePredictor(pred_conf)
# TODO: after release training, finish this
# imgs = [dp[0] for dp in data_points]
# cams = [dp[1] for dp in data_points]
# batch_prob_map, batch_coarse_depth, batch_refine_depth = pred_func(imgs, cams)
#
# for i in range(len(batch_prob_map)):
# imgs, cams = data_points[i]
# prob_map, coarse_depth, refine_depth = batch_prob_map[i], batch_coarse_depth[i], batch_refine_depth[i]
# ref_img, ref_cam = imgs[0], cams[0]
# rgb = cv2.cvtColor(ref_img, cv2.COLOR_BGR2RGB)
# plt.imsave(path.join(out_dir, str(i) + '_prob.png'), np.squeeze(prob_map), cmap='rainbow')
# plt.imsave(path.join(out_dir, str(i) + '_depth.png'), np.squeeze(coarse_depth), cmap='rainbow')
# plt.imsave(path.join(out_dir, str(i) + '_rgb.png'), np.squeeze(rgb.astype('uint8')))
# Cam.write_cam(ref_cam, path.join(out_dir, str(i) + '_cam.txt'))
#
# intrinsic, *_ = Cam.get_depth_meta(ref_cam, 'intrinsic')
# print(intrinsic)
# ma = np.ma.masked_equal(coarse_depth, 0.0, copy=False)
# logger.info('value range: %f -> %f' % (ma.min(), ma.max()))
# downsample_rgb = cv2.resize(rgb, None, fx=0.25, fy=0.25, interpolation=cv2.INTER_LINEAR)
# depth_point_list = PointCloudGenerator.gen_3d_point_with_rgb(coarse_depth, downsample_rgb, intrinsic)
# PointCloudGenerator.write_as_obj(depth_point_list, path.join(out_dir, '%s_depth.obj' % str(i)))
# logger.info('len of data_points: %d' % len(data_points))
# Function here assumes batch = 1
dir_count = 0
view_num_count = 0
for idx, dp in enumerate(data_points):
out_dir = path.join(out_base, str(dir_count))
if not path.exists(out_dir):
os.makedirs(out_dir)
imgs, cams = dp
batch_prob_map, batch_coarse_depth, batch_refine_depth, batch_reg_cost_volume = \
pred_func(np.expand_dims(imgs, 0), np.expand_dims(cams, 0))
logger.info('shape of batch_prob_map: {}'.format(batch_prob_map.shape))
# size of reg_cost_volume: d, h/4, w/4
prob_map, coarse_depth, refine_depth, reg_cost_volume = np.squeeze(batch_prob_map), \
np.squeeze(batch_coarse_depth), \
np.squeeze(batch_refine_depth), \
np.squeeze(batch_reg_cost_volume)
quality_depth = np.where(prob_map > args.threshold, refine_depth, np.zeros_like(refine_depth))
mask_mat = np.where(prob_map > args.threshold, np.ones_like(refine_depth), np.zeros_like(refine_depth))
ref_img, ref_cam = imgs[0], cams[0]
rgb = cv2.cvtColor(ref_img, cv2.COLOR_BGR2RGB)
logger.info('shape of ref_img: {}'.format(ref_img.shape))
logger.info('shape of imgs: {}'.format(imgs.shape))
downsample_rgb = cv2.resize(rgb, None, fx=0.25, fy=0.25, interpolation=cv2.INTER_LINEAR)
quality_depth_upsample = cv2.resize(quality_depth, None, fx=4, fy=4, interpolation=cv2.INTER_NEAREST)
depth_upsample = cv2.resize(refine_depth, None, fx=4, fy=4, interpolation=cv2.INTER_NEAREST)
cv2.imwrite(path.join(out_dir, str(view_num_count) + '_depth_quality.exr'), quality_depth_upsample)
cv2.imwrite(path.join(out_dir, str(view_num_count) + '_depth.exr'), depth_upsample)
plt.imsave(path.join(out_dir, str(view_num_count) + '_prob.png'), prob_map, cmap='rainbow')
plt.imsave(path.join(out_dir, str(view_num_count) + '_depth.png'), coarse_depth, cmap='rainbow')
plt.imsave(path.join(out_dir, str(view_num_count) + '_depth_quality.png'), quality_depth, cmap='rainbow')
plt.imsave(path.join(out_dir, str(view_num_count) + '_rgb.png'), rgb.astype('uint8'))
# save the reg_cost_volume for future process
np.save(path.join(out_dir, str(view_num_count) + '_reg_cost_volume'), reg_cost_volume)
Cam.write_cam(ref_cam, path.join(out_dir, str(view_num_count) + '_cam.txt'), intrinsic_scale=4.)
rainbow_depth_quality = cv2.imread(path.join(out_dir, str(view_num_count) + '_depth_quality.png'))
rainbow_depth_quality = cv2.cvtColor(rainbow_depth_quality, cv2.COLOR_BGR2RGB)
rainbow_depth_quality = np.tile(np.expand_dims(mask_mat, -1), [1, 1, 3]) * rainbow_depth_quality
rainbow_depth_quality_up_sample = cv2.resize(rainbow_depth_quality, None, fx=4, fy=4,
interpolation=cv2.INTER_NEAREST)
alpha = 0.5
fused_rgb = (1 - alpha) * rgb + alpha * rainbow_depth_quality_up_sample
plt.imsave(path.join(out_dir, str(view_num_count) + '_fused_rgb.png'), fused_rgb.astype('uint8'))
intrinsic, *_ = Cam.get_depth_meta(ref_cam, 'intrinsic')
ma = np.ma.masked_equal(coarse_depth, 0.0, copy=False)
logger.info('value range: %f -> %f' % (ma.min(), ma.max()))
depth_point_list = PointCloudGenerator.gen_3d_point_with_rgb(coarse_depth, downsample_rgb, intrinsic, prob_map,
args.threshold)
PointCloudGenerator.write_as_obj(depth_point_list, path.join(out_dir, '%s_depth.obj' % str(view_num_count)))
view_num_count += 1
# FIXME: only works if the dir contains 5 imgs
if view_num_count == 5:
view_num_count = 0
dir_count += 1
def make_test_data(data_dir, view_num, max_h, max_w, max_d,
interval_scale):
"""
the data_dir should be organized like:
* images
* cams
* pair.txt
:param data_dir:
:return:
"""
dir_files = os.listdir(data_dir)
assert 'images' in dir_files and 'cams' in dir_files and 'pair.txt' in dir_files
sample_list = gen_test_input_sample_list(data_dir, view_num)
logger.info('sample_list: %s' % sample_list)
for data in sample_list:
imgs = []
cams = []
for view in range(view_num):
# read_image
# // [fixedTODO]: center image is left to augmentor or tf Graph
# // I have done it here
img = cv2.imread(data[2 * view])
# print(img.shape)
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# load cam and do basic interval_scale
cam = Cam(data[2 * view + 1],
max_d=max_d,
interval_scale=interval_scale)
imgs.append(img)
cams.append(cam.get_mat_form())
logger.info('range: {} {} {} {}'.format(cams[0][1, 3, 0],
cams[0][1, 3, 1],
cams[0][1, 3,
2], cams[0][1, 3,
3]))
general_h_scale = -1.
general_w_scale = -1.
# 选取较大的scale的好处是，宁愿 crop 也不要 padding
for view in range(view_num):
# print(imgs[view].shape)
h, w, _ = imgs[view].shape
height_scale = float(max_h) / h
width_scale = float(max_w) / w
general_h_scale = height_scale if height_scale > general_h_scale else general_h_scale
general_w_scale = width_scale if width_scale > general_w_scale else general_w_scale
assert height_scale < 1 and width_scale < 1, 'max_h, max_w shall be less than h, w'
resize_scale = general_h_scale if general_h_scale > general_w_scale else general_w_scale
logger.info('resize scale is %.2f' % resize_scale)
# first scale
imgs, cams = scale_mvs_input(imgs, cams, scale=resize_scale)
# then crop to fit the nn input
imgs, cams = crop_mvs_input(imgs,
cams,
max_h,
max_w,
base_image_size=8)
# then scale the cam and img, because the final resolution is not full-res
cams = [scale_camera(cam, 0.25) for cam in cams]
# imgs, cams = scale_mvs_input(imgs, cams, scale=0.25)
ref_cam = cams[0]
depth_min, depth_interval, depth_max = Cam.get_depth_meta(
ref_cam, 'depth_min', 'depth_interval', 'depth_max')
# view_num, h, w, 3
imgs = np.array(imgs)
# (view_num, )
cams = np.array(cams)
logger.info('d_min = %f, interval: %f, d_max = %f.' %
(depth_min, depth_interval, depth_max))
assert cams.shape == (view_num, 2, 4, 4)
yield imgs, cams