Merge remote-tracking branch 'origin/develop' into sdk_limitations

# Conflicts:
#	src/superannotate/lib/app/interface/sdk_interface.py
#	src/superannotate/lib/core/service_types.py

Author: Vaghinak Basentsyan

src/superannotate/lib/app/interface/sdk_interface.py

@@ -43,6 +43,7 @@
 from lib.core.exceptions import AppValidationException
 from lib.core.types import AttributeGroup
 from lib.core.types import ClassesJson
+from lib.core.types import MLModel
 from lib.core.types import Project
 from lib.infrastructure.controller import Controller
 from plotly.subplots import make_subplots
@@ -2269,7 +2270,6 @@ def download_image(
     )
     if response.errors:
         raise AppException(response.errors)
-    logger.info(f"Downloaded image {image_name} to {local_dir_path} ")
     return response.data
 
 
@@ -2734,7 +2734,7 @@ def stop_model_training(model: dict):
 
 @Trackable
 @validate_arguments
-def download_model(model: dict, output_dir: Union[str, Path]):
+def download_model(model: MLModel, output_dir: Union[str, Path]):
     """Downloads the neural network and related files
     which are the <model_name>.pth/pkl. <model_name>.json, <model_name>.yaml, classes_mapper.json
 
@@ -2745,8 +2745,9 @@ def download_model(model: dict, output_dir: Union[str, Path]):
     :return: the metadata of the model
     :rtype: dict
     """
-
-    res = controller.download_ml_model(model_data=model, download_path=output_dir)
+    res = controller.download_ml_model(
+        model_data=model.dict(), download_path=output_dir
+    )
     if res.errors:
         logger.error("\n".join([str(error) for error in res.errors]))
     else:

src/superannotate/lib/core/service_types.py

@@ -1,6 +1,8 @@
+from typing import Any
 from typing import Dict
 from typing import Optional
 from typing import Union
+from typing import List
 
 from pydantic import BaseModel
 from pydantic import Extra
@@ -45,11 +47,35 @@ def __init__(self, **data):
         super().__init__(**data)
 
 
+class DownloadMLModelAuthData(BaseModel):
+    access_key: str
+    secret_key: str
+    session_token: str
+    region: str
+    bucket: str
+    paths: List[str]
+
+    class Config:
+        extra = Extra.allow
+        fields = {
+            "access_key": "accessKeyId",
+            "secret_key": "secretAccessKey",
+            "session_token": "sessionToken",
+            "region": "region",
+        }
+
+    def __init__(self, **data):
+        credentials = data["tokens"]
+        data.update(credentials)
+        del data["tokens"]
+        super().__init__(**data)
+
+
 class ServiceResponse(BaseModel):
     status: int
     reason: str
     content: Union[bytes, str]
