added gpu_id as parameter, refactored some notebooks

examples/ffhq-filtering-caching.ipynb

@@ -7,7 +7,7 @@
    "outputs": [],
    "source": [
     "from loreal_poc.dataloaders.loaders import DataLoaderFFHQ\n",
-    "from loreal_poc.dataloaders.wrappers import FilteringDataLoader, CachedDataLoader\n",
+    "from loreal_poc.dataloaders.wrappers import FilteredDataLoader, CachedDataLoader\n",
     "from loreal_poc.dataloaders.base import DataLoaderWrapper\n",
     "from loreal_poc.visualisation.draw import draw_marks\n",
     "\n",
@@ -215,7 +215,7 @@
     }
    ],
    "source": [
-    "odds = FilteringDataLoader(dl, lambda elt: elt[2][\"type\"] == \"odd\")\n",
+    "odds = FilteredDataLoader(dl, lambda elt: elt[2][\"type\"] == \"odd\")\n",
     "len(odds)"
    ]
   },
@@ -256,7 +256,7 @@
     }
    ],
    "source": [
-    "evens = FilteringDataLoader(dl, lambda elt: elt[2][\"type\"] == \"even\")\n",
+    "evens = FilteredDataLoader(dl, lambda elt: elt[2][\"type\"] == \"even\")\n",
     "len(evens)"
    ]
   },

examples/ffhq-filtering.ipynb

@@ -7,7 +7,7 @@
    "outputs": [],
    "source": [
     "from loreal_poc.dataloaders.loaders import DataLoaderFFHQ\n",
-    "from loreal_poc.dataloaders.wrappers import FilteringDataLoader, CachedDataLoader\n",
+    "from loreal_poc.dataloaders.wrappers import FilteredDataLoader, CachedDataLoader\n",
     "from loreal_poc.dataloaders.base import DataLoaderWrapper\n",
     "from loreal_poc.visualisation.draw import draw_marks\n",
     "\n",
@@ -35,7 +35,7 @@
    "source": [
     "img, marks, meta = dl[0]\n",
     "\n",
-    "fdl = FilteringDataLoader(dl, lambda elt: elt[2][\"faceAttributes\"][\"headPose\"][\"yaw\"] < 0)\n",
+    "fdl = FilteredDataLoader(dl, lambda elt: elt[2][\"faceAttributes\"][\"headPose\"][\"yaw\"] < 0)\n",
     "len(fdl)"
    ]
   },

loreal_poc/dataloaders/wrappers.py

@@ -159,15 +159,24 @@ def __init__(self, dataloader: DataIteratorBase, predicate: Callable[[SingleLand
 
 
 class HeadPoseDataLoader(DataLoaderWrapper):
-    def __init__(self, dataloader: DataIteratorBase) -> None:
+    def __init__(self, dataloader: DataIteratorBase, gpu_id: int = -1) -> None:
+        """A dataloader that estimates the head pose in images using the SixDRepNet model
+
+        Args:
+            dataloader (DataIteratorBase): the wrapped dataloader.
+            gpu_id (int, optional): Enable the usage of GPUs. Defaults to -1 (CPU).
+
+        Raises:
+            GiskardImportError: Error to signal a missing package
+        """
         try:
             from sixdrepnet import SixDRepNet
         except ImportError as e:
             raise GiskardImportError("sixdrepnet") from e
 
         super().__init__(dataloader)
 
-        self.pose_detection_model = SixDRepNet(gpu_id=-1)
+        self.pose_detection_model = SixDRepNet(gpu_id=gpu_id)
 
     @property
     def name(self):

loreal_poc/utils/errors.py

@@ -1,3 +1,4 @@
 class GiskardImportError(ImportError):
     def __init__(self, missing_package: str) -> None:
+        super().__init__()
         self.msg = f"The '{missing_package}' Python package is not installed; please execute 'pip install {missing_package}' to obtain it."