Adding Filtering dataloader

Author: Hartorn

loreal_poc/dataloaders/base.py

@@ -5,6 +5,9 @@
 
 import numpy as np
 
+SingleLandmarkData = Tuple[np.ndarray, np.ndarray, Optional[Dict[Any, Any]]]
+BatchedLandmarkData = Tuple[Tuple[np.ndarray], np.ndarray, Tuple[Optional[Dict[Any, Any]]]]
+
 
 class DataIteratorBase(ABC):
     batch_size: int
@@ -61,12 +64,10 @@ def get_meta_with_default(self, idx: int) -> np.ndarray:
         marks = marks if marks is not None else self.meta_none()
         return marks
 
-    def get_single_element(self, idx) -> Tuple[np.ndarray, Optional[np.ndarray], Optional[Dict[Any, Any]]]:
+    def get_single_element(self, idx) -> SingleLandmarkData:
         return self.get_image(idx), self.get_marks_with_default(idx), self.get_meta_with_default(idx)
 
-    def __getitem__(
-        self, idx: int
-    ) -> Tuple[np.ndarray, Optional[np.ndarray], Optional[Dict[Any, Any]]]:  # (image, marks, meta)
+    def __getitem__(self, idx: int) -> BatchedLandmarkData:
         return self._collate_fn(
             [self.get_single_element(i) for i in self.idx_sampler[idx * self.batch_size : (idx + 1) * self.batch_size]]
         )
@@ -84,30 +85,22 @@ def all_marks(self) -> np.ndarray:  # (marks)
     def all_meta(self) -> List:  # (meta)
         return [self.get_meta_with_default(idx) for idx in self.idx_sampler]
 
-    def __next__(self) -> Tuple[np.ndarray, np.ndarray]:
+    def __next__(self) -> BatchedLandmarkData:
         if self.idx >= len(self):
             raise StopIteration
         elt = self[self.idx]
         self.idx += 1
         return elt
 
-    def _collate_fn(
-        self, elements: List[Tuple[np.ndarray, Optional[np.ndarray], Optional[Dict[Any, Any]]]]
-    ) -> Tuple[np.ndarray, Optional[np.ndarray], Optional[Dict[Any, Any]]]:
-        batched_elements = list(zip(*elements))
-        batched_elements[1] = np.array(batched_elements[1])
-
+    def _collate_fn(self, elements: List[SingleLandmarkData]) -> BatchedLandmarkData:
+        batched_elements = list(zip(*elements, strict=True))
         # INFO: Restore if we want to concatenate all meta under one dict instead of keeping them as records (list of dicts)
         # meta_keys = next((list(elt.keys()) for elt in batched_elements[2] if elt is not None), [])
         # batched_elements[2] = {
         #    key: [meta[key] if (meta is not None and key in meta) else None for meta in batched_elements[2]]
         #    for key in meta_keys
         # }
-
-        # if len(batched_elements[0]) != self.batch_size:
-        #     raise StopIteration
-
-        return batched_elements
+        return batched_elements[0], np.array(batched_elements[1]), batched_elements[2]
 
 
 class DataLoaderBase(DataIteratorBase):

loreal_poc/dataloaders/wrappers.py

@@ -1,10 +1,10 @@
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 
 from ..marks.facial_parts import FacialPart
 from ..transformation_functions.crop import crop_image_from_mark, crop_mark
-from .base import DataIteratorBase, DataLoaderWrapper
+from .base import DataIteratorBase, DataLoaderWrapper, SingleLandmarkData
 
 
 class CroppedDataLoader(DataLoaderWrapper):
@@ -51,3 +51,22 @@ def __getitem__(self, idx: int) -> Tuple[np.ndarray, Optional[np.ndarray], Optio
         if len(self._cache_idxs) > self._max_size:
             self._cache.pop(self._cache_idxs.pop(-1))
         return self._cache[idx]
+
+
+class FilteringDataLoader(DataLoaderWrapper):
+    @property
+    def name(self):
+        return f"({self._wrapped_dataloader.name}) filtered using {self._predicate_name}"
+
+    @property
+    def idx_sampler(self) -> np.ndarray:
+        return self._reindex
+
+    def __init__(self, dataloader: DataIteratorBase, predicate: Callable[[SingleLandmarkData], bool]):
+        super().__init__(dataloader)
+        self._predicate_name = predicate.__name__ if hasattr(predicate, "__name__") else str(predicate)
+        self._reindex = [
+            idx
+            for idx in self._wrapped_dataloader.idx_sampler
+            if predicate(self._wrapped_dataloader.get_single_element(idx))
+        ]