"""
import os
import sys
import numpy as np
from torch.utils.data import Dataset
import scipy.io as sio # to load .mat files for depth points
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = os.path.dirname(BASE_DIR)
sys.path.append(BASE_DIR)
sys.path.append(os.path.join(ROOT_DIR, 'utils'))
import pc_util
import sunrgbd_utils
from model_util_sunrgbd import SunrgbdDatasetConfig
DC = SunrgbdDatasetConfig() # dataset specific config
MAX_NUM_OBJ = 64 # maximum number of objects allowed per scene
MEAN_COLOR_RGB = np.array([0.5, 0.5, 0.5]) # sunrgbd color is in 0~1
class SunrgbdDetectionVotesDataset(Dataset):
def __init__(self,
split_set='train',
num_points=20000,
use_color=False,
use_height=False,
use_v1=False,
augment=False,
scan_idx_list=None):
assert (num_points <= 50000)
DUMP_FOUT.write(out_str + "\n")
DUMP_FOUT.flush()
print(out_str)
# Init datasets and dataloaders
def my_worker_init_fn(worker_id):
np.random.seed(np.random.get_state()[1][0] + worker_id)
if FLAGS.dataset == "sunrgbd":
sys.path.append(os.path.join(ROOT_DIR, "sunrgbd"))
from sunrgbd_detection_dataset import SunrgbdDetectionVotesDataset, MAX_NUM_OBJ
from model_util_sunrgbd import SunrgbdDatasetConfig
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(
from model_util_sunrgbd import SunrgbdDatasetConfig
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 = 20000
DATASET_CONFIG = SunrgbdDatasetConfig()
TRAIN_DATASET = SunrgbdDetectionVotesDataset('train',
num_points=NUM_POINT,
augment=False,
use_color=False,
use_height=True,
use_v1=True)
TRAIN_DATALOADER = DataLoader(
TRAIN_DATASET,
batch_size=1,
shuffle=False,
num_workers=6,
)
import tensorflow as tf BASE_PATH = os.path.dirname((os.path.abspath(__file__))) ROOT_PATH = os.path.dirname(BASE_PATH) sys.path.append(os.path.join(ROOT_PATH, 'pointnet2/utils')) sys.path.append(os.path.join(ROOT_PATH, 'sunrgbd')) sys.path.append(os.path.join(ROOT_PATH, 'utils')) from pointnet_util import pointnet_sa_module, pointnet_fp_module from model_util_sunrgbd import SunrgbdDatasetConfig from nn_distance import nn_distance from nn_distance import huber_loss import tf_util import numpy as np # constant for dataset DC = SunrgbdDatasetConfig() NUM_CLASS = 10 NUM_HEADING_BIN = 12 NUM_SIZE_CLUSTER = 10 NUM_PROPOSAL = 256 MEAN_SIZE_ARR = DC.mean_size_arr MAX_NUM_OBJ = 64 # constant for loss calculation FAR_THRESHOLD = 0.6 NEAR_THRESHOLD = 0.3 GT_VOTE_FACTOR = 3 # number of GT votes per point #OBJECTNESS_CLS_WEIGHTS = [0.2,0.8] # put larger weights on positive objectness OBJECTNESS_POS_WEIGHTS = 4 def placeholder_inputs(batch_size, num_point):
def forward(self, xyz, features, center_z, z_feature, center_xy,
xy_feature, center_line, line_feature, end_points):
"""
Args:
xyz: (B,K,3)
features: (B,C,K)
Returns:
scores: (B,num_proposal,2+3+NH*2+NS*4)
"""
if self.sampling == 'vote_fps':
original_features = features
xyz, features, fps_inds = self.vote_aggregation(xyz, features)
sample_inds = fps_inds
elif self.sampling == 'seed_fps':
# FPS on seed and choose the votes corresponding to the seeds
# This gets us a slightly better coverage of *object* votes than vote_fps (which tends to get more cluster votes)
sample_inds = pointnet2_utils.furthest_point_sample(
end_points['seed_xyz'], self.num_proposal)
xyz, features, _ = self.vote_aggregation(xyz, features,
sample_inds)
elif self.sampling == 'random':
# Random sampling from the votes
num_seed = end_points['seed_xyz'].shape[1]
sample_inds = torch.randint(0,
num_seed,
(batch_size, self.num_proposal),
dtype=torch.int).cuda()
xyz, features, _ = self.vote_aggregation(xyz, features,
sample_inds)
else:
log_string('Unknown sampling strategy: %s. Exiting!' %
(self.sampling))
exit()
end_points[
'aggregated_vote_xyzcenter'] = xyz # (batch_size, num_proposal, 3)
end_points[
'aggregated_vote_inds'] = sample_inds # (batch_size, num_proposal,) # should be 0,1,2,...,num_proposal
batch_size = xyz.shape[0]
object_proposal = xyz.shape[1]
# --------- PROPOSAL GENERATION ---------
net = F.relu(self.bn1(self.conv1(features)))
net = F.relu(self.bn2(self.conv2(net)))
original_feature = features.contiguous()
net = self.conv3(
net
) # (batch_size, 2+3+num_heading_bin*2+num_size_cluster*4, num_proposal)
final_feature = net
center_vote, size_vote, sizescore_vote, end_points = decode_scores(
net,
end_points,
self.num_class,
self.num_heading_bin,
self.num_size_cluster,
self.mean_size_arr,
mode='center')
### Create surface center here
### Extract surface points and features here
ind_normal_z = self.softmax_normal(end_points["pred_flag_z"])
end_points["pred_z_ind"] = (ind_normal_z[:, 1, :] >
SURFACE_THRESH).detach().float()
z_sel = (ind_normal_z[:, 1, :] <= SURFACE_THRESH).detach().float()
offset = torch.ones_like(center_z) * UPPER_THRESH
z_center = center_z + offset * z_sel.unsqueeze(-1)
z_sem = end_points["sem_cls_scores_z"]
ind_normal_xy = self.softmax_normal(end_points["pred_flag_xy"])
end_points["pred_xy_ind"] = (ind_normal_xy[:, 1, :] >
SURFACE_THRESH).detach().float()
xy_sel = (ind_normal_xy[:, 1, :] <= SURFACE_THRESH).detach().float()
offset = torch.ones_like(center_xy) * UPPER_THRESH
xy_center = center_xy + offset * xy_sel.unsqueeze(-1)
xy_sem = end_points["sem_cls_scores_xy"]
surface_center_pred = torch.cat((z_center, xy_center), dim=1)
end_points['surface_center_pred'] = surface_center_pred
end_points['surface_sem_pred'] = torch.cat((z_sem, xy_sem), dim=1)
surface_center_feature_pred = torch.cat((z_feature, xy_feature), dim=2)
surface_center_feature_pred = torch.cat((torch.zeros(
(batch_size, 6, surface_center_feature_pred.shape[2])).cuda(),
surface_center_feature_pred),
dim=1)
### Extract line points and features here
ind_normal_line = self.softmax_normal(end_points["pred_flag_line"])
end_points["pred_line_ind"] = (ind_normal_line[:, 1, :] >
LINE_THRESH).detach().float()
line_sel = (ind_normal_line[:, 1, :] <=
SURFACE_THRESH).detach().float()
offset = torch.ones_like(center_line) * UPPER_THRESH
line_center = center_line + offset * line_sel.unsqueeze(-1)
end_points['line_center_pred'] = line_center
end_points['line_sem_pred'] = end_points["sem_cls_scores_line"]
end_points['aggregated_vote_xyzopt'] = xyz
### Extract the object center here
obj_center = center_vote.contiguous()
size_residual = size_vote.contiguous()
pred_size_class = torch.argmax(sizescore_vote.contiguous(), -1)
pred_size_residual = torch.gather(
size_vote.contiguous(), 2,
pred_size_class.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, 1, 3))
mean_size_class_batched = torch.ones_like(
size_residual) * torch.from_numpy(
self.mean_size_arr.astype(
np.float32)).cuda().unsqueeze(0).unsqueeze(0)
pred_size_avg = torch.gather(
mean_size_class_batched, 2,
pred_size_class.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, 1, 3))
obj_size = (pred_size_avg.squeeze(2) +
pred_size_residual.squeeze(2)).detach()
pred_heading_class = torch.argmax(end_points['heading_scores' +
'center'].detach(),
-1) # B,num_proposal
pred_heading_residual = torch.gather(
end_points['heading_residuals' + 'center'].detach(), 2,
pred_heading_class.unsqueeze(-1)) # B,num_proposal,1
pred_heading_residual.squeeze_(2)
# if self.num_class == 18:
if not self.with_angle:
pred_heading = torch.zeros_like(pred_heading_class)
else:
config = SunrgbdDatasetConfig()
pred_heading = pred_heading_class.float() * (2 * np.pi / float(
config.num_heading_bin)) + pred_heading_residual
obj_surface_center, obj_line_center = get_surface_line_points_batch_pytorch(
obj_size, pred_heading, obj_center)
obj_surface_feature = original_feature.repeat(1, 1, 6)
end_points['surface_center_object'] = obj_surface_center
# Add an indicator for different surfaces
obj_upper_indicator = torch.zeros(
(batch_size, object_proposal, 6)).cuda()
obj_upper_indicator[:, :, 0] = 1
obj_lower_indicator = torch.zeros(
(batch_size, object_proposal, 6)).cuda()
obj_lower_indicator[:, :, 1] = 1
obj_front_indicator = torch.zeros(
(batch_size, object_proposal, 6)).cuda()
obj_front_indicator[:, :, 2] = 1
obj_back_indicator = torch.zeros(
(batch_size, object_proposal, 6)).cuda()
obj_back_indicator[:, :, 3] = 1
obj_left_indicator = torch.zeros(
(batch_size, object_proposal, 6)).cuda()
obj_left_indicator[:, :, 4] = 1
obj_right_indicator = torch.zeros(
(batch_size, object_proposal, 6)).cuda()
obj_right_indicator[:, :, 5] = 1
obj_surface_indicator = torch.cat(
(obj_upper_indicator, obj_lower_indicator, obj_front_indicator,
obj_back_indicator, obj_left_indicator, obj_right_indicator),
dim=1).transpose(2, 1).contiguous()
obj_surface_feature = torch.cat(
(obj_surface_indicator, obj_surface_feature), dim=1)
obj_line_feature = original_feature.repeat(1, 1, 12)
end_points['line_center_object'] = obj_line_center
# Add an indicator for different lines
obj_line_indicator0 = torch.zeros(
(batch_size, 12, object_proposal)).cuda()
obj_line_indicator0[:, 0, :] = 1
obj_line_indicator1 = torch.zeros(
(batch_size, 12, object_proposal)).cuda()
obj_line_indicator1[:, 1, :] = 1
obj_line_indicator2 = torch.zeros(
(batch_size, 12, object_proposal)).cuda()
obj_line_indicator2[:, 2, :] = 1
obj_line_indicator3 = torch.zeros(
(batch_size, 12, object_proposal)).cuda()
obj_line_indicator3[:, 3, :] = 1
obj_line_indicator4 = torch.zeros(
(batch_size, 12, object_proposal)).cuda()
obj_line_indicator4[:, 4, :] = 1
obj_line_indicator5 = torch.zeros(
(batch_size, 12, object_proposal)).cuda()
obj_line_indicator5[:, 5, :] = 1
obj_line_indicator6 = torch.zeros(
(batch_size, 12, object_proposal)).cuda()
obj_line_indicator6[:, 6, :] = 1
obj_line_indicator7 = torch.zeros(
(batch_size, 12, object_proposal)).cuda()
obj_line_indicator7[:, 7, :] = 1
obj_line_indicator8 = torch.zeros(
(batch_size, 12, object_proposal)).cuda()
obj_line_indicator8[:, 8, :] = 1
obj_line_indicator9 = torch.zeros(
(batch_size, 12, object_proposal)).cuda()
obj_line_indicator9[:, 9, :] = 1
obj_line_indicator10 = torch.zeros(
(batch_size, 12, object_proposal)).cuda()
obj_line_indicator10[:, 10, :] = 1
obj_line_indicator11 = torch.zeros(
(batch_size, 12, object_proposal)).cuda()
obj_line_indicator11[:, 11, :] = 1
obj_line_indicator = torch.cat(
(obj_line_indicator0, obj_line_indicator1, obj_line_indicator2,
obj_line_indicator3, obj_line_indicator4, obj_line_indicator5,
obj_line_indicator6, obj_line_indicator7, obj_line_indicator8,
obj_line_indicator9, obj_line_indicator10, obj_line_indicator11),
dim=2)
obj_line_feature = torch.cat((obj_line_indicator, obj_line_feature),
dim=1)
surface_xyz, surface_features, _ = self.match_surface_center(
torch.cat((obj_surface_center, surface_center_pred), dim=1),
torch.cat((obj_surface_feature, surface_center_feature_pred),
dim=2))
line_feature = torch.cat((torch.zeros(
(batch_size, 12, line_feature.shape[2])).cuda(), line_feature),
dim=1)
line_xyz, line_features, _ = self.match_line_center(
torch.cat((obj_line_center, line_center), dim=1),
torch.cat((obj_line_feature, line_feature), dim=2))
combine_features = torch.cat(
(surface_features.contiguous(), line_features.contiguous()), dim=2)
match_features = F.relu(
self.bn_match1(self.conv_match1(combine_features)))
match_score = self.conv_match2(match_features)
end_points["match_scores"] = match_score.transpose(2, 1).contiguous()
match_features_sem = F.relu(
self.bn_match_sem1(self.conv_match_sem1(combine_features)))
match_score_sem = self.conv_match_sem2(match_features_sem)
end_points["match_scores_sem"] = match_score_sem.transpose(
2, 1).contiguous()
surface_features = F.relu(
self.bn_surface1(self.conv_surface1(surface_features)))
surface_features = F.relu(
self.bn_surface2(self.conv_surface2(surface_features)))
line_features = F.relu(self.bn_line1(self.conv_line1(line_features)))
line_features = F.relu(self.bn_line2(self.conv_line2(line_features)))
surface_features = surface_features.view(batch_size, -1, 6,
object_proposal).contiguous()
line_features = line_features.view(batch_size, -1, 12,
object_proposal).contiguous()
# Combine all surface and line features
surface_pool_feature = surface_features.view(
batch_size, -1, object_proposal).contiguous()
line_pool_feature = line_features.view(batch_size, -1,
object_proposal).contiguous()
combine_feature = torch.cat((surface_pool_feature, line_pool_feature),
dim=1)
net = F.relu(self.bn_refine1(self.conv_refine1(combine_feature)))
net += original_feature
net = F.relu(self.bn_refine2(self.conv_refine2(net)))
net = F.relu(self.bn_refine3(self.conv_refine3(net)))
net = self.conv_refine4(
net
) # (batch_size, 2+3+num_heading_bin*2+num_size_cluster*4, num_proposal)
end_points = decode_scores(net,
end_points,
self.num_class,
self.num_heading_bin,
self.num_size_cluster,
self.mean_size_arr,
mode='opt')
return end_points
