Fix aws region - root folder id

Author: dshabin

superannotate/db/images.py

@@ -25,6 +25,7 @@
     search_annotation_classes
 )
 from .project_api import get_project_and_folder_metadata, get_project_metadata_bare
+from .utils import _get_boto_session_by_credentials
 
 logger = logging.getLogger("superannotate-python-sdk")
 
@@ -38,16 +39,17 @@ def get_project_root_folder_id(project):
     int
         Root folder ID
     """
-    params = {'team_id': project['team_id']}
-    response = _api.send_request(
-        req_type='GET', path=f'/project/{project["id"]}', params=params
-    )
+    params = {
+        'team_id': project['team_id'],
+        'project_id': project['id'],
+        'is_root': 1
+    }
+    response = _api.send_request(req_type='GET', path='/folders', params=params)
     if not response.ok:
         raise SABaseException(response.status_code, response.text)
 
     response = response.json()
-
-    return response['folder_id']
+    return response['data'][0]['id']
 
 
 def search_images(
@@ -232,7 +234,10 @@ def get_image_metadata(project, image_names, return_dict_on_single_output=True):
         project_folder_id = None
 
     chunk_size = 500
-    chunks = [image_names[i:i + chunk_size] for i in range(0, len(image_names), chunk_size)]
+    chunks = [
+        image_names[i:i + chunk_size]
+        for i in range(0, len(image_names), chunk_size)
+    ]
 
     json_req = {
         'project_id': project['id'],
@@ -249,13 +254,16 @@ def get_image_metadata(project, image_names, return_dict_on_single_output=True):
             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)
+            raise SABaseException(
+                response.status_code,
+                "Couldn't get image metadata. " + response.text
+            )
         metadata_raw += response.json()
 
     metadata_without_deleted = []
-    metadata_without_deleted = [ i for i in metadata_raw if i['delete'] != 1 ]
+    metadata_without_deleted = [i for i in metadata_raw if i['delete'] != 1]
 
     if len(metadata_without_deleted) == 0:
         raise SABaseException(
@@ -1001,11 +1009,7 @@ def upload_image_annotations(
         )
     res = response.json()
     res_json = res['annotation_json_path']
-    s3_session = boto3.Session(
-        aws_access_key_id=res_json['accessKeyId'],
-        aws_secret_access_key=res_json['secretAccessKey'],
-        aws_session_token=res_json['sessionToken']
-    )
+    s3_session = _get_boto_session_by_credentials(res_json)
     s3_resource = s3_session.resource('s3')
     bucket = s3_resource.Bucket(res_json["bucket"])
     bucket.put_object(
@@ -1018,11 +1022,7 @@ def upload_image_annotations(
             with open(mask, "rb") as f:
                 mask = io.BytesIO(f.read())
         res_mask = res['annotation_bluemap_path']
-        s3_session = boto3.Session(
-            aws_access_key_id=res_mask['accessKeyId'],
-            aws_secret_access_key=res_mask['secretAccessKey'],
-            aws_session_token=res_mask['sessionToken']
-        )
+        s3_session = _get_boto_session_by_credentials(res_mask)
         bucket = s3_resource.Bucket(res_mask["bucket"])
         bucket.put_object(Key=res_mask['filePath'], Body=mask)
 

superannotate/db/project_images.py

@@ -20,7 +20,7 @@
     get_project_default_image_quality_in_editor, upload_image_array_to_s3,
     _get_available_image_counts
 )
-from .utils import _get_upload_auth_token
+from .utils import _get_upload_auth_token, _get_boto_session_by_credentials
 
 logger = logging.getLogger("superannotate-python-sdk")
 _api = API.get_instance()
@@ -104,11 +104,7 @@ def upload_image_to_project(
     params = {'team_id': team_id, 'folder_id': folder_id}
     res = _get_upload_auth_token(params=params, project_id=project_id)
     prefix = res['filePath']
-    s3_session = boto3.Session(
-        aws_access_key_id=res['accessKeyId'],
-        aws_secret_access_key=res['secretAccessKey'],
-        aws_session_token=res['sessionToken']
-    )
+    s3_session = _get_boto_session_by_credentials(res)
     s3_resource = s3_session.resource('s3')
     bucket = s3_resource.Bucket(res["bucket"])
     try:
@@ -329,9 +325,11 @@ def move_images(
     )
     if image_names is None:
         image_names = search_images((source_project, source_project_folder))
-    copy_images((source_project, source_project_folder), image_names,
-                (destination_project, destination_project_folder),
-                include_annotations, copy_annotation_status, copy_pin)
+    copy_images(
+        (source_project, source_project_folder), image_names,
+        (destination_project, destination_project_folder), include_annotations,
+        copy_annotation_status, copy_pin
+    )
     delete_images((source_project, source_project_folder), image_names)
     logger.info(
         "Moved images %s from project %s to project %s", image_names,

superannotate/db/projects.py

@@ -40,7 +40,7 @@
     get_project_metadata_with_users
 )
 from .users import get_team_contributor_metadata
-from .utils import _get_upload_auth_token
+from .utils import _get_upload_auth_token, _get_boto_session_by_credentials
 
 logger = logging.getLogger("superannotate-python-sdk")
 
@@ -700,11 +700,7 @@ def __upload_images_to_aws_thread(
         from_s3 = from_session.resource('s3')
     if start_index >= len_img_paths:
         return
-    s3_session = boto3.Session(
-        aws_access_key_id=res['accessKeyId'],
-        aws_secret_access_key=res['secretAccessKey'],
-        aws_session_token=res['sessionToken']
-    )
+    s3_session = _get_boto_session_by_credentials(res)
     s3_resource = s3_session.resource('s3')
     bucket = s3_resource.Bucket(res["bucket"])
     prefix = res['filePath']
@@ -1095,11 +1091,7 @@ def __attach_image_urls_to_project_thread(
     end_index = start_index + chunksize
     if start_index >= len_img_paths:
         return
-    s3_session = boto3.Session(
-        aws_access_key_id=res['accessKeyId'],
-        aws_secret_access_key=res['secretAccessKey'],
-        aws_session_token=res['sessionToken']
-    )
+    s3_session = _get_boto_session_by_credentials(res)
     s3_resource = s3_session.resource('s3')
     bucket = s3_resource.Bucket(res["bucket"])
     prefix = res['filePath']
@@ -1520,11 +1512,7 @@ def __upload_annotations_thread(
             continue
         res = response.json()
         aws_creds = res["creds"]
-        s3_session = boto3.Session(
-            aws_access_key_id=aws_creds['accessKeyId'],
-            aws_secret_access_key=aws_creds['secretAccessKey'],
-            aws_session_token=aws_creds['sessionToken']
-        )
+        s3_session = _get_boto_session_by_credentials(aws_creds)
         s3_resource = s3_session.resource('s3')
         bucket = s3_resource.Bucket(aws_creds["bucket"])
         for image_id, image_info in res['images'].items():

superannotate/db/utils.py

@@ -1,9 +1,11 @@
+import boto3
 from ..api import API
 from ..exceptions import SABaseException, SAImageSizeTooLarge
+
 _api = API.get_instance()
 
 
-def _get_upload_auth_token(params,project_id):
+def _get_upload_auth_token(params, project_id):
     response = _api.send_request(
         req_type='GET',
         path=f'/project/{project_id}/sdkImageUploadToken',
@@ -13,7 +15,15 @@ def _get_upload_auth_token(params,project_id):
         raise SABaseException(
             response.status_code, "Couldn't get upload token " + response.text
         )
-    
+
     res = response.json()
     return res
 
+
+def _get_boto_session_by_credentials(credentials):
+    return boto3.Session(
+        aws_access_key_id=credentials['accessKeyId'],
+        aws_secret_access_key=credentials['secretAccessKey'],
+        aws_session_token=credentials['sessionToken'],
+        region_name=credentials['region']
+    )

superannotate/ml/ml_funcs.py

@@ -8,22 +8,20 @@
 import os
 import time
 
-import boto3
-import pandas as pd
-import plotly.express as px
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 
 from ..api import API
 from ..common import (
     _AVAILABLE_SEGMENTATION_MODELS, model_training_status_int_to_str,
-    project_type_str_to_int, upload_state_int_to_str,_MODEL_TRAINING_TASKS
+    project_type_str_to_int, upload_state_int_to_str, _MODEL_TRAINING_TASKS
 )
-from ..db.images import get_image_metadata, search_images
+from ..db.images import get_image_metadata
 from ..exceptions import SABaseException
 from ..parameter_decorators import model_metadata, project_metadata
 from .defaults import DEFAULT_HYPERPARAMETERS, NON_PLOTABLE_KEYS
 from .utils import log_process, make_plotly_specs, reformat_metrics_json
+from ..db.utils import _get_boto_session_by_credentials
 
 logger = logging.getLogger("superannotate-python-sdk")
 _api = API.get_instance()
@@ -476,11 +474,7 @@ def download_model(model, output_dir):
         raise SABaseException(0, "Could not get model info ")
 
     tokens = response["tokens"]
-    s3_session = boto3.Session(
-        aws_access_key_id=tokens["accessKeyId"],
-        aws_secret_access_key=tokens["secretAccessKey"],
-        aws_session_token=tokens["sessionToken"]
-    )
+    s3_session = _get_boto_session_by_credentials(tokens)
     s3_resource = s3_session.resource('s3')
 
     bucket = s3_resource.Bucket(tokens["bucket"])