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')]))
示例#2
0
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))
示例#3
0
文件: eval.py 项目: SirWyver/votenet
    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),
示例#4
0
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'))