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),
is_eval=True,
quick=FLAGS.val_quick,
num_points=NUM_POINT,
augment=False,
use_color=FLAGS.use_color,
use_height=(not FLAGS.no_height),
use_v1=(not FLAGS.use_sunrgbd_v2),
labeled_sample_list=FLAGS.labeled_sample_list,
test_transductive=FLAGS.transductive)
elif FLAGS.dataset == 'scannet':
sys.path.append(os.path.join(ROOT_DIR, 'scannet'))
from scannet_detection_dataset import ScannetDetectionDataset
from model_util_scannet import ScannetDatasetConfig
DATASET_CONFIG = ScannetDatasetConfig()
TEST_DATASET = ScannetDetectionDataset(
split_set,
num_points=NUM_POINT,
augment=False,
use_color=FLAGS.use_color,
use_height=(not FLAGS.no_height),
labeled_sample_list=FLAGS.labeled_sample_list,
test_transductive=FLAGS.transductive)
else:
print('Unknown dataset %s. Exiting...' % (FLAGS.dataset))
exit(-1)
print(len(TEST_DATASET))
TEST_DATALOADER = DataLoader(TEST_DATASET,
batch_size=BATCH_SIZE,
shuffle=FLAGS.shuffle_dataset,
num_workers=4,
worker_init_fn=my_worker_init_fn)
# Init the model and optimzier
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
use_v1=(not FLAGS.use_sunrgbd_v2))
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()
TRAIN_DATASET = ScannetDetectionDataset('train',
num_points=NUM_POINT,
augment=True,
use_color=FLAGS.use_color,
use_height=(not FLAGS.no_height))
TEST_DATASET = ScannetDetectionDataset('val',
num_points=NUM_POINT,
augment=False,
use_color=FLAGS.use_color,
use_height=(not FLAGS.no_height))
else:
print('Unknown dataset %s. Exiting...' % (FLAGS.dataset))
exit(-1)
print(len(TRAIN_DATASET), len(TEST_DATASET))
TRAIN_DATALOADER = DataLoader(TRAIN_DATASET,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=4,
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))
rotmat[2,2] = 1
cosval = np.cos(heading_angle)
sinval = np.sin(heading_angle)
rotmat[0:2,0:2] = np.array([[cosval, -sinval],[sinval, cosval]])
