feat: Adds update_dataset_from_dir

Author: nicolastomeo

nucleus/__init__.py

@@ -1239,29 +1239,43 @@ def _set_api_key(self, api_key):
 
         return api_key
 
-    def create_dataset_from_dir(
+    def _create_or_update_dataset_from_dir(
         self,
         dirname: str,
+        dataset_id: Optional[str] = None,
         dataset_name: Optional[str] = None,
         use_privacy_mode: bool = False,
         privacy_mode_proxy: str = "",
         allowed_file_types: Tuple[str, ...] = ("png", "jpg", "jpeg"),
         skip_size_warning: bool = False,
+        update_items: bool = False,
     ) -> Union[Dataset, None]:
         """
-        Create a dataset by recursively crawling through a directory.
+        Create or update dataset by recursively crawling through a directory.
         A DatasetItem will be created for each unique image found.
+        In case of update the existing items are skipped or updated depending on update_items param
 
-        Parameters:
+        Args:
             dirname: Where to look for image files, recursively
+            dataset_id: Dataset Id for update
             dataset_name: If none is given, the parent folder name is used
-            use_privacy_mode: Whether the dataset should be treated as privacy
+            use_privacy_mode: Whether the dataset should be treated as privacy (ignored if dataset being updated)
             privacy_mode_proxy: Endpoint that serves image files for privacy mode, ignore if not using privacy mode.
                 The proxy should work based on the relative path of the images in the directory.
             allowed_file_types: Which file type extensions to search for, ie: ('jpg', 'png')
             skip_size_warning: If False, it will throw an error if the script globs more than 500 images. This is a safety check in case the dirname has a typo, and grabs too much data.
-        """
+            update_items: Whether to update items in existing dataset
+
+        Returns:
+            :class: `Union[Dataset, None]`: Updated dataset or None if there wasn't a dataset to create
 
+        """
+        if dataset_id:
+            existing_dataset = self.get_dataset(dataset_id)
+            # fetch dataset use_privacy_mode for existence check
+            use_privacy_mode = existing_dataset.use_privacy_mode
+        else:
+            existing_dataset = None
         if use_privacy_mode:
             assert (
                 privacy_mode_proxy
@@ -1273,9 +1287,6 @@ def create_dataset_from_dir(
             raise ValueError(
                 f"Given directory name: {dirname} does not exists. Searched in {_dirname}"
             )
-
-        folder_name = os.path.basename(_dirname.rstrip("/"))
-        dataset_name = dataset_name or folder_name
         items = create_items_from_folder_crawl(
             _dirname,
             allowed_file_types,
@@ -1285,15 +1296,88 @@ def create_dataset_from_dir(
 
         if len(items) == 0:
             print(f"Did not find any items in {dirname}")
-            return None
+            return existing_dataset
 
         if len(items) > GLOB_SIZE_THRESHOLD_CHECK and not skip_size_warning:
             raise Exception(
                 f"Found over {GLOB_SIZE_THRESHOLD_CHECK} items in {dirname}. If this is intended, set skip_size_warning=True when calling this function."
             )
 
-        dataset = self.create_dataset(
-            name=dataset_name, use_privacy_mode=use_privacy_mode
-        )
-        dataset.append(items, asynchronous=False)
+        if existing_dataset:
+            dataset = existing_dataset
+        else:
+            folder_name = os.path.basename(_dirname.rstrip("/"))
+            dataset_name = dataset_name or folder_name
+            dataset = self.create_dataset(
+                name=dataset_name, use_privacy_mode=use_privacy_mode
+            )
+        dataset.append(items, asynchronous=False, update=update_items)
         return dataset
+
+    def create_dataset_from_dir(
+        self,
+        dirname: str,
+        dataset_name: Optional[str] = None,
+        use_privacy_mode: bool = False,
+        privacy_mode_proxy: str = "",
+        allowed_file_types: Tuple[str, ...] = ("png", "jpg", "jpeg"),
+        skip_size_warning: bool = False,
+    ) -> Union[Dataset, None]:
+        """
+        Create a dataset by recursively crawling through a directory.
+        A DatasetItem will be created for each unique image found.
+
+        Parameters:
+            dirname: Where to look for image files, recursively
+            dataset_name: If none is given, the parent folder name is used
+            use_privacy_mode: Whether the dataset should be treated as privacy
+            privacy_mode_proxy: Endpoint that serves image files for privacy mode, ignore if not using privacy mode.
+                The proxy should work based on the relative path of the images in the directory.
+            allowed_file_types: Which file type extensions to search for, ie: ('jpg', 'png')
+            skip_size_warning: If False, it will throw an error if the script globs more than 500 images. This is a safety check in case the dirname has a typo, and grabs too much data.
+        """
+        return self._create_or_update_dataset_from_dir(
+            dirname,
+            dataset_name=dataset_name,
+            use_privacy_mode=use_privacy_mode,
+            privacy_mode_proxy=privacy_mode_proxy,
+            allowed_file_types=allowed_file_types,
+            skip_size_warning=skip_size_warning,
+        )
+
+    def update_dataset_from_dir(
+        self,
+        dirname: str,
+        dataset_id: str,
+        privacy_mode_proxy: str = "",
+        allowed_file_types: Tuple[str, ...] = ("png", "jpg", "jpeg"),
+        skip_size_warning: bool = False,
+        update_items: bool = False,
+    ) -> Dataset:
+        """
+        update dataset by recursively crawling through a directory.
+        A DatasetItem will be created for each unique image found.
+        The existing items are skipped or updated depending on update_items param
+
+        Args:
+            dirname: Where to look for image files, recursively
+            dataset_id: ID of existing dataset to update
+            privacy_mode_proxy: Endpoint that serves image files for privacy mode, ignore if not using privacy mode.
+                The proxy should work based on the relative path of the images in the directory.
+            allowed_file_types: Which file type extensions to search for, ie: ('jpg', 'png')
+            skip_size_warning: If False, it will throw an error if the script globs more than 500 images. This is a safety check in case the dirname has a typo, and grabs too much data.
+            update_items: Whether to update items in existing dataset
+
+        Returns:
+            :class:`Dataset`: Updated dataset
+        """
+        updated_dataset = self._create_or_update_dataset_from_dir(
+            dirname,
+            dataset_id=dataset_id,
+            privacy_mode_proxy=privacy_mode_proxy,
+            allowed_file_types=allowed_file_types,
+            skip_size_warning=skip_size_warning,
+            update_items=update_items,
+        )
+        assert updated_dataset is not None
+        return updated_dataset

tests/helpers.py

@@ -459,6 +459,7 @@ def reference_id_from_url(url):
 
 this_dir = os.path.dirname(os.path.realpath(__file__))
 TEST_LOCAL_MASK_URL = os.path.join(this_dir, "testdata/000000000285.png")
+TEST_LOCAL_TESTDIR = os.path.join(this_dir, "testdata/testdir")
 
 
 NUM_VALID_SEGMENTATIONS_IN_MAIN_DATASET = len(TEST_DATASET_ITEMS)

tests/test_dataset.py

@@ -1,5 +1,7 @@
 import copy
+import glob
 import math
+import os
 
 import pytest
 
@@ -44,6 +46,7 @@
     TEST_VIDEO_SCENES,
     assert_partial_equality,
     reference_id_from_url,
+    TEST_LOCAL_TESTDIR,
 )
 
 
@@ -611,3 +614,35 @@ def test_query(CLIENT):
     with pytest.raises(NucleusAPIError):
         for qi in dataset.query_items("annotations.count bad syntax"):
             print(qi)  # unreachable, just need to yield an item from generator
+
+
+@pytest.mark.integration
+def test_create_update_dataset_from_dir(CLIENT):
+    reference_ids = set()
+    for file_type in ["png", "jpeg"]:
+        pathname = os.path.join(TEST_LOCAL_TESTDIR, f"**/*.{file_type}")
+        reference_ids.update(
+            path.replace(TEST_LOCAL_TESTDIR + "/", "")
+            for path in glob.glob(pathname=pathname, recursive=True)
+        )
+    dataset = CLIENT.create_dataset_from_dir(
+        TEST_LOCAL_TESTDIR, allowed_file_types=tuple(["exe"])
+    )
+    assert dataset is None
+    dataset = CLIENT.create_dataset_from_dir(
+        TEST_LOCAL_TESTDIR, allowed_file_types=tuple(["png"])
+    )
+    dataset_items = dataset.items
+    assert len(dataset_items) == 1
+    assert dataset_items[0].reference_id in reference_ids
+    dataset = CLIENT.update_dataset_from_dir(
+        TEST_LOCAL_TESTDIR,
+        dataset_id=dataset.id,
+        allowed_file_types=tuple(["png", "jpeg"]),
+    )
+    dataset_items = dataset.items
+    assert len(dataset_items) == 2
+    for dataset_item in dataset_items:
+        assert dataset_item.reference_id in reference_ids
+        reference_ids.remove(dataset_item.reference_id)
+    CLIENT.delete_dataset(dataset.id)