def angle_to_direction_label()

in models/UN_EPT.py [0:0]

• 22 lines of code
• 8 McCabe index (conditional complexity)

    def angle_to_direction_label(self, angle_map, seg_label_map=None, distance_map=None, num_classes=8, extra_ignore_mask=None, return_tensor=False):

        if return_tensor:
            assert isinstance(angle_map, torch.Tensor)
            assert isinstance(seg_label_map, torch.Tensor) or seg_label_map is None
        else:
            assert isinstance(angle_map, np.ndarray)
            assert isinstance(seg_label_map, np.ndarray) or seg_label_map is None

        _, label_map = self.align_angle(angle_map, num_classes=num_classes, return_tensor=return_tensor)
        if distance_map is not None:
            label_map[distance_map > 5] = num_classes
        if seg_label_map is None:
            if return_tensor:
                ignore_mask = torch.zeros(angle_map.shape, dtype=torch.uint8).to(angle_map.device)
            else:
                ignore_mask = np.zeros(angle_map.shape, dtype=np.bool)
        else:
            ignore_mask = seg_label_map == -1

        if extra_ignore_mask is not None:
            extra_ignore_mask = extra_ignore_mask.unsqueeze(1)
            ignore_mask = ignore_mask | extra_ignore_mask
        label_map[ignore_mask] = -1

        return label_map