Az/model manager bugfix (#353)

* fixed bug with AnnotationModel matching query does not exist during model create request * fixed deleting tmp files
cvat-ai · Mar 20, 2019 · ab013b0 · ab013b0
1 parent 1257a1b
commit ab013b0
Showing 1 changed file with 12 additions and 11 deletions.
diff --git a/cvat/apps/auto_annotation/model_manager.py b/cvat/apps/auto_annotation/model_manager.py
@@ -34,7 +34,7 @@ def _get_file_content(filename):
     def _delete_source_files():
         for f in [model_file, weights_file, labelmap_file, interpretation_file]:
             if f:
-                os.remove(model_file)
+                os.remove(f)
 
     def _run_test(model_file, weights_file, labelmap_file, interpretation_file):
         test_image = np.ones((1024, 1980, 3), np.uint8) * 255
@@ -105,13 +105,12 @@ def _run_test(model_file, weights_file, labelmap_file, interpretation_file):
             dl_model.updated_date = timezone.now()
             dl_model.save()
 
-    if not is_local_storage:
+    if is_local_storage:
         _delete_source_files()
 
     if not test_res:
         raise Exception("Model was not properly created/updated. Test failed: {}".format(message))
 
-@transaction.atomic
 def create_or_update(dl_model_id, name, model_file, weights_file, labelmap_file, interpretation_file, owner, storage, is_shared):
     def get_abs_path(share_path):
         if not share_path:
@@ -134,8 +133,9 @@ def save_file_as_tmp(data):
                 tmp_file.write(chunk)
         os.close(fd)
         return filename
-
-    dl_model = AnnotationModel.objects.get(pk=dl_model_id) if dl_model_id else create_empty(owner=owner)
+    is_create_request = dl_model_id is None
+    if is_create_request:
+        dl_model_id = create_empty(owner=owner)
 
     run_tests = bool(model_file or weights_file or labelmap_file or interpretation_file)
     if storage != "local":
@@ -149,12 +149,12 @@ def save_file_as_tmp(data):
         labelmap_file = save_file_as_tmp(labelmap_file)
         interpretation_file = save_file_as_tmp(interpretation_file)
 
-    rq_id = "auto_annotation.create.{}".format(dl_model.id)
+    rq_id = "auto_annotation.create.{}".format(dl_model_id)
     queue = django_rq.get_queue("default")
     queue.enqueue_call(
         func=_update_dl_model_thread,
         args=(
-            dl_model.id,
+            dl_model_id,
             name,
             is_shared,
             model_file,
@@ -163,13 +163,14 @@ def save_file_as_tmp(data):
             interpretation_file,
             run_tests,
             storage == "local",
-            not bool(dl_model_id),
+            is_create_request,
         ),
         job_id=rq_id
     )
 
     return rq_id
 
+@transaction.atomic
 def create_empty(owner, framework=FrameworkChoice.OPENVINO):
     db_model = AnnotationModel(
         owner=owner,
@@ -181,7 +182,7 @@ def create_empty(owner, framework=FrameworkChoice.OPENVINO):
         shutil.rmtree(model_path)
     os.mkdir(model_path)
 
-    return db_model
+    return db_model.id
 
 @transaction.atomic
 def delete(dl_model_id):
@@ -202,7 +203,7 @@ def get_image_key(item):
     image_list = []
     for root, _, filenames in os.walk(path_to_data):
         for filename in fnmatch.filter(filenames, "*.jpg"):
-                image_list.append(os.path.join(root, filename))
+            image_list.append(os.path.join(root, filename))
 
     image_list.sort(key=get_image_key)
     return ImageLoader(image_list)
@@ -428,4 +429,4 @@ def update_progress(job, progress):
             slogger.glob.exception("exception was occurred during auto annotation of the task {}: {}".format(tid, str(ex)), exc_info=True)
             raise ex
 
-        raise e
+        raise e