Small fixes for Open Images dataset 🔧

Author: manolo-lolo

configs/coco_scene_images_transformer.yaml

@@ -48,13 +48,12 @@ data:
   target: main.DataModuleFromConfig
   params:
     batch_size: 6
-    num_workers: 12
     train:
       target: taming.data.annotated_objects_coco.AnnotatedObjectsCoco
       params:
         data_path: data/coco_annotations_100  # substitute with path to full dataset
         split: train
-        keys: [image, objects_bbox, file_name]
+        keys: [image, objects_bbox, file_name, annotations]
         no_tokens: 8192
         target_image_size: 256
         min_object_area: 0.00001
@@ -69,7 +68,7 @@ data:
       params:
         data_path: data/coco_annotations_100  # substitute with path to full dataset
         split: validation
-        keys: [image, objects_bbox, file_name]
+        keys: [image, objects_bbox, file_name, annotations]
         no_tokens: 8192
         target_image_size: 256
         min_object_area: 0.00001

configs/open_images_scene_images_transformer.yaml

@@ -8,9 +8,9 @@ model:
       params:
         vocab_size: 8192
         block_size: 348  # = 256 + 92 = dim(vqgan_latent_space,16x16) + dim(conditional_builder.embedding_dim)
-        n_layer: 40
+        n_layer: 36
         n_head: 16
-        n_embd: 1408
+        n_embd: 1536
         embd_pdrop: 0.1
         resid_pdrop: 0.1
         attn_pdrop: 0.1
@@ -48,15 +48,16 @@ data:
   target: main.DataModuleFromConfig
   params:
     batch_size: 6
-    num_workers: 12
     train:
       target: taming.data.annotated_objects_open_images.AnnotatedObjectsOpenImages
       params:
         data_path: data/open_images_annotations_100  # substitute with path to full dataset
         split: train
-        keys: [image, objects_bbox, file_name]
+        keys: [image, objects_bbox, file_name, annotations]
         no_tokens: 8192
         target_image_size: 256
+        category_allow_list_target: taming.data.open_images_helper.top_300_classes_plus_coco_compatibility
+        category_mapping_target: taming.data.open_images_helper.open_images_unify_categories_for_coco
         min_object_area: 0.0001
         min_objects_per_image: 2
         max_objects_per_image: 30
@@ -65,13 +66,15 @@ data:
         use_group_parameter: true
         encode_crop: true
     validation:
-      target: taming.data.annotated_objects_coco.AnnotatedObjectsCoco
+      target: taming.data.annotated_objects_open_images.AnnotatedObjectsOpenImages
       params:
         data_path: data/open_images_annotations_100  # substitute with path to full dataset
         split: validation
-        keys: [image, objects_bbox, file_name]
+        keys: [image, objects_bbox, file_name, annotations]
         no_tokens: 8192
         target_image_size: 256
+        category_allow_list_target: taming.data.open_images_helper.top_300_classes_plus_coco_compatibility
+        category_mapping_target: taming.data.open_images_helper.open_images_unify_categories_for_coco
         min_object_area: 0.0001
         min_objects_per_image: 2
         max_objects_per_image: 30

taming/data/annotated_objects_dataset.py

@@ -1,5 +1,6 @@
 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
@@ -8,6 +9,7 @@
 
 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
@@ -17,7 +19,7 @@ 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):
+                 encode_crop: bool, category_allow_list_target: str, category_mapping_target: str):
         self.data_path = data_path
         self.split = split
         self.keys = keys
@@ -40,6 +42,12 @@ def __init__(self, data_path: Union[str, Path], split: SplitType, keys: List[str
         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
+        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)
 
     def build_paths(self, top_level: Union[str, Path]) -> Dict[str, Path]:
         top_level = Path(top_level)
@@ -123,12 +131,22 @@ def conditional_builders(self) -> ObjectsCenterPointsConditionalBuilder:
         return self._conditional_builders
 
     def filter_categories(self) -> None:
-        pass
+        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)

taming/data/annotated_objects_open_images.py

@@ -33,8 +33,8 @@
 }
 
 
-def load_annotations(descriptor_path: Path, min_object_area: float, category_no_for_id: Dict[str, int]) -> \
-        Dict[str, List[Annotation]]:
+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)
@@ -43,6 +43,8 @@ def load_annotations(descriptor_path: Path, min_object_area: float, category_no_
             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(
@@ -114,7 +116,8 @@ def __init__(self, **kwargs):
         self.setup_category_id_and_number()
 
         self.image_descriptions = {}
-        annotations = load_annotations(self.paths['annotations'], self.min_object_area, self.category_number)
+        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())
@@ -129,4 +132,5 @@ 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]:
-        return {'file_path': str(self.get_image_path(image_id))}
+        image_path = self.get_image_path(image_id)
+        return {'file_path': str(image_path), 'file_name': image_path.name}

taming/data/conditional_builder/utils.py

@@ -1,3 +1,4 @@
+import importlib
 from typing import List, Any, Tuple, Optional
 
 from taming.data.helper_types import BoundingBox, Annotation
@@ -94,3 +95,11 @@ def get_circle_size(figure_size: Tuple[int, int]) -> int:
     if max(figure_size) >= 512:
         circle_size = 4
     return circle_size
+
+
+def load_object_from_string(object_string: str) -> Any:
+    """
+    Source: https://stackoverflow.com/a/10773699
+    """
+    module_name, class_name = object_string.rsplit(".", 1)
+    return getattr(importlib.import_module(module_name), class_name)