sd/stablediffusion/src/taming-transformers/taming/data/annotated_objects_dataset.py

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:
        # cannot set this up in init because no_classes is only known after loading data in init of superclass
        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:
            # only return specified 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