| from collections import defaultdict | |
| from csv import DictReader, reader as TupleReader | |
| from pathlib import Path | |
| from typing import Dict, List, Any | |
| import warnings | |
| from taming.data.annotated_objects_dataset import AnnotatedObjectsDataset | |
| from taming.data.helper_types import Annotation, Category | |
| from tqdm import tqdm | |
| OPEN_IMAGES_STRUCTURE = { | |
| 'train': { | |
| 'top_level': '', | |
| 'class_descriptions': 'class-descriptions-boxable.csv', | |
| 'annotations': 'oidv6-train-annotations-bbox.csv', | |
| 'file_list': 'train-images-boxable.csv', | |
| 'files': 'train' | |
| }, | |
| 'validation': { | |
| 'top_level': '', | |
| 'class_descriptions': 'class-descriptions-boxable.csv', | |
| 'annotations': 'validation-annotations-bbox.csv', | |
| 'file_list': 'validation-images.csv', | |
| 'files': 'validation' | |
| }, | |
| 'test': { | |
| 'top_level': '', | |
| 'class_descriptions': 'class-descriptions-boxable.csv', | |
| 'annotations': 'test-annotations-bbox.csv', | |
| 'file_list': 'test-images.csv', | |
| 'files': 'test' | |
| } | |
| } | |
| def load_annotations(descriptor_path: Path, min_object_area: float, category_mapping: Dict[str, str], | |
| category_no_for_id: Dict[str, int]) -> Dict[str, List[Annotation]]: | |
| annotations: Dict[str, List[Annotation]] = defaultdict(list) | |
| with open(descriptor_path) as file: | |
| reader = DictReader(file) | |
| for i, row in tqdm(enumerate(reader), total=14620000, desc='Loading OpenImages annotations'): | |
| width = float(row['XMax']) - float(row['XMin']) | |
| height = float(row['YMax']) - float(row['YMin']) | |
| area = width * height | |
| category_id = row['LabelName'] | |
| if category_id in category_mapping: | |
| category_id = category_mapping[category_id] | |
| if area >= min_object_area and category_id in category_no_for_id: | |
| annotations[row['ImageID']].append( | |
| Annotation( | |
| id=i, | |
| image_id=row['ImageID'], | |
| source=row['Source'], | |
| category_id=category_id, | |
| category_no=category_no_for_id[category_id], | |
| confidence=float(row['Confidence']), | |
| bbox=(float(row['XMin']), float(row['YMin']), width, height), | |
| area=area, | |
| is_occluded=bool(int(row['IsOccluded'])), | |
| is_truncated=bool(int(row['IsTruncated'])), | |
| is_group_of=bool(int(row['IsGroupOf'])), | |
| is_depiction=bool(int(row['IsDepiction'])), | |
| is_inside=bool(int(row['IsInside'])) | |
| ) | |
| ) | |
| if 'train' in str(descriptor_path) and i < 14000000: | |
| warnings.warn(f'Running with subset of Open Images. Train dataset has length [{len(annotations)}].') | |
| return dict(annotations) | |
| def load_image_ids(csv_path: Path) -> List[str]: | |
| with open(csv_path) as file: | |
| reader = DictReader(file) | |
| return [row['image_name'] for row in reader] | |
| def load_categories(csv_path: Path) -> Dict[str, Category]: | |
| with open(csv_path) as file: | |
| reader = TupleReader(file) | |
| return {row[0]: Category(id=row[0], name=row[1], super_category=None) for row in reader} | |
| class AnnotatedObjectsOpenImages(AnnotatedObjectsDataset): | |
| def __init__(self, use_additional_parameters: bool, **kwargs): | |
| """ | |
| @param data_path: is the path to the following folder structure: | |
| open_images/ | |
| │ oidv6-train-annotations-bbox.csv | |
| ├── class-descriptions-boxable.csv | |
| ├── oidv6-train-annotations-bbox.csv | |
| ├── test | |
| │ ├── 000026e7ee790996.jpg | |
| │ ├── 000062a39995e348.jpg | |
| │ └── ... | |
| ├── test-annotations-bbox.csv | |
| ├── test-images.csv | |
| ├── train | |
| │ ├── 000002b66c9c498e.jpg | |
| │ ├── 000002b97e5471a0.jpg | |
| │ └── ... | |
| ├── train-images-boxable.csv | |
| ├── validation | |
| │ ├── 0001eeaf4aed83f9.jpg | |
| │ ├── 0004886b7d043cfd.jpg | |
| │ └── ... | |
| ├── validation-annotations-bbox.csv | |
| └── validation-images.csv | |
| @param: split: one of 'train', 'validation' or 'test' | |
| @param: desired image size (returns square images) | |
| """ | |
| super().__init__(**kwargs) | |
| self.use_additional_parameters = use_additional_parameters | |
| self.categories = load_categories(self.paths['class_descriptions']) | |
| self.filter_categories() | |
| self.setup_category_id_and_number() | |
| self.image_descriptions = {} | |
| annotations = load_annotations(self.paths['annotations'], self.min_object_area, self.category_mapping, | |
| self.category_number) | |
| self.annotations = self.filter_object_number(annotations, self.min_object_area, self.min_objects_per_image, | |
| self.max_objects_per_image) | |
| self.image_ids = list(self.annotations.keys()) | |
| self.clean_up_annotations_and_image_descriptions() | |
| def get_path_structure(self) -> Dict[str, str]: | |
| if self.split not in OPEN_IMAGES_STRUCTURE: | |
| raise ValueError(f'Split [{self.split} does not exist for Open Images data.]') | |
| return OPEN_IMAGES_STRUCTURE[self.split] | |
| def get_image_path(self, image_id: str) -> Path: | |
| return self.paths['files'].joinpath(f'{image_id:0>16}.jpg') | |
| def get_image_description(self, image_id: str) -> Dict[str, Any]: | |
| image_path = self.get_image_path(image_id) | |
| return {'file_path': str(image_path), 'file_name': image_path.name} | |