| from pathlib import Path |
| from typing import Optional, List, Callable, Dict, Any, Union |
| import warnings |
|
|
| import PIL.Image as pil_image |
| from torch import Tensor |
| from torch.utils.data import Dataset |
| from torchvision import transforms |
|
|
| from taming.data.conditional_builder.objects_bbox import ObjectsBoundingBoxConditionalBuilder |
| from taming.data.conditional_builder.objects_center_points import ObjectsCenterPointsConditionalBuilder |
| from taming.data.conditional_builder.utils import load_object_from_string |
| from taming.data.helper_types import BoundingBox, CropMethodType, Image, Annotation, SplitType |
| from taming.data.image_transforms import CenterCropReturnCoordinates, RandomCrop1dReturnCoordinates, \ |
| Random2dCropReturnCoordinates, RandomHorizontalFlipReturn, convert_pil_to_tensor |
|
|
|
|
| class AnnotatedObjectsDataset(Dataset): |
| def __init__(self, data_path: Union[str, Path], split: SplitType, keys: List[str], target_image_size: int, |
| min_object_area: float, min_objects_per_image: int, max_objects_per_image: int, |
| crop_method: CropMethodType, random_flip: bool, no_tokens: int, use_group_parameter: bool, |
| encode_crop: bool, category_allow_list_target: str = "", category_mapping_target: str = "", |
| no_object_classes: Optional[int] = None): |
| self.data_path = data_path |
| self.split = split |
| self.keys = keys |
| self.target_image_size = target_image_size |
| self.min_object_area = min_object_area |
| self.min_objects_per_image = min_objects_per_image |
| self.max_objects_per_image = max_objects_per_image |
| self.crop_method = crop_method |
| self.random_flip = random_flip |
| self.no_tokens = no_tokens |
| self.use_group_parameter = use_group_parameter |
| self.encode_crop = encode_crop |
|
|
| self.annotations = None |
| self.image_descriptions = None |
| self.categories = None |
| self.category_ids = None |
| self.category_number = None |
| self.image_ids = None |
| self.transform_functions: List[Callable] = self.setup_transform(target_image_size, crop_method, random_flip) |
| self.paths = self.build_paths(self.data_path) |
| self._conditional_builders = None |
| self.category_allow_list = None |
| if category_allow_list_target: |
| allow_list = load_object_from_string(category_allow_list_target) |
| self.category_allow_list = {name for name, _ in allow_list} |
| self.category_mapping = {} |
| if category_mapping_target: |
| self.category_mapping = load_object_from_string(category_mapping_target) |
| self.no_object_classes = no_object_classes |
|
|
| def build_paths(self, top_level: Union[str, Path]) -> Dict[str, Path]: |
| top_level = Path(top_level) |
| sub_paths = {name: top_level.joinpath(sub_path) for name, sub_path in self.get_path_structure().items()} |
| for path in sub_paths.values(): |
| if not path.exists(): |
| raise FileNotFoundError(f'{type(self).__name__} data structure error: [{path}] does not exist.') |
| return sub_paths |
|
|
| @staticmethod |
| def load_image_from_disk(path: Path) -> Image: |
| return pil_image.open(path).convert('RGB') |
|
|
| @staticmethod |
| def setup_transform(target_image_size: int, crop_method: CropMethodType, random_flip: bool): |
| transform_functions = [] |
| if crop_method == 'none': |
| transform_functions.append(transforms.Resize((target_image_size, target_image_size))) |
| elif crop_method == 'center': |
| transform_functions.extend([ |
| transforms.Resize(target_image_size), |
| CenterCropReturnCoordinates(target_image_size) |
| ]) |
| elif crop_method == 'random-1d': |
| transform_functions.extend([ |
| transforms.Resize(target_image_size), |
| RandomCrop1dReturnCoordinates(target_image_size) |
| ]) |
| elif crop_method == 'random-2d': |
| transform_functions.extend([ |
| Random2dCropReturnCoordinates(target_image_size), |
| transforms.Resize(target_image_size) |
| ]) |
| elif crop_method is None: |
| return None |
| else: |
| raise ValueError(f'Received invalid crop method [{crop_method}].') |
| if random_flip: |
| transform_functions.append(RandomHorizontalFlipReturn()) |
| transform_functions.append(transforms.Lambda(lambda x: x / 127.5 - 1.)) |
| return transform_functions |
|
|
| def image_transform(self, x: Tensor) -> (Optional[BoundingBox], Optional[bool], Tensor): |
| crop_bbox = None |
| flipped = None |
| for t in self.transform_functions: |
| if isinstance(t, (RandomCrop1dReturnCoordinates, CenterCropReturnCoordinates, Random2dCropReturnCoordinates)): |
| crop_bbox, x = t(x) |
| elif isinstance(t, RandomHorizontalFlipReturn): |
| flipped, x = t(x) |
| else: |
| x = t(x) |
| return crop_bbox, flipped, x |
|
|
| @property |
| def no_classes(self) -> int: |
| return self.no_object_classes if self.no_object_classes else len(self.categories) |
|
|
| @property |
| def conditional_builders(self) -> ObjectsCenterPointsConditionalBuilder: |
| |
| if self._conditional_builders is None: |
| self._conditional_builders = { |
| 'objects_center_points': ObjectsCenterPointsConditionalBuilder( |
| self.no_classes, |
| self.max_objects_per_image, |
| self.no_tokens, |
| self.encode_crop, |
| self.use_group_parameter, |
| getattr(self, 'use_additional_parameters', False) |
| ), |
| 'objects_bbox': ObjectsBoundingBoxConditionalBuilder( |
| self.no_classes, |
| self.max_objects_per_image, |
| self.no_tokens, |
| self.encode_crop, |
| self.use_group_parameter, |
| getattr(self, 'use_additional_parameters', False) |
| ) |
| } |
| return self._conditional_builders |
|
|
| def filter_categories(self) -> None: |
| if self.category_allow_list: |
| self.categories = {id_: cat for id_, cat in self.categories.items() if cat.name in self.category_allow_list} |
| if self.category_mapping: |
| self.categories = {id_: cat for id_, cat in self.categories.items() if cat.id not in self.category_mapping} |
|
|
| def setup_category_id_and_number(self) -> None: |
| self.category_ids = list(self.categories.keys()) |
| self.category_ids.sort() |
| if '/m/01s55n' in self.category_ids: |
| self.category_ids.remove('/m/01s55n') |
| self.category_ids.append('/m/01s55n') |
| self.category_number = {category_id: i for i, category_id in enumerate(self.category_ids)} |
| if self.category_allow_list is not None and self.category_mapping is None \ |
| and len(self.category_ids) != len(self.category_allow_list): |
| warnings.warn('Unexpected number of categories: Mismatch with category_allow_list. ' |
| 'Make sure all names in category_allow_list exist.') |
|
|
| def clean_up_annotations_and_image_descriptions(self) -> None: |
| image_id_set = set(self.image_ids) |
| self.annotations = {k: v for k, v in self.annotations.items() if k in image_id_set} |
| self.image_descriptions = {k: v for k, v in self.image_descriptions.items() if k in image_id_set} |
|
|
| @staticmethod |
| def filter_object_number(all_annotations: Dict[str, List[Annotation]], min_object_area: float, |
| min_objects_per_image: int, max_objects_per_image: int) -> Dict[str, List[Annotation]]: |
| filtered = {} |
| for image_id, annotations in all_annotations.items(): |
| annotations_with_min_area = [a for a in annotations if a.area > min_object_area] |
| if min_objects_per_image <= len(annotations_with_min_area) <= max_objects_per_image: |
| filtered[image_id] = annotations_with_min_area |
| return filtered |
|
|
| def __len__(self): |
| return len(self.image_ids) |
|
|
| def __getitem__(self, n: int) -> Dict[str, Any]: |
| image_id = self.get_image_id(n) |
| sample = self.get_image_description(image_id) |
| sample['annotations'] = self.get_annotation(image_id) |
|
|
| if 'image' in self.keys: |
| sample['image_path'] = str(self.get_image_path(image_id)) |
| sample['image'] = self.load_image_from_disk(sample['image_path']) |
| sample['image'] = convert_pil_to_tensor(sample['image']) |
| sample['crop_bbox'], sample['flipped'], sample['image'] = self.image_transform(sample['image']) |
| sample['image'] = sample['image'].permute(1, 2, 0) |
|
|
| for conditional, builder in self.conditional_builders.items(): |
| if conditional in self.keys: |
| sample[conditional] = builder.build(sample['annotations'], sample['crop_bbox'], sample['flipped']) |
|
|
| if self.keys: |
| |
| sample = {key: sample[key] for key in self.keys} |
| return sample |
|
|
| def get_image_id(self, no: int) -> str: |
| return self.image_ids[no] |
|
|
| def get_annotation(self, image_id: str) -> str: |
| return self.annotations[image_id] |
|
|
| def get_textual_label_for_category_id(self, category_id: str) -> str: |
| return self.categories[category_id].name |
|
|
| def get_textual_label_for_category_no(self, category_no: int) -> str: |
| return self.categories[self.get_category_id(category_no)].name |
|
|
| def get_category_number(self, category_id: str) -> int: |
| return self.category_number[category_id] |
|
|
| def get_category_id(self, category_no: int) -> str: |
| return self.category_ids[category_no] |
|
|
| def get_image_description(self, image_id: str) -> Dict[str, Any]: |
| raise NotImplementedError() |
|
|
| def get_path_structure(self): |
| raise NotImplementedError |
|
|
| def get_image_path(self, image_id: str) -> Path: |
| raise NotImplementedError |
|
|
