An axis aligned bounding box is parameterized by (cx,cy,cz) and (dx,dy,dz)
where (cx,cy,cz) is the center point of the box, dx is the x-axis length of the box.
"""
import os
import sys
import numpy as np
from torch.utils.data import Dataset
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = os.path.dirname(BASE_DIR)
sys.path.append(ROOT_DIR)
sys.path.append(os.path.join(ROOT_DIR, 'utils'))
import pc_util
from model_util_scannet import rotate_aligned_boxes
from model_util_scannet import ScannetDatasetConfig
DC = ScannetDatasetConfig()
MAX_NUM_OBJ = 64
MEAN_COLOR_RGB = np.array([109.8, 97.2, 83.8])
class ScannetDetectionDataset(Dataset):
def __init__(self,
split_set='train',
num_points=20000,
use_color=False,
use_height=False,
augment=False):
self.data_path = os.path.join(BASE_DIR, 'scannet_train_detection_data')
all_scan_names = list(set([os.path.basename(x)[0:12] \
for x in os.listdir(self.data_path) if x.startswith('scene')]))
from tqdm import tqdm
import trimesh
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader
NUM_POINT=40000
DATASET_CONFIG = ScannetDatasetConfig()
TRAIN_DATASET = ScannetDetectionDataset('train',
num_points=NUM_POINT,
augment=False,
use_color=False,
use_height=True,
overfit=False)
TRAIN_DATALOADER = DataLoader(TRAIN_DATASET,
batch_size=1,
shuffle=False,
num_workers=6,)
# -- from pc_utils.py write_oriented_bbox
def heading2rotmat(heading_angle):
rotmat = np.zeros((3,3))
DATASET_CONFIG = SunrgbdDatasetConfig()
TEST_DATASET = SunrgbdDetectionVotesDataset(
"val",
num_points=NUM_POINT,
augment=False,
use_color=FLAGS.use_color,
use_height=(not FLAGS.no_height),
use_v1=(not FLAGS.use_sunrgbd_v2),
)
elif FLAGS.dataset == "scannet":
sys.path.append(os.path.join(ROOT_DIR, "scannet"))
from scannet_detection_dataset import ScannetDetectionDataset, MAX_NUM_OBJ
from model_util_scannet import ScannetDatasetConfig
DATASET_CONFIG = ScannetDatasetConfig()
TEST_DATASET = ScannetDetectionDataset(
"val", num_points=NUM_POINT, augment=False, use_color=FLAGS.use_color, use_height=(not FLAGS.no_height)
)
elif FLAGS.dataset == "fsn":
sys.path.append(os.path.join(ROOT_DIR, "fsn"))
from fsn_detection_dataset import FSNDetectionDataset, MAX_NUM_OBJ
from model_util_fsn import FSNDatasetConfig
DATASET_CONFIG = FSNDatasetConfig()
TEST_DATASET = FSNDetectionDataset(
"val" if not FLAGS.overfit else "overfit",
num_points=NUM_POINT,
augment=False,
use_color=FLAGS.use_color,
use_height=(not FLAGS.no_height),
scene_name = 'scannet_train_detection_data/scene0002_00'
output_folder = 'data_viz_dump'
data = np.load(scene_name + '_vert.npy')
scene_points = data[:, 0:3]
colors = data[:, 3:]
instance_labels = np.load(scene_name + '_ins_label.npy')
semantic_labels = np.load(scene_name + '_sem_label.npy')
instance_bboxes = np.load(scene_name + '_bbox.npy')
print(np.unique(instance_labels))
print(np.unique(semantic_labels))
input()
if not os.path.exists(output_folder):
os.mkdir(output_folder)
# Write scene as OBJ file for visualization
pc_util.write_ply_rgb(scene_points, colors,
os.path.join(output_folder, 'scene.obj'))
pc_util.write_ply_color(scene_points, instance_labels,
os.path.join(output_folder, 'scene_instance.obj'))
pc_util.write_ply_color(scene_points, semantic_labels,
os.path.join(output_folder, 'scene_semantic.obj'))
from model_util_scannet import ScannetDatasetConfig
DC = ScannetDatasetConfig()
print(instance_bboxes.shape)
bbox, bbox_color = DC.param2bbox(instance_bboxes, True)
pc_util.write_ply_color(bbox, bbox_color,
os.path.join(output_folder, 'scene_bbox.obj'))