Merge pull request #19 from superannotateai/folder-limitations

Fix training

Author: dshabin

superannotate/common.py

@@ -51,6 +51,14 @@
     "Failed": 4
 }
 
+_MODEL_TRAINING_TASKS = {
+    "Instance Segmentation for Pixel Projects" : "instance_segmentation_pixel",
+    "Instance Segmentation for Vector Projects" : "instance_segmentation_vector",
+    "Keypoint Detection for Vector Projects" : "keypoint_detection_vector",
+    "Object Detection for Vector Projects" : "object_detection_vector",
+    "Semantic Segmentation for Pixel Projects" : "semantic_segmentation_pixel"
+}
+
 
 def prediction_segmentation_status_from_str_to_int(status):
     return _PREDICTION_SEGMENTATION_STATUSES[status]

superannotate/db/search_projects.py

@@ -23,9 +23,11 @@ def search_projects(
     :rtype: list of strs or dicts
     """
     result_list = []
+    limit = 1000
     params = {
         'team_id': str(_api.team_id),
         'offset': 0,
+        'limit' : limit,
         'completeImagesCount': include_complete_image_count
     }
     if name is not None:
@@ -37,9 +39,9 @@ def search_projects(
         if response.ok:
             new_results = response.json()
             result_list += new_results["data"]
-            if new_results["count"] <= len(result_list):
+            params["offset"] += limit
+            if params["offset"] >= new_results["count"]:
                 break
-            params["offset"] = len(result_list)
         else:
             raise SABaseException(
                 response.status_code,

superannotate/ml/ml_funcs.py

@@ -17,7 +17,7 @@
 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
+    project_type_str_to_int, upload_state_int_to_str,_MODEL_TRAINING_TASKS
 )
 from ..db.images import get_image_metadata, search_images
 from ..exceptions import SABaseException
@@ -253,11 +253,11 @@ def run_training(
             hyperparameters[item] = DEFAULT_HYPERPARAMETERS[item]
     complete_image_count = 0
     for proj in project:
-        complete_image_count += proj['completedImagesCount']
+        complete_image_count += proj['rootFolderCompletedImagesCount']
 
     hyperparameters["name"] = model_name
     hyperparameters["description"] = model_description
-    hyperparameters["task"] = task
+    hyperparameters["task"] = _MODEL_TRAINING_TASKS[task]
     hyperparameters["base_model_id"] = base_model["id"]
     hyperparameters["project_ids"] = project_ids
     hyperparameters["image_count"] = complete_image_count