Siammask tracker as DL serverless function

CHANGELOG.md

@@ -6,6 +6,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [1.1.0] - Unreleased
 ### Added
+- Siammask tracker as DL serverless function (<https://github.com/opencv/cvat/pull/1988>)
 - [Datumaro] Added model info and source info commands (<https://github.com/opencv/cvat/pull/1973>)
 
 ### Changed

cvat/apps/lambda_manager/views.py

@@ -155,6 +155,12 @@ def invoke(self, db_task, data):
                     payload.update({
                         "max_distance": max_distance
                     })
+            elif self.kind == LambdaType.TRACKER:
+                payload.update({
+                    "image": self._get_image(db_task, data["frame"], quality),
+                    "shape": data.get("shape", None),
+                    "state": data.get("state", None)
+                })
             else:
                 raise ValidationError(
                     '`{}` lambda function has incorrect type: {}'
@@ -469,7 +475,6 @@ def __call__(function, task, quality, cleanup, **kwargs):
             LambdaJob._call_reid(function, db_task, quality,
                 kwargs.get("threshold"), kwargs.get("max_distance"))
 
-
 def return_response(success_code=status.HTTP_200_OK):
     def wrap_response(func):
         @wraps(func)

serverless/deploy.sh

@@ -46,4 +46,8 @@ nuctl deploy --project-name cvat \
     --path $SCRIPT_DIR/tensorflow/faster_rcnn_inception_v2_coco/nuclio \
     --platform local
 
+nuctl deploy --project-name cvat \
+    --path $SCRIPT_DIR/pytorch/foolwood/siammask/nuclio \
+    --platform local
+
 nuctl get function

serverless/pytorch/foolwood/siammask/nuclio/function.yaml

@@ -26,21 +26,20 @@ spec:
         value: /opt/nuclio
       - kind: RUN
         value: conda create -y -n siammask python=3.6
-      - kind: RUN
-        value: source activate siammask
+      - kind: SHELL
+        value: '["conda", "run", "-n", "siammask", "/bin/bash", "-c"]'
       - kind: RUN
         value: git clone https://github.com/foolwood/SiamMask.git
       - kind: RUN
-        value: pip install -r SiamMask/requirements.txt
+        value: pip install -r SiamMask/requirements.txt jsonpickle
       - kind: RUN
         value: conda install -y gcc_linux-64
       - kind: RUN
         value: cd SiamMask && bash make.sh && cd -
       - kind: RUN
         value: wget -P SiamMask/experiments/siammask_sharp http://www.robots.ox.ac.uk/~qwang/SiamMask_DAVIS.pth
-
-      - kind: WORKDIR
-        value: /opt/nuclio/pysot
+      - kind: ENTRYPOINT
+        value: '["conda", "run", "-n", "siammask"]'
 
   triggers:
     myHttpTrigger:

serverless/pytorch/foolwood/siammask/nuclio/model_handler.py

@@ -4,14 +4,18 @@
 
 from tools.test import *
 import os
+from copy import copy
+import jsonpickle
+import numpy as np
 
 class ModelHandler:
     def __init__(self):
         # Setup device
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         torch.backends.cudnn.benchmark = True
 
-        base_dir = "/opt/nuclio/SiamMask/experiments/siammask_sharp"
+        base_dir = os.environ.get("MODEL_PATH",
+            "/opt/nuclio/SiamMask/experiments/siammask_sharp")
         class configPath:
             config = os.path.join(base_dir, "config_davis.json")
 
@@ -21,18 +25,42 @@ class configPath:
         self.siammask = load_pretrain(siammask, os.path.join(base_dir, "SiamMask_DAVIS.pth"))
         self.siammask.eval().to(self.device)
 
+    def encode_state(self, state):
+        state['net.zf'] = state['net'].zf
+        state.pop('net', None)
+        state.pop('mask', None)
+
+        for k,v in state.items():
+            state[k] = jsonpickle.encode(v)
+
+        return state
+
+    def decode_state(self, state):
+        for k,v in state.items():
+            state[k] = jsonpickle.decode(v)
+
+        state['net'] = copy(self.siammask)
+        state['net'].zf = state['net.zf']
+        del state['net.zf']
+
+        return state
 
     def infer(self, image, shape, state):
+        image = np.array(image)
         if state is None: # init tracking
-            x, y, w, h = shape
-            target_pos = np.array([x + w / 2, y + h / 2])
-            target_sz = np.array([w, h])
-            state = siamese_init(image, target_pos, target_sz, self.siammask,
+            xtl, ytl, xbr, ybr = shape
+            target_pos = np.array([(xtl + xbr) / 2, (ytl + ybr) / 2])
+            target_sz = np.array([xbr - xtl, ybr - ytl])
+            siammask = copy(self.siammask) # don't modify self.siammask
+            state = siamese_init(image, target_pos, target_sz, siammask,
                 self.config['hp'], device=self.device)
+            state = self.encode_state(state)
         else: # track
-            state = siamese_track(state, image, mask_enable=True, refine_enable=True,
-                device=self.device)
-            shape = state['ploygon'].flatten()
+            state = self.decode_state(state)
+            state = siamese_track(state, image, mask_enable=True,
+                refine_enable=True, device=self.device)
+            shape = state['ploygon'].flatten().tolist()
+            state = self.encode_state(state)
 
         return {"shape": shape, "state": state}