From 2d75101c73e5f828ffbef04c2f4e56b449d73cf9 Mon Sep 17 00:00:00 2001 From: Xiong Jie <32722944+xiong-jie-y@users.noreply.github.com> Date: Wed, 29 Sep 2021 20:39:11 +0900 Subject: [PATCH] Added HRNet18 click based interactive segmentation (#3729) (#3740) * Added HRNet18 click based interactive segmentation (#3729) The commit adds serverless container that has a click based interactive segmentation model(HRNet18) that performs better than f-BRS method statistically. * fixup! Added HRNet18 click based interactive segmentation (#3729) Co-authored-by: Nikita Manovich --- CHANGELOG.md | 1 + README.md | 1 + .../saic-vul/hrnet/nuclio/function-gpu.yaml | 78 +++++++++++++++++++ .../pytorch/saic-vul/hrnet/nuclio/main.py | 33 ++++++++ .../saic-vul/hrnet/nuclio/model_handler.py | 68 ++++++++++++++++ .../en/docs/manual/advanced/ai-tools.md | 8 ++ 6 files changed, 189 insertions(+) create mode 100644 serverless/pytorch/saic-vul/hrnet/nuclio/function-gpu.yaml create mode 100644 serverless/pytorch/saic-vul/hrnet/nuclio/main.py create mode 100644 serverless/pytorch/saic-vul/hrnet/nuclio/model_handler.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c1c49816..116b32ae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - cvat-ui: support cloud storages () +- interactor: add HRNet interactive segmentation serverless function () - Added GPU implementation for SiamMask, reworked tracking approach () ### Changed diff --git a/README.md b/README.md index 697256d1..0e940675 100644 --- a/README.md +++ b/README.md @@ -88,6 +88,7 @@ For more information about supported formats look at the | [Text detection v4](/serverless/openvino/omz/intel/text-detection-0004/nuclio) | detector | OpenVINO | X | | | [SiamMask](/serverless/pytorch/foolwood/siammask/nuclio) | tracker | PyTorch | X | | | [f-BRS](/serverless/pytorch/saic-vul/fbrs/nuclio) | interactor | PyTorch | X | | +| [HRNet](/serverless/pytorch/saic-vul/hrnet/nuclio) | interactor | PyTorch | | X | | [Inside-Outside Guidance](/serverless/pytorch/shiyinzhang/iog/nuclio) | interactor | PyTorch | X | | | [Faster RCNN](/serverless/tensorflow/faster_rcnn_inception_v2_coco/nuclio) | detector | TensorFlow | X | X | | [Mask RCNN](/serverless/tensorflow/matterport/mask_rcnn/nuclio) | detector | TensorFlow | X | X | diff --git a/serverless/pytorch/saic-vul/hrnet/nuclio/function-gpu.yaml b/serverless/pytorch/saic-vul/hrnet/nuclio/function-gpu.yaml new file mode 100644 index 00000000..c50bb77d --- /dev/null +++ b/serverless/pytorch/saic-vul/hrnet/nuclio/function-gpu.yaml @@ -0,0 +1,78 @@ +metadata: + name: pth-saic-vul-hrnet + namespace: cvat + annotations: + name: HRNET + type: interactor + spec: + framework: pytorch + min_pos_points: 1 + min_neg_points: 0 + help_message: The interactor allows to get a mask for an object using positive points, and negative points + +spec: + description: HRNet18 for click based interactive segmentation + runtime: 'python:3.8' + handler: main:handler + eventTimeout: 30s + env: + - name: PYTHONPATH + value: /opt/nuclio/hrnet + + build: + image: cvat/pth.saic-vul.hrnet + baseImage: ubuntu:20.04 + + directives: + preCopy: + - kind: ENV + value: DEBIAN_FRONTEND=noninteractive + - kind: RUN + value: apt-get update && apt-get install software-properties-common -y + - kind: RUN + value: add-apt-repository ppa:deadsnakes/ppa + - kind: RUN + value: apt-get update && apt-get install -y --no-install-recommends build-essential git curl libglib2.0-0 software-properties-common python3 python3.6-dev python3-pip python3-tk + - kind: RUN + value: pip3 install --upgrade pip + - kind: WORKDIR + value: /opt/nuclio + - kind: RUN + value: git clone https://github.com/saic-vul/ritm_interactive_segmentation.git hrnet + - kind: WORKDIR + value: /opt/nuclio/hrnet + - kind: RUN + value: apt-get install -y --no-install-recommends wget + - kind: RUN + value: wget https://github.com/saic-vul/ritm_interactive_segmentation/releases/download/v1.0/coco_lvis_h18_itermask.pth + - kind: RUN + value: pip3 install setuptools + - kind: RUN + value: pip3 install -r requirements.txt + - kind: RUN + value: apt update && apt install -y libgl1-mesa-glx + - kind: RUN + value: pip3 uninstall torch torch vision -y + - kind: RUN + value: pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html + - kind: WORKDIR + value: /opt/nuclio + + triggers: + myHttpTrigger: + maxWorkers: 1 + kind: 'http' + workerAvailabilityTimeoutMilliseconds: 10000 + attributes: + maxRequestBodySize: 33554432 # 32MB + + resources: + limits: + nvidia.com/gpu: 1 + + platform: + attributes: + restartPolicy: + name: always + maximumRetryCount: 3 + mountMode: volume diff --git a/serverless/pytorch/saic-vul/hrnet/nuclio/main.py b/serverless/pytorch/saic-vul/hrnet/nuclio/main.py new file mode 100644 index 00000000..297d1020 --- /dev/null +++ b/serverless/pytorch/saic-vul/hrnet/nuclio/main.py @@ -0,0 +1,33 @@ +# Copyright (C) 2021 Intel Corporation +# +# SPDX-License-Identifier: MIT + +import json +import base64 +from PIL import Image +import io +from model_handler import ModelHandler + +def init_context(context): + context.logger.info("Init context... 0%") + + model = ModelHandler() + context.user_data.model = model + + context.logger.info("Init context...100%") + +def handler(context, event): + context.logger.info("call handler") + data = event.body + pos_points = data["pos_points"] + neg_points = data["neg_points"] + threshold = data.get("threshold", 0.5) + buf = io.BytesIO(base64.b64decode(data["image"])) + image = Image.open(buf) + + polygon = context.user_data.model.handle(image, pos_points, + neg_points, threshold) + return context.Response(body=json.dumps(polygon), + headers={}, + content_type='application/json', + status_code=200) diff --git a/serverless/pytorch/saic-vul/hrnet/nuclio/model_handler.py b/serverless/pytorch/saic-vul/hrnet/nuclio/model_handler.py new file mode 100644 index 00000000..1d7c2b47 --- /dev/null +++ b/serverless/pytorch/saic-vul/hrnet/nuclio/model_handler.py @@ -0,0 +1,68 @@ +# Copyright (C) 2021 Intel Corporation +# +# SPDX-License-Identifier: MIT + +import torch +import numpy as np +import cv2 +import os + +from isegm.inference import utils +from isegm.inference.predictors import get_predictor +from isegm.inference.clicker import Clicker, Click + +def convert_mask_to_polygon(mask): + mask = np.array(mask, dtype=np.uint8) + cv2.normalize(mask, mask, 0, 255, cv2.NORM_MINMAX) + contours = None + if int(cv2.__version__.split('.')[0]) > 3: + contours = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)[0] + else: + contours = cv2.findContours(mask, 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.') + + polygon = [] + for point in contours: + polygon.append([int(point[0]), int(point[1])]) + + return polygon + +class ModelHandler: + def __init__(self): + torch.backends.cudnn.deterministic = True + base_dir = os.path.abspath(os.environ.get("MODEL_PATH", "/opt/nuclio/hrnet")) + model_path = os.path.join(base_dir) + + self.net = None + self.device = 'cuda' if torch.cuda.is_available() else 'cpu' + + checkpoint_path = utils.find_checkpoint(model_path, "coco_lvis_h18_itermask.pth") + self.net = utils.load_is_model(checkpoint_path, self.device) + + def handle(self, image, pos_points, neg_points, threshold): + image_nd = np.array(image) + + clicker = Clicker() + for x, y in pos_points: + click = Click(is_positive=True, coords=(y, x)) + clicker.add_click(click) + + for x, y in neg_points: + click = Click(is_positive=False, coords=(y, x)) + clicker.add_click(click) + + predictor = get_predictor(self.net, 'NoBRS', device=self.device, prob_thresh=0.49) + predictor.set_input_image(image_nd) + + object_prob = predictor.get_prediction(clicker) + if self.device == 'cuda': + torch.cuda.empty_cache() + object_mask = object_prob > threshold + polygon = convert_mask_to_polygon(object_mask) + + return polygon diff --git a/site/content/en/docs/manual/advanced/ai-tools.md b/site/content/en/docs/manual/advanced/ai-tools.md index 61f03534..fdb4c98b 100644 --- a/site/content/en/docs/manual/advanced/ai-tools.md +++ b/site/content/en/docs/manual/advanced/ai-tools.md @@ -69,6 +69,14 @@ It is recommended to run the model on GPU, if possible. ![](/images/fbrs_example.gif) +### High Resolution Net (HRNet) + +The model allows to get a mask for an object using positive points +(should be left-clicked on the foreground), and negative points +(should be right-clicked on the background, if necessary). +It is recommended to run the model on GPU, if possible. + + ### Inside-Outside-Guidance The model uses a bounding box and inside/outside points to create a mask.