Function get_ShapeNet_model_list

pointersect/data/genlist.py:52–143

Given setting (train/test), generate a list of obj file names in the shapenet dataset. Args: rootpath: the root directory of the ShapeNet dataset setting: train/test num_classes: number of class num_sample_in_class: number of mesh sample in each class

(
        rootpath: str,
        setting: str,
        num_classes: int = 3,
        num_sample_in_class: int = 2,
        rnd_seed: int = 987,
        use_sub_classlist: bool = False,
)

def get_ShapeNet_model_list(
        rootpath: str,
        setting: str,
        num_classes: int = 3,
        num_sample_in_class: int = 2,
        rnd_seed: int = 987,
        use_sub_classlist: bool = False,
) -> list:
    """
    Given setting (train/test), generate a list of obj file names in the shapenet dataset.

    Args:
        rootpath: the root directory of the ShapeNet dataset
        setting: train/test
        num_classes: number of class
        num_sample_in_class: number of mesh sample in each class
        rnd_seed: random seed
        use_sub_classlist: use specific class set for train/ test

    Returns:
        a list of filenames
    """
    np.random.seed(rnd_seed)
    img_types = ('*.jpg', '*.png', '*.jpeg', '*.JPG')
    obj_path_list = []

    folder_list = sorted([fname for fname in os.listdir(rootpath) if os.path.isdir(os.path.join(rootpath, fname))])
    with open(os.path.join(rootpath, 'taxonomy.json')) as f:
        class_json_list = json.load(f)
        all_class_list = [(class_json['name'].split(',')[0].replace(' ', '_'), class_json['synsetId']) for class_json in
                          class_json_list if class_json['synsetId'] in folder_list]
        # all_class_list = [ c[0].split(',')[0].replace(' ','_') for c in all_class_list ]
        class_to_folder = dict(all_class_list)
        """
            class_to_folder = dict(
                airplane='02691156',
                bench='02828884',
                cabinet='02933112',
                car='02958343',
                chair='03001627',
                display='03211117',
                lamp='03636649',
                loudspeaker='03691459',
                rifle='04090263',
                sofa='04256520',
                table='04379243',
                cellular_telephone='04401088',
                vessel='04530566',
            )
        """

        # In other works, the first 3 classes below are split to training set and the other 10 classes are for testing
        fixed_class = ['airplane', 'car', 'chair', 'bench', 'rifle', 'vessel', 'cabinet', 'display', 'lamp',
                       'loudspeaker', 'sofa', 'table', 'cellular_telephone']
        if setting == 'train':
            class_list = fixed_class[:3]
        elif setting == 'test':
            class_list = fixed_class[3:]