Merge pull request #211 from superannotateai/ml_model_prediction

Ml model prediction

Author: dshabin

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

@@ -270,6 +270,9 @@ def clone_project(
     )
     if response.errors:
         raise AppException(response.errors)
+    logger.info(
+        f"Created project {project_name} (ID {response.data.uuid} ) with type { constances.ProjectType.get_name(response.data.project_type)}."
+    )
     return ProjectSerializer(response.data).serialize()
 
 

src/superannotate/lib/core/usecases.py

@@ -22,7 +22,6 @@
 import numpy as np
 import pandas as pd
 import requests
-from boto3.exceptions import Boto3Error
 from botocore.exceptions import ClientError
 from lib.app.analytics.common import aggregate_annotations_as_df
 from lib.app.analytics.common import consensus_plot
@@ -3437,7 +3436,9 @@ def execute(self):
                 )
                 if self._project.project_type == constances.ProjectType.PIXEL.value:
                     mask_path = None
-                    png_path = self._annotation_path.replace("___pixel.json", "___save.png")
+                    png_path = self._annotation_path.replace(
+                        "___pixel.json", "___save.png"
+                    )
                     if os.path.exists(png_path) and not self._mask:
                         mask_path = png_path
                     elif self._mask:
@@ -4122,36 +4123,44 @@ def __init__(
         self._backend_service = backend_service_provider
         self._team_id = team_id
 
+    def validate_training_status(self):
+        if self._model.training_status not in [
+            constances.TrainingStatus.COMPLETED.value,
+            constances.TrainingStatus.FAILED_AFTER_EVALUATION_WITH_SAVE_MODEL.value,
+        ]:
+            raise AppException("Unable to download.")
+
     def execute(self):
-        metrics_name = os.path.basename(self._model.path).replace(".pth", ".json")
-        mapper_path = self._model.config_path.replace(
-            os.path.basename(self._model.config_path), "classes_mapper.json"
-        )
-        metrics_path = self._model.config_path.replace(
-            os.path.basename(self._model.config_path), metrics_name
-        )
+        if self.is_valid():
+            metrics_name = os.path.basename(self._model.path).replace(".pth", ".json")
+            mapper_path = self._model.config_path.replace(
+                os.path.basename(self._model.config_path), "classes_mapper.json"
+            )
+            metrics_path = self._model.config_path.replace(
+                os.path.basename(self._model.config_path), metrics_name
+            )
 
-        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=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(auth_response.data.bucket)
+            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=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(auth_response.data.bucket)
 
-        bucket.download_file(
-            self._model.config_path, os.path.join(self._download_path, "config.yaml")
-        )
-        bucket.download_file(
-            self._model.path,
-            os.path.join(self._download_path, os.path.basename(self._model.path)),
-        )
-        if self._model.is_global:
+            bucket.download_file(
+                self._model.config_path,
+                os.path.join(self._download_path, "config.yaml"),
+            )
+            bucket.download_file(
+                self._model.path,
+                os.path.join(self._download_path, os.path.basename(self._model.path)),
+            )
             try:
                 bucket.download_file(
                     metrics_path, os.path.join(self._download_path, metrics_name)
@@ -4160,11 +4169,11 @@ def execute(self):
                     mapper_path,
                     os.path.join(self._download_path, "classes_mapper.json"),
                 )
-            except Boto3Error:
-                self._response.errors = AppException(
+            except ClientError:
+                logger.info(
                     "The specified model does not contain a classes_mapper and/or a metrics file."
                 )
-        self._response.data = self._model
+            self._response.data = self._model
         return self._response
 
 
@@ -4531,13 +4540,13 @@ def execute(self):
                 success_images = [
                     img.name
                     for img in images_metadata
-                    if img.segmentation_status
+                    if img.prediction_status
                     == constances.SegmentationStatus.COMPLETED.value
                 ]
                 failed_images = [
                     img.name
                     for img in images_metadata
-                    if img.segmentation_status
+                    if img.prediction_status
                     == constances.SegmentationStatus.FAILED.value
                 ]
 

src/superannotate/lib/infrastructure/controller.py

@@ -1402,6 +1402,8 @@ def download_ml_model(self, model_data: dict, download_path: str):
             path=model_data["path"],
             config_path=model_data["config_path"],
             team_id=model_data["team_id"],
+            training_status=model_data["training_status"],
+            is_global=model_data["is_global"],
         )
         use_case = usecases.DownloadMLModelUseCase(
             model=model,

src/superannotate/lib/infrastructure/repositories.py

@@ -468,4 +468,5 @@ def dict2entity(data: dict):
             path=data["path"],
             config_path=data["config_path"],
             is_global=data["is_global"],
+            training_status=data["training_status"],
         )

src/superannotate/version.py

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