def __init__()

in wypr/dataset/scannet/scannet.py [0:0]

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

    def __init__(self, split_set='train', num_points=20000, use_color=False, use_height=False, 
                augment=False, use_normal=False, precomputed_prop="", sampling_method='random',
                determinstic_eval=False):

        self.CONFIG = ScannetDatasetConfig()
        self.data_path = os.path.join(BASE_DIR, 'scannet_all_points')
        all_scan_names = list(set([os.path.basename(x)[0:12]
                                   for x in os.listdir(self.data_path)
                                   if x.startswith('scene')]))
        if split_set in ['train', 'val']:
            split_filenames = os.path.join(BASE_DIR, 'meta_data', 'scannetv2_{}.txt'.format(split_set))
            with open(split_filenames, 'r') as f:  
                scan_names = f.read().splitlines()   
            # remove unavailiable scans
            self.scan_names = [sname for sname in scan_names if sname in all_scan_names]
            print('kept {} scans out of {}'.format(len(self.scan_names), len(scan_names)))
        elif split_set == 'debug':
            split_filenames = os.path.join(BASE_DIR, 'meta_data/scannetv2_val.txt')
            with open(split_filenames, 'r') as f: scan_names = f.read().splitlines()[:8]
            self.scan_names = [sname for sname in scan_names if sname in all_scan_names]
            print('kept {} scans out of {}'.format(len(self.scan_names), len(scan_names)))
        else:
            raise ValueError('illegal split name')
        self.split_set = split_set
        self.determinstic_eval = determinstic_eval
        self.num_points_sampled = num_points
        self.use_color = use_color        
        self.use_height = use_height
        self.use_normal = use_normal
        self.augment = augment
        self.sampling_method = sampling_method

        # proposals
        self.precomputed_prop = precomputed_prop 
        self.prop_path = os.path.join(BASE_DIR, 'proposals')