Search images in subfolders

Author: hovnatan

docs/source/superannotate.sdk.rst

@@ -75,6 +75,7 @@ ______
 
 .. _ref_search_images:
 .. autofunction:: superannotate.search_images
+.. autofunction:: superannotate.search_images_all_folders
 .. autofunction:: superannotate.get_image_metadata
 .. autofunction:: superannotate.get_image_bytes
 .. autofunction:: superannotate.download_image

superannotate/__init__.py

@@ -49,8 +49,8 @@ def consensus(*args, **kwargs):
     create_fuse_image, delete_image, download_image, download_image_annotations,
     download_image_preannotations, get_image_annotations, get_image_bytes,
     get_image_metadata, get_image_preannotations, search_images,
-    set_image_annotation_status, set_images_annotation_statuses,
-    upload_image_annotations
+    search_images_all_folders, set_image_annotation_status,
+    set_images_annotation_statuses, upload_image_annotations
 )
 from .db.project_api import (
     create_folder, delete_folders, get_folder_metadata,

superannotate/db/images.py

@@ -24,13 +24,32 @@
     get_annotation_classes_id_to_name, get_annotation_classes_name_to_id,
     search_annotation_classes
 )
-from .project_api import get_project_and_folder_metadata
+from .project_api import get_project_and_folder_metadata, get_project_metadata_bare
 
 logger = logging.getLogger("superannotate-python-sdk")
 
 _api = API.get_instance()
 
 
+def _get_project_root_folder_id(project):
+    """Get root folder ID
+    Returns
+    -------
+    int
+        Root folder ID
+    """
+    params = {'team_id': project['team_id']}
+    response = _api.send_request(
+        req_type='GET', path=f'/project/{project["id"]}', params=params
+    )
+    if not response.ok:
+        raise SABaseException(response.status_code, response.text)
+
+    response = response.json()
+
+    return response['folder_id']
+
+
 def search_images(
     project,
     image_name_prefix=None,
@@ -63,7 +82,7 @@ def search_images(
     if project_folder is not None:
         project_folder_id = project_folder["id"]
     else:
-        project_folder_id = None
+        project_folder_id = _get_project_root_folder_id(project)
 
     result_list = []
     params = {
@@ -79,7 +98,7 @@ def search_images(
     total_images = 0
     while True:
         response = _api.send_request(
-            req_type='GET', path='/images', params=params
+            req_type='GET', path='/images-folders', params=params
         )
         if not response.ok:
             raise SABaseException(
@@ -115,6 +134,83 @@ def process_result(x):
         return result_list
 
 
+def search_images_all_folders(
+    project,
+    image_name_prefix=None,
+    annotation_status=None,
+    return_metadata=False
+):
+    """Search images by name_prefix (case-insensitive) and annotation status in 
+    project and all of its folders
+
+    :param project: project name
+    :type project: str
+    :param image_name_prefix: image name prefix for search
+    :type image_name_prefix: str
+    :param annotation_status: if not None, annotation statuses of images to filter,
+                              should be one of NotStarted InProgress QualityCheck Returned Completed Skipped
+    :type annotation_status: str
+
+    :param return_metadata: return metadata of images instead of names
+    :type return_metadata: bool
+
+    :return: metadata of found images or image names
+    :rtype: list of dicts or strs
+    """
+    project = get_project_metadata_bare(project)
+    team_id, project_id = project["team_id"], project["id"]
+    if annotation_status is not None:
+        annotation_status = common.annotation_status_str_to_int(
+            annotation_status
+        )
+
+    project_folder_id = _get_project_root_folder_id(project)
+
+    result_list = []
+    params = {
+        'team_id': team_id,
+        'project_id': project_id,
+        'annotation_status': annotation_status,
+        'offset': 0,
+        'folder_id': project_folder_id
+    }
+    if image_name_prefix is not None:
+        params['name'] = image_name_prefix
+    total_images = 0
+    while True:
+        response = _api.send_request(
+            req_type='GET', path='/images', params=params
+        )
+        if not response.ok:
+            raise SABaseException(
+                response.status_code, "Couldn't search images " + response.text
+            )
+        response = response.json()
+        results_images = response["data"]
+        for r in results_images:
+            if return_metadata:
+                result_list.append(r)
+            else:
+                result_list.append(r["name"])
+
+        total_images += len(results_images)
+        if response["count"] <= total_images:
+            break
+        params["offset"] = total_images
+
+    if return_metadata:
+
+        def process_result(x):
+            x["annotation_status"] = common.annotation_status_int_to_str(
+                x["annotation_status"]
+            )
+            return x
+
+        return list(map(process_result, result_list))
+    else:
+        return result_list
+
+
 def get_image_metadata(project, image_names, return_dict_on_single_output=True):
     """Returns image metadata
 

superannotate/db/projects.py

@@ -33,7 +33,7 @@
     fill_class_and_attribute_ids, get_annotation_classes_name_to_id,
     search_annotation_classes
 )
-from .images import get_image_metadata, search_images
+from .images import get_image_metadata, search_images, search_images_all_folders
 from .project_api import (
     get_project_and_folder_metadata, get_project_metadata_bare,
     get_project_metadata_with_users
@@ -186,16 +186,21 @@ def rename_project(project, new_name):
     )
 
 
-def get_project_image_count(project):
+def get_project_image_count(project, with_all_subfolders=False):
     """Returns number of images in the project.
 
     :param project: project name or folder path (e.g., "project1/folder1")
     :type project: str
+    :param with_all_subfolders: enables recursive folder counting
+    :type with_all_subfolders: bool
 
     :return: number of images in the project
     :rtype: int
     """
-    return len(search_images(project))
+    if not with_all_subfolders:
+        return len(search_images(project))
+    else:
+        return len(search_images_all_folders(project))
 
 
 def upload_video_to_project(

superannotate/version.py

@@ -1 +1 @@
-__version__ = "4.1.0b11"
+__version__ = "4.1.0b12"

tests/test_folders.py

@@ -42,6 +42,9 @@ def test_basic_folders(tmpdir):
     images = sa.search_images(project + "/folder1", "example_image_1")
     assert len(images) == 0
 
+    images = sa.search_images_all_folders(project, "example_image_1")
+    assert len(images) == 1
+
     folder = sa.get_folder_metadata(project, "folder1")
     assert isinstance(folder, dict)
     assert folder["name"] == "folder1"
@@ -210,6 +213,9 @@ def test_project_folder_image_count(tmpdir):
     num_images = sa.get_project_image_count(project + "/folder1")
     assert num_images == 4
 
+    num_images = sa.get_project_image_count(project, with_all_subfolders=True)
+    assert num_images == 8
+
 
 def test_delete_images(tmpdir):
     PROJECT_NAME = "test delete folder images"
@@ -319,6 +325,17 @@ def test_move_images(tmpdir):
     num_images = sa.get_project_image_count(project)
     assert num_images == 3
 
+    num_images = sa.get_project_image_count(
+        PROJECT_NAME, with_all_subfolders=True
+    )
+    assert num_images == 4
+
+    images = sa.search_images_all_folders(PROJECT_NAME)
+    assert images == [
+        "example_image_1.jpg", "example_image_2.jpg", "example_image_3.jpg",
+        "example_image_4.jpg"
+    ]
+
 
 def test_move_images2(tmpdir):
     PROJECT_NAME = "test move folder images2"