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

示例#4
0
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