* 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 <nikita.manovich@intel.com>main
Xiong Jie
4 years ago
committed by
GitHub
parent
0473da064d
commit
2d75101c73
@ -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
|
||||
@ -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)
|
||||
@ -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
|
||||
Loading…
Reference in New Issue