-    data: Optional[Union[UserLimits, UploadAnnotationAuthData, ErrorMessage]]
+    data: Any
 
     def __init__(self, response, content_type):
         data = {

src/superannotate/lib/core/serviceproviders.py

@@ -286,7 +286,9 @@ def delete_model(self, team_id: int, model_id: int):
     def stop_model_training(self, team_id: int, model_id: int):
         raise NotImplementedError
 
-    def get_ml_model_download_tokens(self, team_id: int, model_id: int):
+    def get_ml_model_download_tokens(
+        self, team_id: int, model_id: int
+    ) -> ServiceResponse:
         raise NotImplementedError
 
     def run_segmentation(

src/superannotate/lib/core/types.py

@@ -127,3 +127,14 @@ class Project(BaseModel):
 
     class Config:
         extra = Extra.allow
+
+
+class MLModel(BaseModel):
+    name: NotEmptyStr
+    id: int
+    path: NotEmptyStr
+    config_path: NotEmptyStr
+    team_id: Optional[int]
+
+    class Config:
+        extra = Extra.allow

src/superannotate/lib/core/usecases.py

@@ -2206,6 +2206,7 @@ def __init__(
         image_name: str,
         images: BaseManageableRepository,
         destination: str,
+        annotation_classes: BaseManageableRepository,
     ):
         super().__init__()
         self._service = service
@@ -2214,6 +2215,7 @@ def __init__(
         self._image_name = image_name
         self._images = images
         self._destination = destination
+        self._annotation_classes = annotation_classes
 
     @property
     def image_use_case(self):
@@ -2231,6 +2233,85 @@ def validate_project_type(self):
                 constances.LIMITED_FUNCTIONS[self._project.project_type]
             )
 
+    @property
+    def annotation_classes_name_map(self) -> dict:
+        classes_data = defaultdict(dict)
+        annotation_classes = self._annotation_classes.get_all()
+        for annotation_class in annotation_classes:
+            class_info = {"id": annotation_class.uuid}
+            if annotation_class.attribute_groups:
+                for attribute_group in annotation_class.attribute_groups:
+                    attribute_group_data = defaultdict(dict)
+                    for attribute in attribute_group["attributes"]:
+                        attribute_group_data[attribute["name"]] = attribute["id"]
+                    class_info["attribute_groups"] = {
+                        attribute_group["name"]: {
+                            "id": attribute_group["id"],
+                            "attributes": attribute_group_data,
+                        }
+                    }
+            classes_data[annotation_class.name] = class_info
+        return classes_data
+
+    def get_templates_mapping(self):
+        templates = self._service.get_templates(team_id=self._project.team_id).get(
+            "data", []
+        )
+        templates_map = {}
+        for template in templates:
+            templates_map[template["name"]] = template["id"]
+        return templates_map
+
+    def fill_classes_data(self, annotations: dict):
+        annotation_classes = self.annotation_classes_name_map
+        if "instances" not in annotations:
+            return
+        unknown_classes = {}
+        for annotation in [i for i in annotations["instances"] if "className" in i]:
+            if "className" not in annotation:
+                return
+            annotation_class_name = annotation["className"]
+            if annotation_class_name not in annotation_classes:
+                if annotation_class_name not in unknown_classes:
+                    unknown_classes[annotation_class_name] = {
+                        "id": -(len(unknown_classes) + 1),
+                        "attribute_groups": {},
+                    }
+        annotation_classes.update(unknown_classes)
+        templates = self.get_templates_mapping()
+        for annotation in (
+            i for i in annotations["instances"] if i.get("type", None) == "template"
+        ):
+            annotation["templateId"] = templates.get(
+                annotation.get("templateName", ""), -1
+            )
+
+        for annotation in [i for i in annotations["instances"] if "className" in i]:
+            annotation_class_name = annotation["className"]
+            if annotation_class_name not in annotation_classes:
+                continue
+            annotation["classId"] = annotation_classes[annotation_class_name]["id"]
+            for attribute in annotation["attributes"]:
+                if (
+                    attribute["groupName"]
+                    not in annotation_classes[annotation_class_name]["attribute_groups"]
+                ):
+                    continue
+                attribute["groupId"] = annotation_classes[annotation_class_name][
+                    "attribute_groups"
+                ][attribute["groupName"]]["id"]
+                if (
+                    attribute["name"]
+                    not in annotation_classes[annotation_class_name][
+                        "attribute_groups"
+                    ][attribute["groupName"]]["attributes"]
+                ):
+                    del attribute["groupId"]
+                    continue
+                attribute["id"] = annotation_classes[annotation_class_name][
+                    "attribute_groups"
+                ][attribute["groupName"]]["attributes"]
+
     def execute(self):
         if self.is_valid():
             data = {
@@ -2283,6 +2364,7 @@ def execute(self):
                     logger.info("There is no blue-map for the image.")
 
             json_path = Path(self._destination) / data["annotation_json_filename"]
+            self.fill_classes_data(data["annotation_json"])
             with open(json_path, "w") as f:
                 json.dump(data["annotation_json"], f, indent=4)
 
@@ -3039,6 +3121,7 @@ def __init__(
         images: BaseManageableRepository,
         classes: BaseManageableRepository,
         backend_service_provider: SuerannotateServiceProvider,
+        annotation_classes: BaseReadOnlyRepository,
         download_path: str,
         image_variant: str = "original",
         include_annotations: bool = False,
@@ -3065,6 +3148,7 @@ def __init__(
             image_name=self._image.name,
             images=images,
             destination=download_path,
+            annotation_classes=annotation_classes,
         )
         self.get_annotation_classes_ues_case = GetAnnotationClassesUseCase(
             classes=classes,
@@ -3098,31 +3182,35 @@ def validate_include_annotations(self):
 
     def execute(self):
         if self.is_valid():
+            fuse_image = None
+            annotations = None
+
             image_bytes = self.get_image_use_case.execute().data
             download_path = f"{self._download_path}/{self._image.name}"
             if self._image_variant == "lores":
                 download_path = download_path + "___lores.jpg"
             with open(download_path, "wb") as image_file:
                 image_file.write(image_bytes.getbuffer())
 
-            annotations = None
             if self._include_annotations:
                 annotations = self.download_annotation_use_case.execute().data
 
-            fuse_image = None
             if self._include_annotations and self._include_fuse:
                 classes = self.get_annotation_classes_ues_case.execute().data
-                fuse_image_use_case = CreateFuseImageUseCase(
-                    project_type=constances.ProjectType.get_name(
-                        self._project.project_type
-                    ),
-                    image_path=download_path,
-                    classes=[
-                        annotation_class.to_dict() for annotation_class in classes
-                    ],
-                    generate_overlay=self._include_overlay,
+                fuse_image = (
+                    CreateFuseImageUseCase(
+                        project_type=constances.ProjectType.get_name(
+                            self._project.project_type
+                        ),
+                        image_path=download_path,
+                        classes=[
+                            annotation_class.to_dict() for annotation_class in classes
+                        ],
+                        generate_overlay=self._include_overlay,
+                    )
+                    .execute()
+                    .data
                 )
-                fuse_image = fuse_image_use_case.execute().data
 
             self._response.data = (
                 download_path,
@@ -3912,16 +4000,18 @@ def execute(self):
             os.path.basename(self._model.config_path), metrics_name
         )
 
-        download_token = self._backend_service.get_ml_model_download_tokens(
+        auth_response = self._backend_service.get_ml_model_download_tokens(
             self._team_id, self._model.uuid
         )
+        if not auth_response.ok:
+            raise AppException(auth_response.error)
         s3_session = boto3.Session(
-            aws_access_key_id=download_token["tokens"]["accessKeyId"],
-            aws_secret_access_key=download_token["tokens"]["secretAccessKey"],
-            aws_session_token=download_token["tokens"]["sessionToken"],
-            region_name=download_token["tokens"]["region"],
+            aws_access_key_id=auth_response.data.access_key,
+            aws_secret_access_key=auth_response.data.secret_key,
+            aws_session_token=auth_response.data.session_token,
+            region_name=auth_response.data.region,
         )
-        bucket = s3_session.resource("s3").Bucket(download_token["tokens"]["bucket"])
+        bucket = s3_session.resource("s3").Bucket(auth_response.data.bucket)
 
         bucket.download_file(
             self._model.config_path, os.path.join(self._download_path, "config.yaml")

src/superannotate/lib/infrastructure/controller.py

@@ -994,6 +994,9 @@ def download_image_annotations(
             image_name=image_name,
             images=ImageRepository(service=self._backend_client),
             destination=destination,
+            annotation_classes=AnnotationClassRepository(
+                service=self._backend_client, project=project
+            ),
         )
         return use_case.execute()
 
@@ -1230,6 +1233,9 @@ def download_image(
             include_annotations=include_annotations,
             include_fuse=include_fuse,
             include_overlay=include_overlay,
+            annotation_classes=AnnotationClassRepository(
+                service=self._backend_client, project=project
+            ),
         )
         return use_case.execute()
 

src/superannotate/lib/infrastructure/services.py

@@ -10,6 +10,7 @@
 import lib.core as constance
 import requests.packages.urllib3
 from lib.core.exceptions import AppException
+from lib.core.service_types import DownloadMLModelAuthData
 from lib.core.service_types import ServiceResponse
 from lib.core.service_types import UploadAnnotationAuthData
 from lib.core.service_types import UserLimits
@@ -934,8 +935,12 @@ def get_ml_model_download_tokens(self, team_id: int, model_id: int):
         get_token_url = urljoin(
             self.api_url, self.URL_GET_ML_MODEL_DOWNLOAD_TOKEN.format(model_id)
         )
-        res = self._request(get_token_url, "get", params={"team_id": team_id})
-        return res.json()
+        return self._request(
+            get_token_url,
+            "get",
+            params={"team_id": team_id},
+            content_type=DownloadMLModelAuthData,
+        )
 
     def run_segmentation(
         self, team_id: int, project_id: int, model_name: str, image_ids: list

src/superannotate/version.py

@@ -1 +1 @@
-__version__ = "5.0.0b25"
+__version__ = "5.0.0b27"

tests/integration/test_interface.py

@@ -97,3 +97,21 @@ def test_search_project(self):
         sa.set_image_annotation_status(self.PROJECT_NAME, self.EXAMPLE_IMAGE_1, "Completed")
         data = sa.search_projects(self.PROJECT_NAME, return_metadata=True, include_complete_image_count=True)
         self.assertIsNotNone(data[0]['completed_images_count'])
+
+    def test_overlay_fuse(self):
+        sa.upload_image_to_project(self.PROJECT_NAME, f"{self.folder_path}/{self.EXAMPLE_IMAGE_1}")
+        sa.create_annotation_classes_from_classes_json(self.PROJECT_NAME, f"{self.folder_path}/classes/classes.json")
+        sa.upload_image_annotations(
+            self.PROJECT_NAME, self.EXAMPLE_IMAGE_1, f"{self.folder_path}/{self.EXAMPLE_IMAGE_1}___objects.json"
+        )
+        with tempfile.TemporaryDirectory() as temp_dir:
+            paths = sa.download_image(
+                self.PROJECT_NAME,
+                self.EXAMPLE_IMAGE_1,
+                temp_dir,
+                include_annotations=True,
+                include_fuse=True,
+                include_overlay=True,
+            )
+            self.assertIsNotNone(paths)
+