fix auto annotation to not eat all RAM (#1328)

Co-authored-by: Nikita Manovich <40690625+nmanovic@users.noreply.github.com>

CHANGELOG.md

@@ -55,6 +55,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - React UI is the primary UI
 
 ### Fixed
+- Cleaned up memory in Auto Annotation to enable long running tasks on videos
 - New shape is added when press ``esc`` when drawing instead of cancellation
 - Dextr segmentation doesn't work.
 - `FileNotFoundError` during dump after moving format files

cvat/apps/auto_annotation/inference.py

@@ -1,6 +1,6 @@
+import itertools
 from .model_loader import ModelLoader
 from cvat.apps.engine.utils import import_modules, execute_python_code
-import itertools
 
 def _process_detections(detections, path_to_conv_script, restricted=True):
     results = Results()
@@ -31,6 +31,17 @@ def _process_detections(detections, path_to_conv_script, restricted=True):
 
     return results
 
+def _process_attributes(shape_attributes, label_attr_spec):
+    attributes = []
+    for attr_text, attr_value in shape_attributes.items():
+        if attr_text in label_attr_spec:
+            attributes.append({
+                "spec_id": label_attr_spec[attr_text],
+                "value": attr_value,
+            })
+
+    return attributes
+
 class Results():
     def __init__(self):
         self._results = {
@@ -84,36 +95,20 @@ class Results():
             "attributes": attributes or {},
         }
 
-def run_inference_engine_annotation(data, model_file, weights_file,
+class InferenceAnnotationRunner:
-       labels_mapping, attribute_spec, convertation_file, job=None, update_progress=None, restricted=True):
+    def __init__(self, data, model_file, weights_file, labels_mapping,
-    def process_attributes(shape_attributes, label_attr_spec):
+    attribute_spec, convertation_file):
-        attributes = []
+        self.data = iter(data)
-        for attr_text, attr_value in shape_attributes.items():
+        self.data_len = len(data)
-            if attr_text in label_attr_spec:
+        self.model = ModelLoader(model=model_file, weights=weights_file)
-                attributes.append({
+        self.frame_counter = 0
-                    "spec_id": label_attr_spec[attr_text],
+        self.attribute_spec = attribute_spec
-                    "value": attr_value,
+        self.convertation_file = convertation_file
-                })
+        self.iteration_size = 128
+        self.labels_mapping = labels_mapping
 
-        return attributes
-
-    def add_shapes(shapes, target_container):
-        for shape in shapes:
-            if shape["label"] not in labels_mapping:
-                    continue
-            db_label = labels_mapping[shape["label"]]
-            label_attr_spec = attribute_spec.get(db_label)
-            target_container.append({
-                "label_id": db_label,
-                "frame": shape["frame"],
-                "points": shape["points"],
-                "type": shape["type"],
-                "z_order": 0,
-                "group": None,
-                "occluded": False,
-                "attributes": process_attributes(shape["attributes"], label_attr_spec),
-            })
 
+    def run(self, job=None, update_progress=None, restricted=True):
         result = {
             "shapes": [],
             "tracks": [],
@@ -121,28 +116,48 @@ def run_inference_engine_annotation(data, model_file, weights_file,
             "version": 0
         }
 
-    data_len = len(data)
-    model = ModelLoader(model=model_file, weights=weights_file)
-
-    frame_counter = 0
-
         detections = []
-    for frame in data:
+        for _ in range(self.iteration_size):
+            try:
+                frame = next(self.data)
+            except StopIteration:
+                break
+
             orig_rows, orig_cols = frame.shape[:2]
 
             detections.append({
-            "frame_id": frame_counter,
+                "frame_id": self.frame_counter,
                 "frame_height": orig_rows,
                 "frame_width": orig_cols,
-            "detections": model.infer(frame),
+                "detections": self.model.infer(frame),
             })
 
-        frame_counter += 1
+            self.frame_counter += 1
-        if job and update_progress and not update_progress(job, frame_counter * 100 / data_len):
+            if job and update_progress and not update_progress(job, self.frame_counter * 100 / self.data_len):
-            return None
+                return None, False
+
+        processed_detections = _process_detections(detections, self.convertation_file, restricted=restricted)
+
+        self._add_shapes(processed_detections.get_shapes(), result["shapes"])
 
-    processed_detections = _process_detections(detections, convertation_file, restricted=restricted)
+        more_items = self.frame_counter != self.data_len
 
-    add_shapes(processed_detections.get_shapes(), result["shapes"])
+        return result, more_items
 
-    return result
+    def _add_shapes(self, shapes, target_container):
+        for shape in shapes:
+            if shape["label"] not in self.labels_mapping:
+                    continue
+
+            db_label = self.labels_mapping[shape["label"]]
+            label_attr_spec = self.attribute_spec.get(db_label)
+            target_container.append({
+                "label_id": db_label,
+                "frame": shape["frame"],
+                "points": shape["points"],
+                "type": shape["type"],
+                "z_order": 0,
+                "group": None,
+                "occluded": False,
+                "attributes": _process_attributes(shape["attributes"], label_attr_spec),
+            })

cvat/apps/auto_annotation/inference_engine.py

@@ -10,7 +10,6 @@ import platform
 
 _IE_PLUGINS_PATH = os.getenv("IE_PLUGINS_PATH", None)
 
-
 def _check_instruction(instruction):
     return instruction == str.strip(
         subprocess.check_output(
@@ -24,7 +23,7 @@ def make_plugin_or_core():
     use_core_openvino = False
     try:
         major, minor, reference = [int(x) for x in version.split('.')]
-        if major >= 2 and minor >= 1 and reference >= 37988:
+        if major >= 2 and minor >= 1:
             use_core_openvino = True
     except Exception:
         pass

cvat/apps/auto_annotation/model_manager.py

@@ -23,7 +23,7 @@ from cvat.apps.engine.frame_provider import FrameProvider
 from .models import AnnotationModel, FrameworkChoice
 from .model_loader import load_labelmap
 from .image_loader import ImageLoader
-from .inference import run_inference_engine_annotation
+from .inference import InferenceAnnotationRunner
 
 
 def _remove_old_file(model_file_field):
@@ -44,15 +44,15 @@ def _update_dl_model_thread(dl_model_id, name, is_shared, model_file, weights_fi
         test_image = np.ones((1024, 1980, 3), np.uint8) * 255
         try:
             dummy_labelmap = {key: key for key in load_labelmap(labelmap_file).keys()}
-            run_inference_engine_annotation(
+            runner = InferenceAnnotationRunner(
                 data=[test_image,],
                 model_file=model_file,
                 weights_file=weights_file,
                 labels_mapping=dummy_labelmap,
                 attribute_spec={},
-                convertation_file=interpretation_file,
+                convertation_file=interpretation_file)
-                restricted=restricted
+
-            )
+            runner.run(restricted=restricted)
         except Exception as e:
             return False, str(e)
 
@@ -227,17 +227,19 @@ def run_inference_thread(tid, model_file, weights_file, labels_mapping, attribut
 
         result = None
         slogger.glob.info("auto annotation with openvino toolkit for task {}".format(tid))
-        result = run_inference_engine_annotation(
+        more_data = True
+        runner = InferenceAnnotationRunner(
             data=ImageLoader(FrameProvider(db_task.data)),
             model_file=model_file,
             weights_file=weights_file,
             labels_mapping=labels_mapping,
             attribute_spec=attributes,
-            convertation_file= convertation_file,
+            convertation_file= convertation_file)
+        while more_data:
+            result, more_data = runner.run(
                 job=job,
                 update_progress=update_progress,
-            restricted=restricted
+                restricted=restricted)
-        )
 
             if result is None:
                 slogger.glob.info("auto annotation for task {} canceled by user".format(tid))