Args: categories (iterable): List of categories to use. if not specified, will use TEST_CATEGORIES. num_images (int): Default number of images in each batch.
(
self,
CO3D_DIR,
CO3D_ANNOTATION_DIR,
categories: Union[str, list, None] = None,
split_name: str = "test",
min_num_images: int = 50,
sort_by_filename: bool = False,
)
| 109 | |
| 110 | class Co3dDataset(Dataset): |
| 111 | def __init__( |
| 112 | self, |
| 113 | CO3D_DIR, |
| 114 | CO3D_ANNOTATION_DIR, |
| 115 | categories: Union[str, list, None] = None, |
| 116 | split_name: str = "test", |
| 117 | min_num_images: int = 50, |
| 118 | sort_by_filename: bool = False, |
| 119 | ): |
| 120 | """ |
| 121 | Args: |
| 122 | categories (iterable): List of categories to use. if not specified, will use TEST_CATEGORIES. |
| 123 | num_images (int): Default number of images in each batch. |
| 124 | """ |
| 125 | categories = TEST_CATEGORIES if categories is None else categories |
| 126 | if isinstance(categories, str): |
| 127 | if categories == "test": |
| 128 | categories = TEST_CATEGORIES |
| 129 | elif categories == "debug": |
| 130 | categories = DEBUG_CATEGORIES |
| 131 | elif categories == "train": |
| 132 | categories = TRAINING_CATEGORIES |
| 133 | elif categories == "all": |
| 134 | categories = TRAINING_CATEGORIES + TEST_CATEGORIES |
| 135 | else: |
| 136 | raise ValueError(f"Unknown str category: {categories}") |
| 137 | elif isinstance(categories, list): |
| 138 | categories = categories |
| 139 | else: |
| 140 | raise ValueError(f"Unknown categories: {categories}") |
| 141 | self.split_name = split_name |
| 142 | self.categories = categories |
| 143 | |
| 144 | self.low_quality_translations = [] |
| 145 | self.metadata = {} |
| 146 | self.category_map = {} |
| 147 | |
| 148 | print(f"[CO3DV2] CO3D_DIR is {CO3D_DIR}") |
| 149 | self.CO3D_DIR = CO3D_DIR |
| 150 | self.CO3D_ANNOTATION_DIR = CO3D_ANNOTATION_DIR |
| 151 | self.min_num_images = min_num_images |
| 152 | |
| 153 | for c in categories: |
| 154 | annotation_file = osp.join(self.CO3D_ANNOTATION_DIR, f"{c}_{split_name}.jgz") |
| 155 | with gzip.open(annotation_file, "r") as fin: |
| 156 | annotation = json.loads(fin.read()) |
| 157 | |
| 158 | counter = 0 |
| 159 | for seq_name, seq_data in annotation.items(): |
| 160 | counter += 1 |
| 161 | if len(seq_data) < min_num_images: |
| 162 | print(f"[CO3DV2] sequence {seq_name} in category {c} has only {len(seq_data)} images, filter it.") |
| 163 | continue |
| 164 | |
| 165 | filtered_data = [] |
| 166 | self.category_map[seq_name] = c |
| 167 | bad_seq = False |
| 168 | for data in seq_data: |