update labeldict

mwalmsley · Jan 15, 2024 · d3b4398 · d3b4398
1 parent 23062a5
commit d3b4398
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/zoobot/pytorch/datasets/webdatamodule.py b/zoobot/pytorch/datasets/webdatamodule.py
@@ -217,3 +217,11 @@ def label_transform(label_dict):
         return label_transform
     else:
         return identity  # do nothing
+
+def dict_to_filled_dict_factory(label_cols):
+    # might be a little slow, but very safe
+    def label_transform(label_dict: dict):
+        # modifies inplace with 0 iff key missing
+        [label_dict.setdefault(col, 0) for col in label_cols]
+        return label_dict
+    return label_transform