Semi Automatic Segmentation

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - The ReID application for automatic bounding box merging has been added (#299)
 - Keyboard shortcuts to switch next/previous default shape type (box, polygon etc) [Alt + <, Alt + >] (#316)
 - Converter for VOC now supports interpolation tracks 
+- Semi-automatic semantic segmentation with the [Deep Extreme Cut](http://www.vision.ee.ethz.ch/~cvlsegmentation/dextr/) work
 
 ### Changed
 - Propagation setup has been moved from settings to bottom player panel

Dockerfile

@@ -125,6 +125,16 @@ RUN if [ "$OPENVINO_TOOLKIT" = "yes" ]; then \
         wget https://download.01.org/openvinotoolkit/2018_R5/open_model_zoo/person-reidentification-retail-0079/FP32/person-reidentification-retail-0079.bin -O reid/reid.bin; \
     fi
 
+# TODO: CHANGE URL
+ARG WITH_DEXTR
+ENV WITH_DEXTR=${WITH_DEXTR}
+ENV DEXTR_MODEL_DIR=${HOME}/models/dextr
+RUN if [ "$WITH_DEXTR" = "yes" ]; then \
+        mkdir ${DEXTR_MODEL_DIR} -p && \
+        wget https://download.01.org/openvinotoolkit/models_contrib/cvat/dextr_model_v1.zip -O ${DEXTR_MODEL_DIR}/dextr.zip && \
+        unzip ${DEXTR_MODEL_DIR}/dextr.zip -d ${DEXTR_MODEL_DIR} && rm ${DEXTR_MODEL_DIR}/dextr.zip; \
+    fi
+
 COPY ssh ${HOME}/.ssh
 COPY cvat/ ${HOME}/cvat
 COPY tests ${HOME}/tests

README.md

@@ -82,6 +82,7 @@ docker-compose -f docker-compose.yml -f components/cuda/docker-compose.cuda.yml
 -   [Analytics: management and monitoring of data annotation team](components/analytics/README.md)
 -   [TF Object Detection API: auto annotation](components/tf_annotation/README.md)
 -   [Support for NVIDIA GPUs](components/cuda/README.md)
+-   [Semi-automatic segmentation with Deep Extreme Cut](cvat/apps/dextr_segmentation/README.md)
 
 ### Create superuser account
 

cvat/apps/dextr_segmentation/README.md

@@ -0,0 +1,26 @@
+# Semi-Automatic Segmentation with [Deep Extreme Cut](http://www.vision.ee.ethz.ch/~cvlsegmentation/dextr/)
+
+## About the application
+
+The application allows to use deep learning model for semi-automatic semantic and instance segmentation.
+You can get a segmentation polygon from four (or more) extreme points of object.
+This application uses the pre-trained DEXTR model which has been converted to Inference Engine format.
+
+We are grateful to K.K. Maninis, S. Caelles, J. Pont-Tuset, and L. Van Gool who permitted using their models in our tool
+
+## Installation
+
+```bash
+# Build image with Deep Extreme Cut application (OpenVINO component is also needed)
+docker-compose -f docker-compose.yml -f components/openvino/docker-compose.openvino.yml -f cvat/apps/dextr_segmentation/docker-compose.dextr.yml build
+```
+
+## Running
+
+1.  Open a job
+2.  Select "Auto Segmentation" in the list of shapes
+3.  Run the draw mode as usually (by press the "Create Shape" button or by "N" shortcut)
+4.  Click four-six (or more if need) extreme points of an object
+5.  Close the draw mode (the same way like start it)
+6.  Wait a moment and you will receive class agnostic annotation polygon
+7.  You can close an annotation request if it is so long (in case if it is queued to rq worker and all workers are busy)

cvat/apps/dextr_segmentation/__init__.py

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

cvat/apps/dextr_segmentation/apps.py

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

cvat/apps/dextr_segmentation/dextr.py

@@ -0,0 +1,112 @@
+
+# Copyright (C) 2018 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+from openvino.inference_engine import IENetwork, IEPlugin
+
+import os
+import cv2
+import PIL
+import numpy as np
+
+_IE_CPU_EXTENSION = os.getenv("IE_CPU_EXTENSION", "libcpu_extension_avx2.so")
+_IE_PLUGINS_PATH = os.getenv("IE_PLUGINS_PATH", None)
+
+_DEXTR_MODEL_DIR = os.getenv("DEXTR_MODEL_DIR", None)
+_DEXTR_PADDING = 50
+_DEXTR_TRESHOLD = 0.9
+_DEXTR_SIZE = 512
+
+class DEXTR_HANDLER:
+    _plugin = None
+    _network = None
+    _exec_network = None
+    _input_blob = None
+    _output_blob = None
+    _xml = None
+    _bin = None
+
+    def __init__(self):
+        if not _IE_PLUGINS_PATH:
+            raise Exception("IE_PLUGINS_PATH is not defined")
+        if not _DEXTR_MODEL_DIR:
+            raise Exception("DEXTR_MODEL_DIR is not defined")
+        self._xml = os.path.join(_DEXTR_MODEL_DIR, 'dextr.xml')
+        self._bin = os.path.join(_DEXTR_MODEL_DIR, 'dextr.bin')
+
+    def handle(self, im_path, points):
+        # Lazy initialization
+        if not self._plugin:
+            self._plugin = IEPlugin(device="CPU", plugin_dirs=[_IE_PLUGINS_PATH])
+            self._plugin.add_cpu_extension(os.path.join(_IE_PLUGINS_PATH, _IE_CPU_EXTENSION))
+            self._network = IENetwork.from_ir(model=self._xml, weights=self._bin)
+            self._input_blob = next(iter(self._network.inputs))
+            self._output_blob = next(iter(self._network.outputs))
+            self._exec_network = self._plugin.load(network=self._network)
+
+        image = PIL.Image.open(im_path)
+        numpy_image = np.array(image)
+        points = np.asarray([[int(p["x"]), int(p["y"])] for p in points], dtype=int)
+        bounding_box = (
+            max(min(points[:, 0]) - _DEXTR_PADDING, 0),
+            max(min(points[:, 1]) - _DEXTR_PADDING, 0),
+            min(max(points[:, 0]) + _DEXTR_PADDING, numpy_image.shape[1] - 1),
+            min(max(points[:, 1]) + _DEXTR_PADDING, numpy_image.shape[0] - 1)
+        )
+
+        # Prepare an image
+        numpy_cropped = np.array(image.crop(bounding_box))
+        resized = cv2.resize(numpy_cropped, (_DEXTR_SIZE, _DEXTR_SIZE),
+            interpolation = cv2.INTER_CUBIC).astype(np.float32)
+
+        # Make a heatmap
+        points = points - [min(points[:, 0]), min(points[:, 1])] + [_DEXTR_PADDING, _DEXTR_PADDING]
+        points = (points * [_DEXTR_SIZE / numpy_cropped.shape[1], _DEXTR_SIZE / numpy_cropped.shape[0]]).astype(int)
+        heatmap = np.zeros(shape=resized.shape[:2], dtype=np.float64)
+        for point in points:
+            gaussian_x_axis = np.arange(0, _DEXTR_SIZE, 1, float) - point[0]
+            gaussian_y_axis = np.arange(0, _DEXTR_SIZE, 1, float)[:, np.newaxis] - point[1]
+            gaussian = np.exp(-4 * np.log(2) * ((gaussian_x_axis ** 2 + gaussian_y_axis ** 2) / 100)).astype(np.float64)
+            heatmap = np.maximum(heatmap, gaussian)
+        cv2.normalize(heatmap,  heatmap, 0, 255, cv2.NORM_MINMAX)
+
+        # Concat an image and a heatmap
+        input_dextr = np.concatenate((resized, heatmap[:, :, np.newaxis].astype(resized.dtype)), axis=2)
+        input_dextr = input_dextr.transpose((2,0,1))
+
+        np.set_printoptions(threshold=np.nan)
+        pred = self._exec_network.infer(inputs={self._input_blob: input_dextr[np.newaxis, ...]})[self._output_blob][0, 0, :, :]
+        pred = cv2.resize(pred, tuple(reversed(numpy_cropped.shape[:2])), interpolation = cv2.INTER_CUBIC)
+        result = np.zeros(numpy_image.shape[:2])
+        result[bounding_box[1]:bounding_box[1] + pred.shape[0], bounding_box[0]:bounding_box[0] + pred.shape[1]] = pred > _DEXTR_TRESHOLD
+
+        # Convert a mask to a polygon
+        result = np.array(result, dtype=np.uint8)
+        cv2.normalize(result,result,0,255,cv2.NORM_MINMAX)
+        contours = None
+        if int(cv2.__version__.split('.')[0]) > 3:
+            contours = cv2.findContours(result, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)[0]
+        else:
+            contours = cv2.findContours(result, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)[1]
+
+        contours = max(contours, key=lambda arr: arr.size)
+        if contours.shape.count(1):
+            contours = np.squeeze(contours)
+        if contours.size < 3 * 2:
+            raise Exception('Less then three point have been detected. Can not build a polygon.')
+
+        result = ""
+        for point in contours:
+            result += "{},{} ".format(int(point[0]), int(point[1]))
+        result = result[:-1]
+
+        return result
+
+    def __del__(self):
+        if self._exec_network:
+            del self._exec_network
+        if self._network:
+            del self._network
+        if self._plugin:
+            del self._plugin

cvat/apps/dextr_segmentation/docker-compose.dextr.yml

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