Model manager for DL models to preannotate images (#289)

cvat/apps/auto_annotation/README.md

@@ -2,8 +2,9 @@
 
 ### Description
 
-The application will be enabled automatically if OpenVINO™ component is
-installed. It allows to use custom models for auto annotation. Only models in
+The application will be enabled automatically if
+[OpenVINO™ component](../../../components/openvino)
+is installed. It allows to use custom models for auto annotation. Only models in
 OpenVINO™ toolkit format are supported. If you would like to annotate a
 task with a custom model please convert it to the intermediate representation
 (IR) format via the model optimizer tool. See [OpenVINO documentation](https://software.intel.com/en-us/articles/OpenVINO-InferEngine) for details.
@@ -14,9 +15,7 @@ To annotate a task with a custom model you need to prepare 4 files:
 1. __Model config__ (*.xml) - a text file with network configuration.
 1. __Model weights__ (*.bin) - a binary file with trained weights.
 1. __Label map__ (*.json) - a simple json file with `label_map` dictionary like
-object with string values for label numbers. Values in `label_map` should be
-exactly equal to labels for the annotation task, otherwise objects with mismatched
-labels will be ignored.
+object with string values for label numbers.
   Example:
     ```json
     {

cvat/apps/auto_annotation/__init__.py

@@ -3,6 +3,10 @@
 #
 # SPDX-License-Identifier: MIT
 
-from cvat.settings.base import JS_3RDPARTY
+from cvat.settings.base import JS_3RDPARTY, CSS_3RDPARTY
 
-JS_3RDPARTY['dashboard'] = JS_3RDPARTY.get('dashboard', []) + ['auto_annotation/js/auto_annotation.js']
+default_app_config = 'cvat.apps.auto_annotation.apps.AutoAnnotationConfig'
+
+JS_3RDPARTY['dashboard'] = JS_3RDPARTY.get('dashboard', []) + ['auto_annotation/js/dashboardPlugin.js']
+
+CSS_3RDPARTY['dashboard'] = CSS_3RDPARTY.get('dashboard', []) + ['auto_annotation/stylesheet.css']

cvat/apps/auto_annotation/apps.py

@@ -7,5 +7,9 @@
 
 
 class AutoAnnotationConfig(AppConfig):
-    name = "auto_annotation"
+    name = "cvat.apps.auto_annotation"
 
+    def ready(self):
+        from .permissions import setup_permissions
+
+        setup_permissions()

cvat/apps/auto_annotation/migrations/0001_initial.py

@@ -0,0 +1,39 @@
+# Generated by Django 2.1.3 on 2019-01-24 14:05
+
+import cvat.apps.auto_annotation.models
+from django.conf import settings
+import django.core.files.storage
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='AnnotationModel',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('name', cvat.apps.auto_annotation.models.SafeCharField(max_length=256)),
+                ('created_date', models.DateTimeField(auto_now_add=True)),
+                ('updated_date', models.DateTimeField(auto_now_add=True)),
+                ('model_file', models.FileField(storage=django.core.files.storage.FileSystemStorage(), upload_to=cvat.apps.auto_annotation.models.upload_path_handler)),
+                ('weights_file', models.FileField(storage=django.core.files.storage.FileSystemStorage(), upload_to=cvat.apps.auto_annotation.models.upload_path_handler)),
+                ('labelmap_file', models.FileField(storage=django.core.files.storage.FileSystemStorage(), upload_to=cvat.apps.auto_annotation.models.upload_path_handler)),
+                ('interpretation_file', models.FileField(storage=django.core.files.storage.FileSystemStorage(), upload_to=cvat.apps.auto_annotation.models.upload_path_handler)),
+                ('shared', models.BooleanField(default=False)),
+                ('primary', models.BooleanField(default=False)),
+                ('framework', models.CharField(default=cvat.apps.auto_annotation.models.FrameworkChoice('openvino'), max_length=32)),
+                ('owner', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, to=settings.AUTH_USER_MODEL)),
+            ],
+            options={
+                'default_permissions': (),
+            },
+        ),
+    ]

cvat/apps/auto_annotation/model_manager.py

@@ -0,0 +1,143 @@
+# Copyright (C) 2018 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+import django_rq
+import os
+import rq
+import shutil
+
+from django.db import transaction
+from django.utils import timezone
+from django.conf import settings
+
+from .models import AnnotationModel, FrameworkChoice
+
+def _remove_old_file(model_file_field):
+    if model_file_field and os.path.exists(model_file_field.name):
+        os.remove(model_file_field.name)
+
+@transaction.atomic
+def _update_dl_model_thread(dl_model_id, model_file, weights_file, labelmap_file, interpretation_file, run_tests):
+    def _get_file_content(filename):
+        return os.path.basename(filename), open(filename, "rb")
+
+    job = rq.get_current_job()
+    job.meta["progress"] = "Saving data"
+    job.save_meta()
+
+    dl_model = AnnotationModel.objects.select_for_update().get(pk=dl_model_id)
+
+    #save files in case of files should be uploaded from share
+    if model_file:
+        _remove_old_file(dl_model.model_file)
+        dl_model.model_file.save(*_get_file_content(model_file))
+    if weights_file:
+        _remove_old_file(dl_model.weights_file)
+        dl_model.weights_file.save(*_get_file_content(weights_file))
+    if labelmap_file:
+        _remove_old_file(dl_model.labelmap_file)
+        dl_model.labelmap_file.save(*_get_file_content(labelmap_file))
+    if interpretation_file:
+        _remove_old_file(dl_model.interpretation_file)
+        dl_model.interpretation_file.save(*_get_file_content(interpretation_file))
+
+    if run_tests:
+        #only for testing
+        import time
+        time.sleep(3)
+        job.meta["progress"] = "Test started"
+        job.save_meta()
+        time.sleep(5)
+        job.meta["progress"] = "Test finished"
+
+@transaction.atomic
+def update_model(dl_model_id, name, model_file, weights_file, labelmap_file, interpretation_file, storage, is_shared):
+
+    def get_abs_path(share_path):
+        if not share_path:
+            return share_path
+        share_root = settings.SHARE_ROOT
+        relpath = os.path.normpath(share_path).lstrip('/')
+        if '..' in relpath.split(os.path.sep):
+            raise Exception('Permission denied')
+        abspath = os.path.abspath(os.path.join(share_root, relpath))
+        if os.path.commonprefix([share_root, abspath]) != share_root:
+            raise Exception('Bad file path on share: ' + abspath)
+        return abspath
+
+    dl_model = AnnotationModel.objects.select_for_update().get(pk=dl_model_id)
+
+    if name:
+        dl_model.name = name
+
+    if is_shared != None:
+        dl_model.shared = is_shared
+
+    run_tests = bool(model_file or weights_file or labelmap_file or interpretation_file)
+    if storage != "local":
+        model_file = get_abs_path(model_file)
+        weights_file = get_abs_path(weights_file)
+        labelmap_file = get_abs_path(labelmap_file)
+        interpretation_file = get_abs_path(interpretation_file)
+    else:
+        if model_file:
+            _remove_old_file(dl_model.model_file)
+            dl_model.model_file = model_file
+            model_file = None
+        if weights_file:
+            _remove_old_file(dl_model.weights_file)
+            dl_model.weights_file = weights_file
+            weights_file = None
+        if labelmap_file:
+            _remove_old_file(dl_model.labelmap_file)
+            dl_model.labelmap_file = labelmap_file
+            labelmap_file = None
+        if interpretation_file:
+            _remove_old_file(dl_model.interpretation_file)
+            dl_model.interpretation_file = interpretation_file
+            interpretation_file = None
+
+    dl_model.updated_date = timezone.now()
+    dl_model.save()
+
+    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,
+            model_file,
+            weights_file,
+            labelmap_file,
+            interpretation_file,
+            run_tests,
+        ),
+        job_id = rq_id
+    )
+
+    return rq_id
+
+def create_empty(owner, framework=FrameworkChoice.OPENVINO):
+    db_model = AnnotationModel(
+        owner=owner,
+    )
+    db_model.save()
+
+    model_path = db_model.get_dirname()
+    if os.path.isdir(model_path):
+        shutil.rmtree(model_path)
+    os.mkdir(model_path)
+
+    return db_model.id
+
+@transaction.atomic
+def delete(dl_model_id):
+    dl_model = AnnotationModel.objects.select_for_update().get(pk=dl_model_id)
+    if dl_model:
+        if dl_model.primary:
+            raise Exception("Can not delete primary model {}".format(dl_model_id))
+
+        dl_model.delete()
+        shutil.rmtree(dl_model.get_dirname(), ignore_errors=True)
+    else:
+        raise Exception("Requested DL model {} doesn't exist".format(dl_model_id))

cvat/apps/auto_annotation/models.py

@@ -2,3 +2,54 @@
 # Copyright (C) 2018 Intel Corporation
 #
 # SPDX-License-Identifier: MIT
+
+from enum import Enum
+
+from django.db import models
+from django.conf import settings
+from django.contrib.auth.models import User
+from django.core.files.storage import FileSystemStorage
+
+fs = FileSystemStorage()
+
+def upload_path_handler(instance, filename):
+    return "{models_root}/{id}/{file}".format(models_root=settings.MODELS_ROOT, id=instance.id, file=filename)
+
+class FrameworkChoice(Enum):
+    OPENVINO = 'openvino'
+    TENSORFLOW = 'tensorflow'
+    PYTORCH = 'pytorch'
+
+    def __str__(self):
+        return self.value
+
+
+class SafeCharField(models.CharField):
+    def get_prep_value(self, value):
+        value = super().get_prep_value(value)
+        if value:
+            return value[:self.max_length]
+        return value
+
+class AnnotationModel(models.Model):
+    name = SafeCharField(max_length=256)
+    owner = models.ForeignKey(User, null=True, blank=True,
+        on_delete=models.SET_NULL)
+    created_date = models.DateTimeField(auto_now_add=True)
+    updated_date = models.DateTimeField(auto_now_add=True)
+    model_file = models.FileField(upload_to=upload_path_handler, storage=fs)
+    weights_file = models.FileField(upload_to=upload_path_handler, storage=fs)
+    labelmap_file = models.FileField(upload_to=upload_path_handler, storage=fs)
+    interpretation_file = models.FileField(upload_to=upload_path_handler, storage=fs)
+    shared = models.BooleanField(default=False)
+    primary = models.BooleanField(default=False)
+    framework = models.CharField(max_length=32, default=FrameworkChoice.OPENVINO)
+
+    class Meta:
+        default_permissions = ()
+
+    def get_dirname(self):
+        return "{models_root}/{id}".format(models_root=settings.MODELS_ROOT, id=self.id)
+
+    def __str__(self):
+        return self.name

cvat/apps/auto_annotation/permissions.py

@@ -0,0 +1,29 @@
+# Copyright (C) 2018 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+import rules
+
+from cvat.apps.authentication.auth import has_admin_role, has_user_role
+
+@rules.predicate
+def is_model_owner(db_user, db_dl_model):
+    return db_dl_model.owner == db_user
+
+@rules.predicate
+def is_shared_model(_, db_dl_model):
+    return db_dl_model.shared
+
+@rules.predicate
+def is_primary_model(_, db_dl_model):
+    return db_dl_model.primary
+
+def setup_permissions():
+    rules.add_perm('auto_annotation.model.create', has_admin_role | has_user_role)
+
+    rules.add_perm('auto_annotation.model.update', (has_admin_role | is_model_owner) & ~is_primary_model)
+
+    rules.add_perm('auto_annotation.model.delete', (has_admin_role | is_model_owner) & ~is_primary_model)
+
+    rules.add_perm('auto_annotation.model.access', has_admin_role | is_model_owner |
+        is_shared_model | is_primary_model)