mirror of
https://github.com/opencv/opencv.git
synced 2024-12-15 09:49:13 +08:00
141 lines
4.4 KiB
Python
141 lines
4.4 KiB
Python
|
import numpy as np
|
||
|
import cv2 as cv
|
||
|
import argparse
|
||
|
|
||
|
# Use source data from this site:
|
||
|
# https://vision.in.tum.de/data/datasets/rgbd-dataset/download
|
||
|
# For example if you use rgbd_dataset_freiburg1_xyz sequence, your prompt should be:
|
||
|
# python /path_to_opencv/samples/python/volume.py --source_folder /path_to_datasets/rgbd_dataset_freiburg1_xyz --algo <some algo>
|
||
|
# so that the folder contains files groundtruth.txt and depth.txt
|
||
|
|
||
|
# for more info about this function look cv::Quat::toRotMat3x3(...)
|
||
|
def quatToMat3(a, b, c, d):
|
||
|
return np.array([
|
||
|
[1 - 2 * (c * c + d * d), 2 * (b * c - a * d) , 2 * (b * d + a * c)],
|
||
|
[2 * (b * c + a * d) , 1 - 2 * (b * b + d * d), 2 * (c * d - a * b)],
|
||
|
[2 * (b * d - a * c) , 2 * (c * d + a * b) , 1 - 2 * (b * b + c * c)]
|
||
|
])
|
||
|
|
||
|
def make_Rt(val):
|
||
|
R = quatToMat3(val[6], val[3], val[4] ,val[5])
|
||
|
t = np.array([ [val[0]], [val[1]], [val[2]] ])
|
||
|
tmp = np.array([0, 0, 0, 1])
|
||
|
|
||
|
Rt = np.append(R, t , axis=1 )
|
||
|
Rt = np.vstack([Rt, tmp])
|
||
|
|
||
|
return Rt
|
||
|
|
||
|
def get_image_info(path, is_depth):
|
||
|
image_info = {}
|
||
|
source = 'depth.txt'
|
||
|
if not is_depth:
|
||
|
source = 'rgb.txt'
|
||
|
with open(path+source) as file:
|
||
|
lines = file.readlines()
|
||
|
for line in lines:
|
||
|
words = line.split(' ')
|
||
|
if words[0] == '#':
|
||
|
continue
|
||
|
image_info[float(words[0])] = words[1][:-1]
|
||
|
return image_info
|
||
|
|
||
|
def get_groundtruth_info(path):
|
||
|
groundtruth_info = {}
|
||
|
with open(path+'groundtruth.txt') as file:
|
||
|
lines = file.readlines()
|
||
|
for line in lines:
|
||
|
words = line.split(' ')
|
||
|
if words[0] == '#':
|
||
|
continue
|
||
|
groundtruth_info[float(words[0])] = [float(i) for i in words[1:]]
|
||
|
return groundtruth_info
|
||
|
|
||
|
def main():
|
||
|
|
||
|
parser = argparse.ArgumentParser()
|
||
|
parser.add_argument(
|
||
|
'--algo',
|
||
|
help="""TSDF - reconstruct data in volume with bounds,
|
||
|
HashTSDF - reconstruct data in volume without bounds (infinite volume),
|
||
|
ColorTSDF - like TSDF but also keeps color data,
|
||
|
default - runs TSDF""",
|
||
|
default="")
|
||
|
parser.add_argument(
|
||
|
'-src',
|
||
|
'--source_folder',
|
||
|
default="")
|
||
|
|
||
|
args = parser.parse_args()
|
||
|
|
||
|
path = args.source_folder
|
||
|
if path[-1] != '/':
|
||
|
path += '/'
|
||
|
|
||
|
depth_info = get_image_info(path, True)
|
||
|
rgb_info = get_image_info(path, False)
|
||
|
groundtruth_info = get_groundtruth_info(path)
|
||
|
|
||
|
volume_type = cv.VolumeType_TSDF
|
||
|
if args.algo == "HashTSDF":
|
||
|
volume_type = cv.VolumeType_HashTSDF
|
||
|
elif args.algo == "ColorTSDF":
|
||
|
volume_type = cv.VolumeType_ColorTSDF
|
||
|
|
||
|
settings = cv.VolumeSettings(volume_type)
|
||
|
volume = cv.Volume(volume_type, settings)
|
||
|
|
||
|
for key in list(depth_info.keys())[:]:
|
||
|
Rt = np.eye(4)
|
||
|
for key1 in groundtruth_info:
|
||
|
if np.abs(key1 - key) < 0.01:
|
||
|
Rt = make_Rt(groundtruth_info[key1])
|
||
|
break
|
||
|
|
||
|
rgb_path = ''
|
||
|
for key1 in rgb_info:
|
||
|
if np.abs(key1 - key) < 0.05:
|
||
|
rgb_path = path + rgb_info[key1]
|
||
|
break
|
||
|
|
||
|
depthPath = path + depth_info[key]
|
||
|
depth = cv.imread(depthPath, cv.IMREAD_ANYDEPTH).astype(np.float32)
|
||
|
if depth.size <= 0:
|
||
|
raise Exception('Failed to load depth file: %s' % depthPath)
|
||
|
|
||
|
rgb = cv.imread(rgb_path, cv.IMREAD_COLOR).astype(np.float32)
|
||
|
if rgb.size <= 0:
|
||
|
raise Exception('Failed to load RGB file: %s' % rgb_path)
|
||
|
|
||
|
if volume_type != cv.VolumeType_ColorTSDF:
|
||
|
volume.integrate(depth, Rt)
|
||
|
else:
|
||
|
volume.integrateColor(depth, rgb, Rt)
|
||
|
|
||
|
size = (480, 640, 4)
|
||
|
|
||
|
points = np.zeros(size, np.float32)
|
||
|
normals = np.zeros(size, np.float32)
|
||
|
colors = np.zeros(size, np.float32)
|
||
|
|
||
|
if volume_type != cv.VolumeType_ColorTSDF:
|
||
|
volume.raycast(Rt, points, normals)
|
||
|
else:
|
||
|
volume.raycastColor(Rt, points, normals, colors)
|
||
|
|
||
|
channels = list(cv.split(points))
|
||
|
|
||
|
cv.imshow("X", np.absolute(channels[0]))
|
||
|
cv.imshow("Y", np.absolute(channels[1]))
|
||
|
cv.imshow("Z", channels[2])
|
||
|
|
||
|
if volume_type == cv.VolumeType_ColorTSDF:
|
||
|
cv.imshow("Color", colors.astype(np.uint8))
|
||
|
|
||
|
#TODO: also display normals
|
||
|
|
||
|
cv.waitKey(10)
|
||
|
|
||
|
if __name__ == '__main__':
|
||
|
main()
|