SDK | Added createdAt, updatedAt for classes in classes.json

Author: Vaghinak Basentsyan

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

@@ -44,8 +44,9 @@
 from lib.core.enums import ImageQuality
 from lib.core.exceptions import AppException
 from lib.core.exceptions import AppValidationException
-from lib.core.types import ClassesJson
 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
@@ -1884,7 +1885,9 @@ def create_annotation_class(
     """
     if isinstance(project, Project):
         project = project.dict()
-    attribute_groups = list(map(lambda x: x.dict(), attribute_groups)) if attribute_groups else None
+    attribute_groups = (
+        list(map(lambda x: x.dict(), attribute_groups)) if attribute_groups else None
+    )
     response = controller.create_annotation_class(
         project_name=project, name=name, color=color, attribute_groups=attribute_groups
     )
@@ -1941,6 +1944,8 @@ def download_annotation_classes_json(project: str, folder: Union[str, Path]):
     response = controller.download_annotation_classes(
         project_name=project, download_path=folder
     )
+    if response.errors:
+        raise AppException(response.errors)
     return response.data
 
 
@@ -2268,7 +2273,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
 
 
@@ -2733,7 +2737,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
 
@@ -2744,8 +2748,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/entities.py

@@ -24,6 +24,22 @@ def to_dict(self):
         raise NotImplementedError
 
 
+class BaseTimedEntity(BaseEntity):
+    def __init__(
+        self, uuid: Any = None, createdAt: str = None, updatedAt: str = None,
+    ):
+        super().__init__(uuid)
+        self.createdAt = createdAt
+        self.updatedAt = updatedAt
+
+    def to_dict(self):
+        return {
+            "id": self.uuid,
+            "createdAt": self.createdAt,
+            "updatedAt": self.updatedAt,
+        }
+
+
 class ConfigEntity(BaseEntity):
     def __init__(self, uuid: str, value: str):
         super().__init__(uuid)
@@ -41,10 +57,12 @@ def to_dict(self):
         return {"key": self.uuid, "value": self.value}
 
 
-class ProjectEntity(BaseEntity):
+class ProjectEntity(BaseTimedEntity):
     def __init__(
         self,
         uuid: int = None,
+        createdAt: str = None,
+        updatedAt: str = None,
         team_id: int = None,
         name: str = None,
         project_type: int = None,
@@ -65,7 +83,7 @@ def __init__(
         completed_images_count: int = None,
         root_folder_completed_images_count: int = None,
     ):
-        super().__init__(uuid)
+        super().__init__(uuid, createdAt, updatedAt)
         self.team_id = team_id
         self.name = name
         self.project_type = project_type
@@ -100,7 +118,7 @@ def __copy__(self):
 
     def to_dict(self):
         return {
-            "id": self.uuid,
+            **super().to_dict(),
             "team_id": self.team_id,
             "name": self.name,
             "type": self.project_type,
@@ -175,17 +193,19 @@ def to_dict(self):
         }
 
 
-class FolderEntity(BaseEntity):
+class FolderEntity(BaseTimedEntity):
     def __init__(
         self,
         uuid: int = None,
+        createdAt: str = None,
+        updatedAt: str = None,
         project_id: int = None,
         parent_id: int = None,
         team_id: int = None,
         name: str = None,
         folder_users: List[dict] = None,
     ):
-        super().__init__(uuid)
+        super().__init__(uuid, createdAt, updatedAt)
         self.team_id = team_id
         self.project_id = project_id
         self.name = name
@@ -194,6 +214,7 @@ def __init__(
 
     def to_dict(self):
         return {
+            **super().to_dict(),
             "id": self.uuid,
             "team_id": self.team_id,
             "name": self.name,
@@ -301,22 +322,26 @@ def to_dict(self):
         return {"uuid": self.uuid, "bytes": self.data, "metadata": self.metadata}
 
 
-class AnnotationClassEntity(BaseEntity):
+class AnnotationClassEntity(BaseTimedEntity):
     def __init__(
         self,
         uuid: int = None,
+        createdAt: str = None,
+        updatedAt: str = None,
         color: str = None,
         count: int = None,
         name: str = None,
         project_id: int = None,
         attribute_groups: Iterable = None,
     ):
-        super().__init__(uuid)
+        super().__init__(uuid, createdAt, updatedAt)
         self.color = color
         self.count = count
         self.name = name
         self.project_id = project_id
         self.attribute_groups = attribute_groups
+        self.createdAt = createdAt
+        self.updatedAt = updatedAt
 
     def __copy__(self):
         return AnnotationClassEntity(
@@ -328,7 +353,7 @@ def __copy__(self):
 
     def to_dict(self):
         return {
-            "id": self.uuid,
+            **super().to_dict(),
             "color": self.color,
             "count": self.count,
             "name": self.name,
@@ -402,12 +427,14 @@ def to_dict(self):
         }
 
 
-class MLModelEntity(BaseEntity):
+class MLModelEntity(BaseTimedEntity):
     def __init__(
         self,
         uuid: int = None,
         team_id: int = None,
         name: str = None,
+        createdAt: str = None,
+        updatedAt: str = None,
         path: str = None,
         config_path: str = None,
         model_type: int = None,
@@ -423,7 +450,7 @@ def __init__(
         is_global: bool = None,
         hyper_parameters: dict = {},
     ):
-        super().__init__(uuid=uuid)
+        super().__init__(uuid, createdAt, updatedAt)
         self.name = name
         self.path = path
         self.team_id = team_id
@@ -443,7 +470,7 @@ def __init__(
 
     def to_dict(self):
         return {
-            "id": self.uuid,
+            **super().to_dict(),
             "name": self.name,
             "team_id": self.team_id,
             "description": self.description,

src/superannotate/lib/core/service_types.py

@@ -1,5 +1,6 @@
+from typing import Any
 from typing import Dict
-from typing import Optional
+from typing import List
 from typing import Union
 
 from lib.core.exceptions import AppException
@@ -52,11 +53,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]]
+    data: Any
 
     def __init__(self, response, content_type):
         data = {
@@ -71,3 +96,7 @@ def __init__(self, response, content_type):
     @property
     def ok(self):
         return 199 < self.status < 300
+
+    @property
+    def error(self):
+        return getattr(self.data, "error", "Unknown error.")

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

@@ -4,8 +4,8 @@
 
 from pydantic import BaseModel
 from pydantic import constr
-from pydantic import StrictStr
 from pydantic import Extra
+from pydantic import StrictStr
 
 
 NotEmptyStr = constr(strict=True, min_length=1)
@@ -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