Spaces:
Runtime error
Runtime error
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 | |
def load_image_from_disk(path: Path) -> Image: | |
return pil_image.open(path).convert('RGB') | |
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 | |
def no_classes(self) -> int: | |
return self.no_object_classes if self.no_object_classes else len(self.categories) | |
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} | |
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 | |