Initial commit
This commit is contained in:
Kijai
@@ -0,0 +1,218 @@
|
||||
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
|
||||
Reference in New Issue
Block a user