Auto segmentation using Mask_RCNN (cvat-ai#767)

.gitignore

@@ -12,7 +12,7 @@
 /ssh/*
 !/ssh/README.md
 node_modules
-
+/Mask_RCNN/
 
 # Ignore temporary files
 docker-compose.override.yml

Dockerfile

@@ -43,7 +43,10 @@ RUN apt-get update && \
         unzip \
         unrar \
         p7zip-full \
-        vim && \
+        vim \
+        git-core \
+        libsm6 \
+        libxext6 && \
     pip3 install -U setuptools && \
     ln -fs /usr/share/zoneinfo/${TZ} /etc/localtime && \
     dpkg-reconfigure -f noninteractive tzdata && \
@@ -78,6 +81,14 @@ RUN if [ "$TF_ANNOTATION" = "yes" ]; then \
         bash -i /tmp/components/tf_annotation/install.sh; \
     fi
 
+# Auto segmentation support. by Mohammad
+ARG AUTO_SEGMENTATION
+ENV AUTO_SEGMENTATION=${AUTO_SEGMENTATION}
+ENV AUTO_SEGMENTATION_PATH=${HOME}/Mask_RCNN
+RUN if [ "$AUTO_SEGMENTATION" = "yes" ]; then \
+    bash -i /tmp/components/auto_segmentation/install.sh; \
+    fi
+
 ARG WITH_TESTS
 RUN if [ "$WITH_TESTS" = "yes" ]; then \
         wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | apt-key add - && \

components/auto_segmentation/README.md

@@ -0,0 +1,38 @@
+## [Keras+Tensorflow Mask R-CNN Segmentation](https://github.com/matterport/Mask_RCNN)
+
+### What is it?
+-   This application allows you automatically to segment many various objects on images.
+-   It's based on Feature Pyramid Network (FPN) and a ResNet101 backbone.
+
+-   It uses a pre-trained model on MS COCO dataset
+-   It supports next classes (use them in "labels" row):
+```python
+'BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
+'bus', 'train', 'truck', 'boat', 'traffic light',
+'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
+'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
+'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
+'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
+'kite', 'baseball bat', 'baseball glove', 'skateboard',
+'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
+'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
+'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
+'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
+'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
+'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
+'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
+'teddy bear', 'hair drier', 'toothbrush'.
+```
+-  Component adds "Run Auto Segmentation" button into dashboard.
+
+### Build docker image
+```bash
+# From project root directory
+docker-compose -f docker-compose.yml -f components/auto_segmentation/docker-compose.auto_segmentation.yml build
+```
+
+### Run docker container
+```bash
+# From project root directory
+docker-compose -f docker-compose.yml -f components/auto_segmentation/docker-compose.auto_segmentation.yml up -d
+```

components/auto_segmentation/docker-compose.auto_segmentation.yml

@@ -0,0 +1,13 @@
+#
+# Copyright (C) 2018 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+#
+version: "2.3"
+
+services:
+  cvat:
+    build:
+      context: .
+      args:
+        AUTO_SEGMENTATION: "yes"

components/auto_segmentation/install.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+#
+
+set -e
+
+cd ${HOME} && \
+git clone https://github.com/matterport/Mask_RCNN.git && \
+wget https://github.com/matterport/Mask_RCNN/releases/download/v2.0/mask_rcnn_coco.h5 && \
+mv mask_rcnn_coco.h5 Mask_RCNN/mask_rcnn_coco.h5
+
+# TODO remove useless files
+# tensorflow and Keras are installed globally

cvat/apps/auto_segmentation/__init__.py

@@ -0,0 +1,9 @@
+
+# Copyright (C) 2018 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+from cvat.settings.base import JS_3RDPARTY
+
+JS_3RDPARTY['dashboard'] = JS_3RDPARTY.get('dashboard', []) + ['auto_segmentation/js/dashboardPlugin.js']
+

cvat/apps/auto_segmentation/admin.py

@@ -0,0 +1,8 @@
+
+# Copyright (C) 2018 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+
+# Register your models here.
+

cvat/apps/auto_segmentation/apps.py

@@ -0,0 +1,11 @@
+
+# Copyright (C) 2018 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+from django.apps import AppConfig
+
+
+class AutoSegmentationConfig(AppConfig):
+    name = 'auto_segmentation'
+

cvat/apps/auto_segmentation/migrations/__init__.py

@@ -0,0 +1,5 @@
+
+# Copyright (C) 2018 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+

cvat/apps/auto_segmentation/models.py

@@ -0,0 +1,8 @@
+
+# Copyright (C) 2018 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+
+# Create your models here.
+

cvat/apps/auto_segmentation/static/auto_segmentation/js/dashboardPlugin.js

@@ -0,0 +1,112 @@
+/*
+ * Copyright (C) 2018 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ */
+
+/* global
+    userConfirm:false
+    showMessage:false
+*/
+
+window.addEventListener('dashboardReady', () => {
+    function checkProcess(tid, button) {
+        function checkCallback() {
+            $.get(`/tensorflow/segmentation/check/task/${tid}`).done((statusData) => {
+                if (['started', 'queued'].includes(statusData.status)) {
+                    const progress = Math.round(statusData.progress) || '0';
+                    button.text(`Cancel Auto Segmentation (${progress}%)`);
+                    setTimeout(checkCallback, 5000);
+                } else {
+                    button.text('Run Auto Segmentation');
+                    button.removeClass('tfAnnotationProcess');
+                    button.prop('disabled', false);
+
+                    if (statusData.status === 'failed') {
+                        const message = `Tensorflow Segmentation failed. Error: ${statusData.stderr}`;
+                        showMessage(message);
+                    } else if (statusData.status !== 'finished') {
+                        const message = `Tensorflow segmentation check request returned status "${statusData.status}"`;
+                        showMessage(message);
+                    }
+                }
+            }).fail((errorData) => {
+                const message = `Can not sent tensorflow segmentation check request. Code: ${errorData.status}. `
+                    + `Message: ${errorData.responseText || errorData.statusText}`;
+                showMessage(message);
+            });
+        }
+
+        setTimeout(checkCallback, 5000);
+    }
+
+
+    function runProcess(tid, button) {
+        $.get(`/tensorflow/segmentation/create/task/${tid}`).done(() => {
+            showMessage('Process has started');
+            button.text('Cancel Auto Segmentation (0%)');
+            button.addClass('tfAnnotationProcess');
+            checkProcess(tid, button);
+        }).fail((errorData) => {
+            const message = `Can not run Auto Segmentation. Code: ${errorData.status}. `
+                + `Message: ${errorData.responseText || errorData.statusText}`;
+            showMessage(message);
+        });
+    }
+
+
+    function cancelProcess(tid, button) {
+        $.get(`/tensorflow/segmentation/cancel/task/${tid}`).done(() => {
+            button.prop('disabled', true);
+        }).fail((errorData) => {
+            const message = `Can not cancel Auto Segmentation. Code: ${errorData.status}. `
+                + `Message: ${errorData.responseText || errorData.statusText}`;
+            showMessage(message);
+        });
+    }
+
+
+    function setupDashboardItem(item, metaData) {
+        const tid = +item.attr('tid');
+        const button = $('<button> Run Auto Segmentation </button>');
+
+        button.on('click', () => {
+            if (button.hasClass('tfAnnotationProcess')) {
+                userConfirm('The process will be canceled. Continue?', () => {
+                    cancelProcess(tid, button);
+                });
+            } else {
+                userConfirm('The current annotation will be lost. Are you sure?', () => {
+                    runProcess(tid, button);
+                });
+            }
+        });
+
+        button.addClass('dashboardTFAnnotationButton regular dashboardButtonUI');
+        button.appendTo(item.find('div.dashboardButtonsUI'));
+
+        if ((tid in metaData) && (metaData[tid].active)) {
+            button.text('Cancel Auto Segmentation');
+            button.addClass('tfAnnotationProcess');
+            checkProcess(tid, button);
+        }
+    }
+
+    const elements = $('.dashboardItem');
+    const tids = Array.from(elements, el => +el.getAttribute('tid'));
+
+    $.ajax({
+        type: 'POST',
+        url: '/tensorflow/segmentation/meta/get',
+        data: JSON.stringify(tids),
+        contentType: 'application/json; charset=utf-8',
+    }).done((metaData) => {
+        elements.each(function setupDashboardItemWrapper() {
+            setupDashboardItem($(this), metaData);
+        });
+    }).fail((errorData) => {
+        const message = `Can not get Auto Segmentation meta info. Code: ${errorData.status}. `
+            + `Message: ${errorData.responseText || errorData.statusText}`;
+        showMessage(message);
+    });
+});

cvat/apps/auto_segmentation/tests.py

@@ -0,0 +1,8 @@
+
+# Copyright (C) 2018 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+
+# Create your tests here.
+

cvat/apps/auto_segmentation/urls.py

@@ -0,0 +1,14 @@
+
+# Copyright (C) 2018 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+from django.urls import path
+from . import views
+
+urlpatterns = [
+    path('create/task/<int:tid>', views.create),
+    path('check/task/<int:tid>', views.check),
+    path('cancel/task/<int:tid>', views.cancel),
+    path('meta/get', views.get_meta_info),
+]