diff --git a/basketballtrainer/data/convert_dataset.py b/basketballtrainer/data/convert_dataset.py index 3691967..679a89c 100644 --- a/basketballtrainer/data/convert_dataset.py +++ b/basketballtrainer/data/convert_dataset.py @@ -18,7 +18,7 @@ def convert_dataset_to_paddleseg_format(dataset_path: str, target_path: str): """ - Copy images and labels annotated in MATLAB to PaddleSeg-compliant directory structure. + Copy images and labels annotated in MATLAB to a PaddleSeg-compliant directory structure. **Please note** that the dataset structure generated by this function still has to be split and shuffled by following the instructions found [here](https://github.com/PaddlePaddle/PaddleSeg/blob/release/2.8/docs/data/marker/marker.md#4-split-a-custom-dataset) :param dataset_path: path string pointing to the root of the MATLAB dataset @@ -26,12 +26,14 @@ def convert_dataset_to_paddleseg_format(dataset_path: str, target_path: str): """ source = pathlib.Path(dataset_path) target = pathlib.Path(target_path) - images, labels = __generate_ordered_filenames_lists(source) + if (source / 'frames').exists() and (source / 'masks').exists(): + images, labels = __generate_ordered_filenames_lists(source) + else: + images, labels = __generate_ordered_filenames_lists(source / '*/*') if not os.path.exists(str(target / 'images')): os.mkdir(str(target / 'images')) if not os.path.exists(str(target / 'labels')): os.mkdir(str(target / 'labels')) - for sample_index in range(len(images)): shutil.copy2(images[sample_index], str(target / f'images/image{sample_index + 1}.png')) shutil.copy2(labels[sample_index], str(target / f'labels/label{sample_index + 1}.png')) @@ -40,7 +42,7 @@ def convert_dataset_to_paddleseg_format(dataset_path: str, target_path: str): def __generate_ordered_filenames_lists(source: pathlib.Path) -> ([str], [str]): images = [] labels = [] - for match_directory_path in glob.iglob(str(source / '*/*')): + for match_directory_path in glob.iglob(str(source)): match_directory = pathlib.Path(match_directory_path) match_image_paths = [ match_image_path