def __getitem__()

in lib/datasets/frustum_dataset.py [0:0]

• 47 lines of code
• 9 McCabe index (conditional complexity)

    def __getitem__(self, index):
        ''' Get index-th element from the picked file dataset. '''
        # ------------------------------ INPUTS ----------------------------
        rot_angle = self.get_center_view_rot_angle(index)

        # Compute one hot vector
        cls_type = self.type_list[index]
        assert(cls_type in ['Car', 'Pedestrian', 'Cyclist'])
        one_hot_vec = np.zeros((3), dtype=np.float32)
        one_hot_vec[self.dataset_helper.type2onehot[cls_type]] = 1

        # Get point cloud
        if self.rotate_to_center:
            point_set = self.get_center_view_point_set(index)
        else:
            point_set = self.input_list[index]

        # Resample
        choice = np.random.choice(point_set.shape[0], self.npoints, replace=True)
        point_set = point_set[choice, :]

        patch_feat = self.rgb_feat_list[index]



        if self.from_rgb_detection:
            point_set = point_set.transpose(1, 0)  # N * C -> C * N
            return point_set, rot_angle, self.prob_list[index], self.id_list[index], \
                   self.type_list[index], self.box2d_list[index], one_hot_vec, patch_feat


        # ------------------------------ LABELS ----------------------------
        # segmentation labels
        seg = self.label_list[index]
        seg = seg[choice]

        # Get center point of 3D box
        if self.rotate_to_center:
            box3d_center = self.get_center_view_box3d_center(index)
        else:
            box3d_center = self.get_box3d_center(index)

        # Heading
        if self.rotate_to_center:
            heading_angle = self.heading_list[index] - rot_angle
        else:
            heading_angle = self.heading_list[index]

        # Size
        size_class, size_residual = self.dataset_helper.size2class(self.size_list[index], self.type_list[index])

        # Data Augmentation
        if self.random_flip:
            # note: rot_angle won't be correct if we have random_flip
            # so do not use it in case of random flipping.
            if np.random.random() > 0.5: # 50% chance flipping
                point_set[:, 0] *= -1
                box3d_center[0] *= -1
                heading_angle = np.pi - heading_angle
                patch_feat_tmp = np.fliplr(patch_feat.transpose(1,2,0))
                patch_feat_new = np.ascontiguousarray(patch_feat_tmp.transpose(2,0,1))
            else:
                patch_feat_new = patch_feat

        if self.random_shift:  # random shift object center
            dist = np.sqrt(np.sum(box3d_center[0]**2+box3d_center[2]**2))
            shift = np.clip(np.random.randn()*dist*0.2, -dist*0.2, dist*0.2)
            point_set[:,2] += shift
            box3d_center[2] += shift

        angle_class, angle_residual = self.dataset_helper.angle2class(heading_angle)

        point_set = point_set.transpose(1, 0)  # N * C -> C * N
        return point_set, seg, box3d_center, angle_class, angle_residual,\
            size_class, size_residual, rot_angle, one_hot_vec, patch_feat_new