[ADD][FEAT] Allow additional params to init DatasetHandler

wbenbihi · Aug 19, 2022 · 6723c8e · 6723c8e
1 parent 414f725
commit 6723c8e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/config/config.default.yaml b/config/config.default.yaml
@@ -48,6 +48,8 @@ data:
 
 dataset:
   object: "hourglass_tensorflow.datasets.HTFBaseDatasetHandler"
+  params:
+    mockup_param: to_pass_to_object_init_method
   sets:
     test: true
     validation: true

diff --git a/hourglass_tensorflow/utils/config.py b/hourglass_tensorflow/utils/config.py
@@ -46,8 +46,16 @@ class HTFDatasetSetsConfig(BaseModel):
     test: bool = False
 
 
+class HTFDatasetParamsConfig(BaseModel):
+    class Config:
+        extra = "allow"
+
+
 class HTFDatasetConfig(BaseModel):
     object: str
+    params: Optional[HTFDatasetParamsConfig] = Field(
+        default_factory=HTFDatasetParamsConfig
+    )
     sets: Optional[HTFDatasetSetsConfig]
     split: Optional[HTFDatasetSplitConfig] = Field(
         default_factory=HTFDatasetSplitConfig
@@ -407,7 +415,10 @@ def _load_dataset_object(self) -> None:
     def prepare_dataset(self) -> None:
         self._load_dataset_object()
         self.dataset_handler: "HTFBaseDatasetHandler" = self._metadata.dataset_object(
-            dataset=self._labels_df, config=self.config.dataset, global_config=self
+            dataset=self._labels_df,
+            config=self.config.dataset,
+            global_config=self,
+            **self.config.dataset.params.dict(),
         )
         self.dataset_handler.split_train_test()