# Copyright (C) 2020-2022 Intel Corporation
# Copyright (C) 2022 CVAT.ai Corporation
#
# SPDX-License-Identifier: MIT
import os
import cv2
import numpy as np
from model_loader import ModelLoader
from shared import to_cvat_mask
class PixelLinkDecoder():
def __init__(self, pixel_threshold, link_threshold):
four_neighbours = False
if four_neighbours:
self._get_neighbours = self._get_neighbours_4
else:
self._get_neighbours = self._get_neighbours_8
self.pixel_conf_threshold = pixel_threshold
self.link_conf_threshold = link_threshold
def decode(self, height, width, detections: dict):
self.image_height = height
self.image_width = width
self.pixel_scores = self._set_pixel_scores(detections['model/segm_logits/add'])
self.link_scores = self._set_link_scores(detections['model/link_logits_/add'])
self.pixel_mask = self.pixel_scores >= self.pixel_conf_threshold
self.link_mask = self.link_scores >= self.link_conf_threshold
self.points = list(zip(*np.where(self.pixel_mask)))
self.h, self.w = np.shape(self.pixel_mask)
self.group_mask = dict.fromkeys(self.points, -1)
self.bboxes = None
self.root_map = None
self.mask = None
self._decode()
def _softmax(self, x, axis=None):
return np.exp(x - self._logsumexp(x, axis=axis, keepdims=True))
# pylint: disable=no-self-use
def _logsumexp(self, a, axis=None, b=None, keepdims=False, return_sign=False):
if b is not None:
a, b = np.broadcast_arrays(a, b)
if np.any(b == 0):
a = a + 0. # promote to at least float
a[b == 0] = -np.inf
a_max = np.amax(a, axis=axis, keepdims=True)
if a_max.ndim > 0:
a_max[~np.isfinite(a_max)] = 0
elif not np.isfinite(a_max):
a_max = 0
if b is not None:
b = np.asarray(b)
tmp = b * np.exp(a - a_max)
else:
tmp = np.exp(a - a_max)
# suppress warnings about log of zero
with np.errstate(divide='ignore'):
s = np.sum(tmp, axis=axis, keepdims=keepdims)
if return_sign:
sgn = np.sign(s)
s *= sgn # /= makes more sense but we need zero -> zero
out = np.log(s)
if not keepdims:
a_max = np.squeeze(a_max, axis=axis)
out += a_max
if return_sign:
return out, sgn
else:
return out
def _set_pixel_scores(self, pixel_scores):
"get softmaxed properly shaped pixel scores"
tmp = np.transpose(pixel_scores, (0, 2, 3, 1))
return self._softmax(tmp, axis=-1)[0, :, :, 1]
def _set_link_scores(self, link_scores):
"get softmaxed properly shaped links scores"
tmp = np.transpose(link_scores, (0, 2, 3, 1))
tmp_reshaped = tmp.reshape(tmp.shape[:-1] + (8, 2))
return self._softmax(tmp_reshaped, axis=-1)[0, :, :, :, 1]
def _find_root(self, point):
root = point
update_parent = False
tmp = self.group_mask[root]
while tmp is not -1:
root = tmp
tmp = self.group_mask[root]
update_parent = True
if update_parent:
self.group_mask[point] = root
return root
def _join(self, p1, p2):
root1 = self._find_root(p1)
root2 = self._find_root(p2)
if root1 != root2:
self.group_mask[root2] = root1
def _get_index(self, root):
if root not in self.root_map:
self.root_map[root] = len(self.root_map) + 1
return self.root_map[root]
def _get_all(self):
self.root_map = {}
self.mask = np.zeros_like(self.pixel_mask, dtype=np.int32)
for point in self.points:
point_root = self._find_root(point)
bbox_idx = self._get_index(point_root)
self.mask[point] = bbox_idx
def _get_neighbours_8(self, x, y):
w, h = self.w, self.h
tmp = [(0, x - 1, y - 1), (1, x, y - 1),
(2, x + 1, y - 1), (3, x - 1, y),
(4, x + 1, y), (5, x - 1, y + 1),
(6, x, y + 1), (7, x + 1, y + 1)]
return [i for i in tmp if i[1] >= 0 and i[1] < w and i[2] >= 0 and i[2] < h]
def _get_neighbours_4(self, x, y):
w, h = self.w, self.h
tmp = [(1, x, y - 1),
(3, x - 1, y),
(4, x + 1, y),
(6, x, y + 1)]
return [i for i in tmp if i[1] >= 0 and i[1] < w and i[2] >= 0 and i[2] < h]
def _mask_to_bboxes(self, min_area=300, min_height=10):
self.bboxes = []
max_bbox_idx = self.mask.max()
mask_tmp = cv2.resize(self.mask, (self.image_width, self.image_height), interpolation=cv2.INTER_NEAREST)
for bbox_idx in range(1, max_bbox_idx + 1):
bbox_mask = mask_tmp == bbox_idx
cnts, _ = cv2.findContours(bbox_mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
if len(cnts) == 0:
continue
cnt = cnts[0]
rect, w, h = self._min_area_rect(cnt)
if min(w, h) < min_height:
continue
if w * h < min_area:
continue
self.bboxes.append(self._order_points(rect))
# pylint: disable=no-self-use
def _min_area_rect(self, cnt):
rect = cv2.minAreaRect(cnt)
w, h = rect[1]
box = cv2.boxPoints(rect)
box = np.int0(box)
return box, w, h
# pylint: disable=no-self-use
def _order_points(self, rect):
""" (x, y)
Order: TL, TR, BR, BL
"""
tmp = np.zeros_like(rect)
sums = rect.sum(axis=1)
tmp[0] = rect[np.argmin(sums)]
tmp[2] = rect[np.argmax(sums)]
diff = np.diff(rect, axis=1)
tmp[1] = rect[np.argmin(diff)]
tmp[3] = rect[np.argmax(diff)]
return tmp
def _decode(self):
for point in self.points:
y, x = point
neighbours = self._get_neighbours(x, y)
for n_idx, nx, ny in neighbours:
link_value = self.link_mask[y, x, n_idx]
pixel_cls = self.pixel_mask[ny, nx]
if link_value and pixel_cls:
self._join(point, (ny, nx))
self._get_all()
self._mask_to_bboxes()
class ModelHandler:
def __init__(self, labels):
base_dir = os.path.abspath(os.environ.get("MODEL_PATH",
"/opt/nuclio/open_model_zoo/intel/text-detection-0004/FP32"))
model_xml = os.path.join(base_dir, "text-detection-0004.xml")
model_bin = os.path.join(base_dir, "text-detection-0004.bin")
self.model = ModelLoader(model_xml, model_bin)
self.labels = labels
def infer(self, image, pixel_threshold, link_threshold):
output_layer = self.model.infer(image)
results = []
obj_class = 1
pcd = PixelLinkDecoder(pixel_threshold, link_threshold)
pcd.decode(image.height, image.width, output_layer)
for box in pcd.bboxes:
mask = pcd.pixel_mask
mask = np.array(mask, dtype=np.uint8)
mask = cv2.resize(mask, dsize=(image.width, image.height), interpolation=cv2.INTER_CUBIC)
cv2.normalize(mask, mask, 0, 255, cv2.NORM_MINMAX)
box = box.ravel().tolist()
x_min = min(box[::2])
x_max = max(box[::2])
y_min = min(box[1::2])
y_max = max(box[1::2])
cvat_mask = to_cvat_mask((x_min, y_min, x_max, y_max), mask)
results.append({
"confidence": None,
"label": self.labels.get(obj_class, "unknown"),
"points": box,
"mask": cvat_mask,
"type": "mask",
})
return results