Skip to content

GitLab

Unverified Commit 50b6bbb3 authored by yinchimaoliang's avatar yinchimaoliang Committed by GitHub
Browse files

[Enhance] Remove unnecessary commands for creating sunrgbd data. (#110)

parent 9b091154
Show whitespace changes
Inline Side-by-side
Showing with 4 additions and 565 deletions
data/sunrgbd/README.md
View file @ 50b6bbb3
...@@ -5,9 +5,7 @@ We follow the procedure in [votenet](https://github.com/facebookresearch/votenet ...@@ -5,9 +5,7 @@ We follow the procedure in [votenet](https://github.com/facebookresearch/votenet
2. Enter the `matlab` folder, Extract point clouds and annotations by running `extract_split.m`, `extract_rgbd_data_v2.m` and `extract_rgbd_data_v1.m`. 2. Enter the `matlab` folder, Extract point clouds and annotations by running `extract_split.m`, `extract_rgbd_data_v2.m` and `extract_rgbd_data_v1.m`.
3. Back to this directory, prepare data by running `python sunrgbd_data.py --gen_v1_data`. 3. Enter the project root directory, Generate training data by running
4. Enter the project root directory, Generate training data by running
```bash ```bash
python tools/create_data.py sunrgbd --root-path ./data/sunrgbd --out-dir ./data/sunrgbd --extra-tag sunrgbd python tools/create_data.py sunrgbd --root-path ./data/sunrgbd --out-dir ./data/sunrgbd --extra-tag sunrgbd
``` ```
...@@ -18,9 +16,7 @@ cd matlab ...@@ -18,9 +16,7 @@ cd matlab
matlab -nosplash -nodesktop -r 'extract_split;quit;' matlab -nosplash -nodesktop -r 'extract_split;quit;'
matlab -nosplash -nodesktop -r 'extract_rgbd_data_v2;quit;' matlab -nosplash -nodesktop -r 'extract_rgbd_data_v2;quit;'
matlab -nosplash -nodesktop -r 'extract_rgbd_data_v1;quit;' matlab -nosplash -nodesktop -r 'extract_rgbd_data_v1;quit;'
cd .. cd ../../..
python sunrgbd_data.py --gen_v1_data
cd ../..
python tools/create_data.py sunrgbd --root-path ./data/sunrgbd --out-dir ./data/sunrgbd --extra-tag sunrgbd python tools/create_data.py sunrgbd --root-path ./data/sunrgbd --out-dir ./data/sunrgbd --extra-tag sunrgbd
``` ```
...@@ -29,8 +25,6 @@ NOTE: SUNRGBDtoolbox.zip should have MD5 hash `18d22e1761d36352f37232cba102f91f` ...@@ -29,8 +25,6 @@ NOTE: SUNRGBDtoolbox.zip should have MD5 hash `18d22e1761d36352f37232cba102f91f`
The directory structure after pre-processing should be as below The directory structure after pre-processing should be as below
``` ```
sunrgbd sunrgbd
├── sunrgbd_utils.py
├── sunrgbd_data.py
├── README.md ├── README.md
├── matlab ├── matlab
│ ├── extract_rgbd_data_v1.m │ ├── extract_rgbd_data_v1.m
...@@ -50,8 +44,6 @@ sunrgbd ...@@ -50,8 +44,6 @@ sunrgbd
│ ├── label │ ├── label
│ ├── seg_label │ ├── seg_label
│ ├── val_data_idx.txt │ ├── val_data_idx.txt
├── sunrgbd_pc_bbox_votes_50k_v1_train
├── sunrgbd_pc_bbox_votes_50k_v1_val
├── points ├── points
├── sunrgbd_infos_train.pkl ├── sunrgbd_infos_train.pkl
├── sunrgbd_infos_val.pkl ├── sunrgbd_infos_val.pkl
......
data/sunrgbd/sunrgbd_data.py deleted 100644 → 0
View file @ 9b091154
# Modified from
# https://github.com/facebookresearch/votenet/blob/master/sunrgbd/sunrgbd_data.py
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
"""Helper class and functions for loading SUN RGB-D objects.
Author: Charles R. Qi
Date: December, 2018
Note: removed unused code for frustum preparation.
Changed a way for data visualization (removed depdency on mayavi).
Load depth with scipy.io
"""
import argparse
import numpy as np
import os
import sunrgbd_utils
import sys
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)
sys.path.append(os.path.join(BASE_DIR, '../utils/'))
DEFAULT_TYPE_WHITELIST = [
'bed', 'table', 'sofa', 'chair', 'toilet', 'desk', 'dresser',
'night_stand', 'bookshelf', 'bathtub'
]
def random_sampling(pc, num_sample, replace=None, return_choices=False):
"""Input is NxC, output is num_samplexC."""
if replace is None:
replace = (pc.shape[0] < num_sample)
choices = np.random.choice(pc.shape[0], num_sample, replace=replace)
if return_choices:
return pc[choices], choices
else:
return pc[choices]
class sunrgbd_object(object):
"""Surrgbd object.
Load and parse object data.
Args:
root_dir(str): Root directory.
split(str): Training or testing.
Default: 'training.
use_v1(bool): Whether to use v1.
Default(False).
"""
def __init__(self, root_dir, split='training', use_v1=False):
self.root_dir = root_dir
self.split = split
assert (self.split == 'training')
self.split_dir = os.path.join(root_dir)
if split == 'training':
self.num_samples = 10335
elif split == 'testing':
self.num_samples = 2860
else:
print('Unknown split: %s' % (split))
exit(-1)
self.image_dir = os.path.join(self.split_dir, 'image')
self.calib_dir = os.path.join(self.split_dir, 'calib')
self.depth_dir = os.path.join(self.split_dir, 'depth')
if use_v1:
self.label_dir = os.path.join(self.split_dir, 'label_v1')
else:
self.label_dir = os.path.join(self.split_dir, 'label')
def __len__(self):
return self.num_samples
def get_image(self, idx):
img_filename = os.path.join(self.image_dir, '%06d.jpg' % (idx))
return sunrgbd_utils.load_image(img_filename)
def get_depth(self, idx):
depth_filename = os.path.join(self.depth_dir, '%06d.mat' % (idx))
return sunrgbd_utils.load_depth_points_mat(depth_filename)
def get_calibration(self, idx):
calib_filename = os.path.join(self.calib_dir, '%06d.txt' % (idx))
return sunrgbd_utils.SUNRGBD_Calibration(calib_filename)
def get_label_objects(self, idx):
label_filename = os.path.join(self.label_dir, '%06d.txt' % (idx))
return sunrgbd_utils.read_sunrgbd_label(label_filename)
def extract_sunrgbd_data(idx_filename,
split,
output_folder,
num_point=20000,
type_whitelist=DEFAULT_TYPE_WHITELIST,
save_votes=False,
use_v1=False,
skip_empty_scene=True):
"""Extract scene point clouds and bounding boxes (centroids, box sizes,
heading angles, semantic classes). Dumped point clouds and boxes are in
upright depth coord.
Args:
idx_filename: a TXT file where each line is an int number (index)
split: training or testing
save_votes: whether to compute and save Ground truth votes.
use_v1: use the SUN RGB-D V1 data
skip_empty_scene: if True, skip scenes that contain no object
(no objet in whitelist)
Dumps:
<id>_pc.npz of (N,6) where N is for number of subsampled points
and 6 is for XYZ and RGB (in 0~1) in upright depth coord
<id>_bbox.npy of (K,8) where K is the number of objects, 8 is for
centroids (cx,cy,cz), dimension (l,w,h), heanding_angle and
semantic_class
<id>_votes.npz of (N,10) with 0/1 indicating whether the point
belongs to an object, then three sets of GT votes for up to
three objects. If the point is only in one object's OBB, then
the three GT votes are the same.
"""
dataset = sunrgbd_object('./sunrgbd_trainval', split, use_v1=use_v1)
data_idx_list = [int(line.rstrip()) for line in open(idx_filename)]
if not os.path.exists(output_folder):
os.mkdir(output_folder)
for data_idx in data_idx_list:
print('------------- ', data_idx)
objects = dataset.get_label_objects(data_idx)
if skip_empty_scene and \
(len(objects) == 0 or
len([obj for obj in objects if
obj.classname in type_whitelist]) == 0):
continue
object_list = []
for obj in objects:
if obj.classname not in type_whitelist:
continue
obb = np.zeros((8))
obb[0:3] = obj.centroid
obb[3:6] = np.array([obj.length, obj.width, obj.height])
obb[6] = obj.heading_angle
obb[7] = sunrgbd_utils.type2class[obj.classname]
object_list.append(obb)
if len(object_list) == 0:
obbs = np.zeros((0, 8))
else:
obbs = np.vstack(object_list) # (K,8)
pc_upright_depth = dataset.get_depth(data_idx)
pc_upright_depth_subsampled = random_sampling(pc_upright_depth,
num_point)
np.savez_compressed(
os.path.join(output_folder, '%06d_pc.npz' % (data_idx)),
pc=pc_upright_depth_subsampled)
np.save(
os.path.join(output_folder, '%06d_bbox.npy' % (data_idx)), obbs)
if save_votes:
N = pc_upright_depth_subsampled.shape[0]
point_votes = np.zeros((N, 10)) # 3 votes and 1 vote mask
point_vote_idx = np.zeros(
(N)).astype(np.int32) # in the range of [0,2]
indices = np.arange(N)
for obj in objects:
if obj.classname not in type_whitelist:
continue
try:
# Find all points in this object's OBB
box3d_pts_3d = sunrgbd_utils.my_compute_box_3d(
obj.centroid,
np.array([obj.length, obj.width, obj.height]),
obj.heading_angle)
pc_in_box3d, inds = sunrgbd_utils.extract_pc_in_box3d(
pc_upright_depth_subsampled, box3d_pts_3d)
point_votes[inds, 0] = 1
votes = np.expand_dims(obj.centroid, 0) - pc_in_box3d[:,
0:3]
sparse_inds = indices[inds]
for i in range(len(sparse_inds)):
j = sparse_inds[i]
point_votes[j,
int(point_vote_idx[j] * 3 +
1):int((point_vote_idx[j] + 1) * 3 +
1)] = votes[i, :]
# Populate votes with the fisrt vote
if point_vote_idx[j] == 0:
point_votes[j, 4:7] = votes[i, :]
point_votes[j, 7:10] = votes[i, :]
point_vote_idx[inds] = np.minimum(2,
point_vote_idx[inds] + 1)
except Exception:
print('ERROR ----', data_idx, obj.classname)
np.savez_compressed(
os.path.join(output_folder, '%06d_votes.npz' % (data_idx)),
point_votes=point_votes)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument(
'--compute_median_size',
action='store_true',
help='Compute median 3D bounding box sizes for each class.')
parser.add_argument(
'--gen_v1_data', action='store_true', help='Generate V1 dataset.')
parser.add_argument(
'--gen_v2_data', action='store_true', help='Generate V2 dataset.')
args = parser.parse_args()
if args.gen_v1_data:
extract_sunrgbd_data(
os.path.join(BASE_DIR, 'sunrgbd_trainval/train_data_idx.txt'),
split='training',
output_folder=os.path.join(BASE_DIR,
'sunrgbd_pc_bbox_votes_50k_v1_train'),
save_votes=True,
num_point=50000,
use_v1=True,
skip_empty_scene=False)
extract_sunrgbd_data(
os.path.join(BASE_DIR, 'sunrgbd_trainval/val_data_idx.txt'),
split='training',
output_folder=os.path.join(BASE_DIR,
'sunrgbd_pc_bbox_votes_50k_v1_val'),
save_votes=True,
num_point=50000,
use_v1=True,
skip_empty_scene=False)
if args.gen_v2_data:
extract_sunrgbd_data(
os.path.join(BASE_DIR, 'sunrgbd_trainval/train_data_idx.txt'),
split='training',
output_folder=os.path.join(BASE_DIR,
'sunrgbd_pc_bbox_votes_50k_v2_train'),
save_votes=True,
num_point=50000,
use_v1=False,
skip_empty_scene=False)
extract_sunrgbd_data(
os.path.join(BASE_DIR, 'sunrgbd_trainval/val_data_idx.txt'),
split='training',
output_folder=os.path.join(BASE_DIR,
'sunrgbd_pc_bbox_votes_50k_v2_val'),
save_votes=True,
num_point=50000,
use_v1=False,
skip_empty_scene=False)
data/sunrgbd/sunrgbd_utils.py deleted 100644 → 0
View file @ 9b091154
# Modified from
# https://github.com/facebookresearch/votenet/blob/master/sunrgbd/sunrgbd_utils.py
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
"""Provides Python helper function to read My SUNRGBD dataset.
Author: Charles R. Qi
Date: October, 2017
Updated by Charles R. Qi
Date: December, 2018
Note: removed basis loading.
"""
import cv2
import numpy as np
from scipy import io as sio
type2class = {
'bed': 0,
'table': 1,
'sofa': 2,
'chair': 3,
'toilet': 4,
'desk': 5,
'dresser': 6,
'night_stand': 7,
'bookshelf': 8,
'bathtub': 9
}
class2type = {type2class[t]: t for t in type2class}
def flip_axis_to_camera(pc):
"""Flip axis to camera.
Flip X-right,Y-forward,Z-up to X-right,Y-down,Z-forward.
Args:
pc (np.ndarray): points in depth axis.
Returns:
np.ndarray: points in camera axis.
"""
pc2 = np.copy(pc)
pc2[:, [0, 1, 2]] = pc2[:, [0, 2, 1]] # cam X,Y,Z = depth X,-Z,Y
pc2[:, 1] *= -1
return pc2
def flip_axis_to_depth(pc):
pc2 = np.copy(pc)
pc2[:, [0, 1, 2]] = pc2[:, [0, 2, 1]] # depth X,Y,Z = cam X,Z,-Y
pc2[:, 2] *= -1
return pc2
class SUNObject3d(object):
def __init__(self, line):
data = line.split(' ')
data[1:] = [float(x) for x in data[1:]]
self.classname = data[0]
self.xmin = data[1]
self.ymin = data[2]
self.xmax = data[1] + data[3]
self.ymax = data[2] + data[4]
self.box2d = np.array([self.xmin, self.ymin, self.xmax, self.ymax])
self.centroid = np.array([data[5], data[6], data[7]])
self.unused_dimension = np.array([data[8], data[9], data[10]])
self.width = data[8]
self.length = data[9]
self.height = data[10]
self.orientation = np.zeros((3, ))
self.orientation[0] = data[11]
self.orientation[1] = data[12]
self.heading_angle = -1 * np.arctan2(self.orientation[1],
self.orientation[0])
class SUNRGBD_Calibration(object):
"""Calibration matrices and utils.
We define five coordinate system in SUN RGBD dataset:
camera coodinate:
Z is forward, Y is downward, X is rightward.
depth coordinate:
Just change axis order and flip up-down axis from camera coord.
upright depth coordinate: tilted depth coordinate by Rtilt such that
Z is gravity direction, Z is up-axis, Y is forward,
X is right-ward.
upright camera coordinate:
Just change axis order and flip up-down axis from upright.
depth coordinate
image coordinate:
----> x-axis (u)
|
v
y-axis (v)
depth points are stored in upright depth coordinate.
labels for 3d box (basis, centroid, size) are in upright
depth coordinate.
2d boxes are in image coordinate
We generate frustum point cloud and 3d box
in upright camera coordinate.
Args:
calib_filepath(str): Path of the calib file.
"""
def __init__(self, calib_filepath):
lines = [line.rstrip() for line in open(calib_filepath)]
Rtilt = np.array([float(x) for x in lines[0].split(' ')])
self.Rtilt = np.reshape(Rtilt, (3, 3), order='F')
K = np.array([float(x) for x in lines[1].split(' ')])
self.K = np.reshape(K, (3, 3), order='F')
self.f_u = self.K[0, 0]
self.f_v = self.K[1, 1]
self.c_u = self.K[0, 2]
self.c_v = self.K[1, 2]
def project_upright_depth_to_camera(self, pc):
"""Convert pc coordinate from depth to image.
Args:
pc (np.ndarray): Point cloud in depth coordinate.
Returns:
pc (np.ndarray): Point cloud in camera coordinate.
"""
# Project upright depth to depth coordinate
pc2 = np.dot(np.transpose(self.Rtilt), np.transpose(pc[:,
0:3])) # (3,n)
return flip_axis_to_camera(np.transpose(pc2))
def project_upright_depth_to_image(self, pc):
"""Convert pc coordinate from depth to image.
Args:
pc (np.ndarray): Point cloud in depth coordinate.
Returns:
np.ndarray: [N, 2] uv.
np.ndarray: [n,] depth.
"""
pc2 = self.project_upright_depth_to_camera(pc)
uv = np.dot(pc2, np.transpose(self.K)) # (n,3)
uv[:, 0] /= uv[:, 2]
uv[:, 1] /= uv[:, 2]
return uv[:, 0:2], pc2[:, 2]
def project_image_to_camera(self, uv_depth):
n = uv_depth.shape[0]
x = ((uv_depth[:, 0] - self.c_u) * uv_depth[:, 2]) / self.f_u
y = ((uv_depth[:, 1] - self.c_v) * uv_depth[:, 2]) / self.f_v
pts_3d_camera = np.zeros((n, 3))
pts_3d_camera[:, 0] = x
pts_3d_camera[:, 1] = y
pts_3d_camera[:, 2] = uv_depth[:, 2]
return pts_3d_camera
def rotz(t):
"""Rotation about the z-axis.
Args:
t (float): Heading angle.
Returns:
np.ndarray: Transforation matrix
"""
c = np.cos(t)
s = np.sin(t)
return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])
def transform_from_rot_trans(R, t):
"""Transforation matrix from rotation matrix and translation vector.
Args:
R (np.ndarray): Rotation matrix.
t (np.ndarray): Translation vector.
Returns:
np.ndarray: Transforation matrix.
"""
R = R.reshape(3, 3)
t = t.reshape(3, 1)
return np.vstack((np.hstack([R, t]), [0, 0, 0, 1]))
def read_sunrgbd_label(label_filename):
lines = [line.rstrip() for line in open(label_filename)]
objects = [SUNObject3d(line) for line in lines]
return objects
def load_image(img_filename):
return cv2.imread(img_filename)
def load_depth_points(depth_filename):
depth = np.loadtxt(depth_filename)
return depth
def load_depth_points_mat(depth_filename):
depth = sio.loadmat(depth_filename)['instance']
return depth
def in_hull(p, hull):
from scipy.spatial import Delaunay
if not isinstance(hull, Delaunay):
hull = Delaunay(hull)
return hull.find_simplex(p) >= 0
def extract_pc_in_box3d(pc, box3d):
"""Extract point cloud in box3d.
Args:
pc (np.ndarray): [N, 3] Point cloud.
box3d (np.ndarray): [8,3] 3d box.
Returns:
np.ndarray: Selected point cloud.
np.ndarray: Indices of selected point cloud.
"""
box3d_roi_inds = in_hull(pc[:, 0:3], box3d)
return pc[box3d_roi_inds, :], box3d_roi_inds
def my_compute_box_3d(center, size, heading_angle):
R = rotz(-1 * heading_angle)
l, w, h = size
x_corners = [-l, l, l, -l, -l, l, l, -l]
y_corners = [w, w, -w, -w, w, w, -w, -w]
z_corners = [h, h, h, h, -h, -h, -h, -h]
corners_3d = np.dot(R, np.vstack([x_corners, y_corners, z_corners]))
corners_3d[0, :] += center[0]
corners_3d[1, :] += center[1]
corners_3d[2, :] += center[2]
return np.transpose(corners_3d)
def compute_box_3d(obj, calib):
"""Takes an object and a projection matrix (P) and projects the 3d bounding
box into the image plane.
Args:
obj (SUNObject3d): Instance of SUNObject3d.
calib (SUNRGBD_Calibration): Instance of SUNRGBD_Calibration.
Returns:
np.ndarray: [8,2] array in image coord.
corners_3d: [8,3] array in in upright depth coord.
"""
center = obj.centroid
# compute rotational matrix around yaw axis
R = rotz(-1 * obj.heading_angle)
# 3d bounding box dimensions
length = obj.length # along heading arrow
width = obj.width # perpendicular to heading arrow
height = obj.height
# rotate and translate 3d bounding box
x_corners = [
-length, length, length, -length, -length, length, length, -length
]
y_corners = [width, width, -width, -width, width, width, -width, -width]
z_corners = [
height, height, height, height, -height, -height, -height, -height
]
corners_3d = np.dot(R, np.vstack([x_corners, y_corners, z_corners]))
corners_3d[0, :] += center[0]
corners_3d[1, :] += center[1]
corners_3d[2, :] += center[2]
# project the 3d bounding box into the image plane
corners_2d, _ = calib.project_upright_depth_to_image(
np.transpose(corners_3d))
return corners_2d, np.transpose(corners_3d)
setup.cfg
View file @ 50b6bbb3
...@@ -8,6 +8,6 @@ line_length = 79 ...@@ -8,6 +8,6 @@ line_length = 79
multi_line_output = 0 multi_line_output = 0
known_standard_library = setuptools known_standard_library = setuptools
known_first_party = mmdet,mmdet3d known_first_party = mmdet,mmdet3d
known_third_party = cv2,load_scannet_data,lyft_dataset_sdk,m2r,matplotlib,mmcv,nuimages,numba,numpy,nuscenes,pandas,plyfile,pycocotools,pyquaternion,pytest,recommonmark,scannet_utils,scipy,seaborn,shapely,skimage,sunrgbd_utils,terminaltables,torch,torchvision,trimesh known_third_party = cv2,load_scannet_data,lyft_dataset_sdk,m2r,matplotlib,mmcv,nuimages,numba,numpy,nuscenes,pandas,plyfile,pycocotools,pyquaternion,pytest,recommonmark,scannet_utils,scipy,seaborn,shapely,skimage,terminaltables,torch,torchvision,trimesh
no_lines_before = STDLIB,LOCALFOLDER no_lines_before = STDLIB,LOCALFOLDER
default_section = THIRDPARTY default_section = THIRDPARTY
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment