From ab013b0da174c53b82e8d49d549e4a4e02ee149c Mon Sep 17 00:00:00 2001
From: Andrey Zhavoronkov <41117609+azhavoro@users.noreply.github.com>
Date: Wed, 20 Mar 2019 22:45:09 +0300
Subject: [PATCH] Az/model manager bugfix (#353)

* fixed bug with AnnotationModel matching query does not exist during model create request

* fixed deleting tmp files
---
 cvat/apps/auto_annotation/model_manager.py | 23 +++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/cvat/apps/auto_annotation/model_manager.py b/cvat/apps/auto_annotation/model_manager.py
index da87317d..eef95628 100644
--- a/cvat/apps/auto_annotation/model_manager.py
+++ b/cvat/apps/auto_annotation/model_manager.py
@@ -34,7 +34,7 @@ def _update_dl_model_thread(dl_model_id, name, is_shared, model_file, weights_fi
     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 _update_dl_model_thread(dl_model_id, name, is_shared, model_file, weights_fi
             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 create_or_update(dl_model_id, name, model_file, weights_file, labelmap_file,
                 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 create_or_update(dl_model_id, name, model_file, weights_file, labelmap_file,
         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 create_or_update(dl_model_id, name, model_file, weights_file, labelmap_file,
             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_data(path_to_data):
     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 run_inference_thread(tid, model_file, weights_file, labels_mapping, attribut
             slogger.glob.exception("exception was occurred during auto annotation of the task {}: {}".format(tid, str(ex)), exc_info=True)
             raise ex
 
-        raise e
\ No newline at end of file
+        raise e