removing the Status Enums

Author: TumanyanV

superannotate/common.py

@@ -33,13 +33,21 @@
     "FailedAfterEvaluationWithSavedModel": 6
 }
 
+_PREDICTION_SEGMENTATION_STATUSES = {
+    "NotStarted": 1,
+    "InProgress": 2,
+    "Completed": 3,
+    "Failed": 4
+}
 
-class PredictionSegmentationStatuses(IntEnum):
-    NotStarted = 1
-    InProgress = 2
-    Completed = 3
-    Failed = 4
+def prediction_segmentation_status_from_str_to_int(status):
+    return _PREDICTION_SEGMENTATION_STATUSES[status]
 
+def prediction_segmentation_status_from_int_to_str(status):
+    for idx, item in _PREDICTION_SEGMENTATION_STATUSES.items():
+        if item == status:
+            return idx
+    raise RuntimeError("NA segmentation/prediction status")
 
 def model_training_status_int_to_str(project_status):
     for item in _MODEL_TRAINING_STATUSES:

superannotate/db/images.py

@@ -160,6 +160,11 @@ def get_image_metadata(project, image_names):
     )
 
     metadata = response.json()
+    for item in metadata:
+        item['annotation_status'] = common.annotation_status_int_to_str(item['annotation_status'])
+        item['prediction_status'] = common.prediction_segmentation_status_from_int_to_str(item['prediction_status'])
+        item['segmentation_status'] = common.prediction_segmentation_status_from_int_to_str(item['segmentation_status'])
+
     if len(metadata) == 1:
         return metadata[0]
     return metadata

superannotate/ml/utils.py

@@ -1,5 +1,5 @@
 from ..db.images import get_image_metadata
-from ..common import PredictionSegmentationStatuses
+from ..common import _PREDICTION_SEGMENTATION_STATUSES
 from .defaults import PLOTTABLE_METRICS
 from ast import literal_eval
 import pandas as pd
@@ -60,13 +60,14 @@ def get_images_prediction_segmentation_status(project, image_names, task):
     metadata = get_image_metadata(project, image_names)
     if isinstance(metadata, dict):
         metadata = [metadata]
+
     success_names = [
         x['name']
-        for x in metadata if x[task] == PredictionSegmentationStatuses.Completed
+        for x in metadata if x[task] == 'Completed'
     ]
     failure_names = [
         x['name']
-        for x in metadata if x[task] == PredictionSegmentationStatuses.Failed
+        for x in metadata if x[task] == 'Failed'
     ]
     return success_names, failure_names