Start work on annotation folders

Author: hovnatan

superannotate/db/folders.py

@@ -1,7 +1,11 @@
+import logging
+
 from ..api import API
 from ..exceptions import SABaseException
 from .project_api import get_project_metadata_bare
 
+logger = logging.getLogger("superannotate-python-sdk")
+
 _api = API.get_instance()
 
 
@@ -27,36 +31,35 @@ def search_folders(project, folder_name=None, return_metadata=False):
         project = get_project_metadata_bare(project)
     team_id, project_id = project["team_id"], project["id"]
     result_list = []
-    params = {'team_id': team_id, 'project_id': project_id, 'offset': 0}
-    total_got = 0
+    params = {
+        'team_id': team_id,
+        'project_id': project_id,
+        'offset': 0,
+        'name': folder_name,
+        'is_root': 0
+    }
     total_folders = 0
     while True:
         response = _api.send_request(
-            req_type='GET', path='/images', params=params
+            req_type='GET', path='/folders', params=params
         )
         if not response.ok:
             raise SABaseException(
-                response.status_code, "Couldn't search images " + response.text
+                response.status_code, "Couldn't search folders " + response.text
             )
         response = response.json()
-        images = response["images"]
-        folders = response["folders"]
-
-        results_folders = folders["data"]
+        results_folders = response["data"]
         for r in results_folders:
-            if folder_name is not None and r["name"] != folder_name:
-                continue
             if return_metadata:
                 result_list.append(r)
             else:
                 result_list.append(r["name"])
 
         total_folders += len(results_folders)
-        if folders["count"] <= total_folders:
+        if response["count"] <= total_folders:
             break
 
-        total_got += len(images["data"]) + len(results_folders)
-        params["offset"] = total_got
+        params["offset"] = total_folders
 
     return result_list
 
@@ -84,4 +87,9 @@ def create_folder(project, folder_name):
             response.status_code, "Couldn't create project " + response.text
         )
     res = response.json()
+    if res["name"] != folder_name:
+        logger.warning(
+            "Created folder has name %s, since folder with name %s already existed.",
+            res["name"], folder_name
+        )
     return res

superannotate/db/images.py

@@ -64,16 +64,19 @@ def search_images(
         )
 
     if project_folder is not None:
-        folder = get_folder_metadata(project, project_folder)
-        folder = folder["id"]
+        if not isinstance(project_folder, dict):
+            project_folder = get_folder_metadata(project, project_folder)
+        project_folder_id = project_folder["id"]
+    else:
+        project_folder_id = None
 
     result_list = []
     params = {
         'team_id': team_id,
         'project_id': project_id,
         'annotation_status': annotation_status,
         'offset': 0,
-        'folder_id': project_folder
+        'folder_id': project_folder_id
     }
     if image_name_prefix is not None:
         params['name'] = image_name_prefix
