def __init__()

in contactopt/loader.py [0:0]

• 12 lines of code
• 5 McCabe index (conditional complexity)

    def __init__(self, data, train=False, min_num_cont=1):
        start_time = time.time()
        self.train = train
        self.aug_vert_jitter = 0.0005

        if isinstance(data, str):
            self.dataset = pickle.load(open(data, 'rb'))    # Load pickle, can take many seconds
        else:
            self.dataset = data

        if 'num_verts_in_contact' in self.dataset[0]:
            print('Cutting samples with less than {} points in contact. Was size {}'.format(min_num_cont, len(self.dataset)))
            self.dataset = [s for s in self.dataset if s['num_verts_in_contact'] >= min_num_cont]

        print('Dataset loaded in {:.2f} sec, {} samples'.format(time.time() - start_time, len(self.dataset)))