@@ -137,17 +140,30 @@ def get_image_metadata(
     if isinstance(image_names, str):
         image_names = [image_names]
 
+    if project_folder is not None:
+        if not isinstance(project_folder, dict):
+            project_folder = get_folder_metadata(project, project_folder)
+        project_folder_id = project_folder["id"]
+    else:
+        project_folder_id = None
+
     json_req = {
         'project_id': project['id'],
         'team_id': _api.team_id,
-        'names': image_names
-        'folder_id': project_folder
+        'names': image_names,
     }
+    if project_folder_id is not None:
+        json_req["folder_id"] = project_folder_id
     response = _api.send_request(
         req_type='POST',
         path='/images/getBulk',
         json_req=json_req,
     )
+    if not response.ok:
+        raise SABaseException(
+            response.status_code,
+            "Couldn't get image metadata. " + response.text
+        )
 
     metadata_raw = response.json()
     metadata_without_deleted = []
@@ -649,7 +665,7 @@ def get_image_preannotations(project, image_name):
     return _get_image_pre_or_annotations(project, image_name, "pre")
 
 
-def get_image_annotations(project, image_name):
+def get_image_annotations(project, image_name, project_folder=None):
     """Get annotations of the image.
 
     :param project: project name or metadata of the project
@@ -664,10 +680,14 @@ def get_image_annotations(project, image_name):
         "annotation_mask_filename": mask filename on server
     :rtype: dict
     """
-    return _get_image_pre_or_annotations(project, image_name, "")
+    return _get_image_pre_or_annotations(
+        project, image_name, "", project_folder
+    )
 
 
-def _get_image_pre_or_annotations(project, image_name, pre, project_type=None):
+def _get_image_pre_or_annotations(
+    project, image_name, pre, project_type=None, project_folder=None
+):
     image = get_image_metadata(project, image_name)
     team_id, project_id, image_id, folder_id = image["team_id"], image[
         "project_id"], image["id"], image['folder_id']

superannotate/db/projects.py

@@ -642,7 +642,7 @@ def get_image_array_to_upload(
 def __upload_images_to_aws_thread(
     res, img_paths, project, annotation_status, prefix, thread_id, chunksize,
     couldnt_upload, uploaded, tried_upload, image_quality_in_editor,
-    from_s3_bucket, folder_id
+    from_s3_bucket, project_folder_id
 ):
     len_img_paths = len(img_paths)
     start_index = thread_id * chunksize
@@ -699,7 +699,7 @@ def __upload_images_to_aws_thread(
                     __create_image(
                         uploaded_imgs_info[0], uploaded_imgs_info[1], project,
                         annotation_status, prefix, uploaded_imgs_info[2],
-                        folder_id
+                        project_folder_id
                     )
                 except SABaseException as e:
                     couldnt_upload[thread_id] += uploaded_imgs
@@ -711,7 +711,7 @@ def __upload_images_to_aws_thread(
     try:
         __create_image(
             uploaded_imgs_info[0], uploaded_imgs_info[1], project,
-            annotation_status, prefix, uploaded_imgs_info[2], folder_id
+            annotation_status, prefix, uploaded_imgs_info[2], project_folder_id
         )
     except SABaseException as e:
         couldnt_upload[thread_id] += uploaded_imgs
@@ -722,7 +722,7 @@ def __upload_images_to_aws_thread(
 
 def __create_image(
     img_names, img_paths, project, annotation_status, remote_dir, sizes,
-    folder_id
+    project_folder_id
 ):
     if len(img_paths) == 0:
         return
@@ -737,8 +737,8 @@ def __create_image(
         "annotation_status": annotation_status,
         "meta": {}
     }
-    if folder_id is not None:
-        data["folder_id"] = folder_id
+    if project_folder_id is not None:
+        data["folder_id"] = project_folder_id
     for img_name, img_path, size in zip(img_names, img_paths, sizes):
         img_name_uuid = Path(img_path).name
         remote_path = remote_dir + f"{img_name_uuid}"
@@ -838,10 +838,10 @@ def upload_images_to_project(
         raise SABaseException(
             response.status_code, "Couldn't get upload token " + response.text
         )
-    if folder is not None:
-        if not isinstance(folder, dict):
-            folder = get_folder_metadata(project, folder)
-        folder = folder["id"]
+    if project_folder is not None:
+        if not isinstance(project_folder, dict):
+            project_folder = get_folder_metadata(project, project_folder)
+        project_folder = project_folder["id"]
     res = response.json()
     prefix = res['filePath']
     tqdm_thread = threading.Thread(
@@ -858,7 +858,7 @@ def upload_images_to_project(
             args=(
                 res, img_paths, project, annotation_status, prefix, thread_id,
                 chunksize, couldnt_upload, uploaded, tried_upload,
-                image_quality_in_editor, from_s3_bucket, folder
+                image_quality_in_editor, from_s3_bucket, project_folder
             ),
             daemon=True
         )
@@ -1172,7 +1172,7 @@ def upload_images_from_azure_blob_to_project(
 def __upload_annotations_thread(
     team_id, project_id, project_type, anns_filenames, folder_path,
     annotation_classes_dict, pre, thread_id, chunksize, missing_images,
-    couldnt_upload, uploaded, from_s3_bucket, project_folder
+    couldnt_upload, uploaded, from_s3_bucket, project_folder_id
 ):
     NUM_TO_SEND = 500
     len_anns = len(anns_filenames)
@@ -1196,7 +1196,9 @@ def __upload_annotations_thread(
             image_name = anns_filenames[j][:-len_postfix_json]
             names.append(image_name)
         try:
-            metadatas = get_image_metadata({"id": project_id}, names, False)
+            metadatas = get_image_metadata(
+                {"id": project_id}, names, False, {"id": project_folder_id}
+            )
         except SABaseException:
             metadatas = []
         names_in_metadatas = [metadata["name"] for metadata in metadatas]
@@ -1216,7 +1218,7 @@ def __upload_annotations_thread(
             "project_id": project_id,
             "team_id": team_id,
             "ids": [metadata["id"] for metadata in metadatas],
-            "folder_id": project_folder
+            "folder_id": project_folder_id
         }
         endpoint = '/images/getAnnotationsPathsAndTokens' if pre == "" else '/images/getPreAnnotationsPathsAndTokens'
         response = _api.send_request(
@@ -1338,9 +1340,16 @@ def _upload_pre_or_annotations_from_folder_to_project(
     if not isinstance(project, dict):
         project = get_project_metadata_bare(project)
 
+    if project_folder is not None:
+        if not isinstance(project_folder, dict):
+            project_folder = get_folder_metadata(project, project_folder)
+        project_folder_id = project_folder["id"]
+    else:
+        project_folder_id = None
+
     return _upload_annotations_from_folder_to_project(
         project, folder_path, pre, from_s3_bucket, recursive_subfolders,
-        project_folder
+        project_folder_id
     )
 
 
@@ -1350,7 +1359,7 @@ def _upload_annotations_from_folder_to_project(
     pre,
     from_s3_bucket=None,
     recursive_subfolders=False,
-    project_folder=None
+    project_folder_id=None
 ):
     return_result = []
     if from_s3_bucket is not None:
@@ -1362,7 +1371,7 @@ def _upload_annotations_from_folder_to_project(
                 if path.is_dir():
                     return_result += _upload_annotations_from_folder_to_project(
                         project, path, pre, from_s3_bucket,
-                        recursive_subfolders, project_folder
+                        recursive_subfolders, project_folder_id
                     )
         else:
             s3_client = boto3.client('s3')
@@ -1374,7 +1383,7 @@ def _upload_annotations_from_folder_to_project(
                 for o in results:
                     return_result += _upload_annotations_from_folder_to_project(
                         project, o.get('Prefix'), pre, from_s3_bucket,
-                        recursive_subfolders, project_folder
+                        recursive_subfolders, project_folder_id
                     )
 
     team_id, project_id, project_type = project["team_id"], project[
@@ -1449,7 +1458,7 @@ def _upload_annotations_from_folder_to_project(
                 team_id, project_id, project_type, annotations_filenames,
                 folder_path, annotation_classes_dict, pre, thread_id, chunksize,
                 missing_image, couldnt_upload, uploaded, from_s3_bucket,
-                project_folder
+                project_folder_id
             ),
             daemon=True
         )

tests/test_folders.py

@@ -1,13 +1,13 @@
 from pathlib import Path
-import json
-from superannotate.exceptions import SABaseException
 
 import pytest
-
 import superannotate as sa
+from superannotate.exceptions import SABaseException
 
 PROJECT_NAME1 = "test folder simple"
-FROM_FOLDER = "./tests/sample_project_vector"
+PROJECT_NAME2 = "test folder annotations"
+
+FROM_FOLDER = Path("./tests/sample_project_vector")
 
 
 def test_basic_folders(tmpdir):
@@ -80,4 +80,63 @@ def test_basic_folders(tmpdir):
             annotation_status="InProgress",
             project_folder="folder2"
         )
-    assert 'Folder not found' in str(e)
+    assert 'Folder not found' in str(e)
+
+    folder_metadata = sa.create_folder(project, "folder2")
+    assert folder_metadata["name"] == "folder2"
+
+    folders = sa.search_folders(project)
+    assert len(folders) == 2
+
+    folders = sa.search_folders(project, folder_name="folder")
+    assert len(folders) == 2
+
+    folders = sa.search_folders(project, folder_name="folder2")
+    assert len(folders) == 1
+    assert folders[0] == "folder2"
+
+    folders = sa.search_folders(project, folder_name="folder1")
+    assert len(folders) == 1
+    assert folders[0] == "folder1"
+
+    folders = sa.search_folders(project, folder_name="old")
+    assert len(folders) == 2
+
+
+def test_folder_annotations(tmpdir):
+    tmpdir = Path(tmpdir)
+
+    projects_found = sa.search_projects(PROJECT_NAME2, return_metadata=True)
+    for pr in projects_found:
+        sa.delete_project(pr)
+
+    project = sa.create_project(PROJECT_NAME2, 'test', 'Vector')
+    sa.upload_images_from_folder_to_project(
+        project, FROM_FOLDER, annotation_status="InProgress"
+    )
+    sa.create_annotation_classes_from_classes_json(
+        project, FROM_FOLDER / "classes" / "classes.json"
+    )
+    folder_metadata = sa.create_folder(project, "folder1")
+    assert folder_metadata["name"] == "folder1"
+
+    folders = sa.search_folders(project, return_metadata=True)
+    assert len(folders) == 1
+
+    sa.upload_images_from_folder_to_project(
+        project,
+        FROM_FOLDER,
+        annotation_status="InProgress",
+        project_folder=folders[0]
+    )
+    sa.upload_annotations_from_folder_to_project(
+        project, FROM_FOLDER, project_folder=folders[0]
+    )
+    print(folders[0])
+    annot = sa.get_image_annotations(project, "example_image_1.jpg")
+    assert len(annot["annotation_json"]["instances"]) == 0
+
+    annot = sa.get_image_annotations(
+        project, "example_image_1.jpg", project_folder="folder1"
+    )
+    assert len(annot["annotation_json"]["instances"]) > 0