Coverage for src/pythia/utils/maskrcnn.py: 51%
65 statements
« prev ^ index » next coverage.py v6.4.4, created at 2022-10-07 19:27 +0000
1"""Simple python MASKRCNN output parser.
3The function `extract_maskrcnn_mask` should be used.
5The core function, `resize_mask_vec`, is based on its cpp counterpart
6`resizeMask`, from deepstream cpp sources. It was ported to python
7and vectorized via numpy to improve speed.
9Author: <Pablo Woolvett pablowoolvett@gmail.com>
11"""
13from typing import List
14from typing import Tuple
15from typing import TypedDict
17import numpy as np
18import pyds
20from pythia.utils.ext import grouped
23def _gen_ranges(
24 original_height: int,
25 original_width: int,
26 target_height: int,
27 target_width: int,
28) -> Tuple[np.ndarray, np.ndarray]:
29 ratio_h = float(original_height / target_height)
30 ratio_w = float(original_width / target_width)
32 height = np.arange(0, original_height, ratio_h)
33 width = np.arange(0, original_width, ratio_w)
34 return height, width
37def _gen_clips(
38 width: np.ndarray,
39 original_width: int,
40 height: np.ndarray,
41 original_height: int,
42) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
44 left = np.clip(np.floor(width), 0.0, original_width - 1)
45 right = np.clip(np.ceil(width), 0.0, original_width - 1)
46 top = np.clip(np.floor(height), 0.0, original_height - 1)
47 bottom = np.clip(np.ceil(height), 0.0, original_height - 1)
48 return left, right, top, bottom
51def _gen_idxs(
52 original_width: int,
53 left: np.ndarray,
54 right: np.ndarray,
55 top: np.ndarray,
56 bottom: np.ndarray,
57) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
59 left_top_idx = np.add.outer(top * original_width, left).astype(int)
60 right_top_idx = np.add.outer(top * original_width, right).astype(int)
61 left_bottom_idx = np.add.outer(bottom * original_width, left).astype(int)
62 right_bottom_idx = np.add.outer(bottom * original_width, right).astype(int)
64 return left_top_idx, right_top_idx, left_bottom_idx, right_bottom_idx
67def _take_vals(
68 src,
69 *idxmats,
70):
71 return tuple(src.take(idxmat) for idxmat in idxmats)
74def _interpolate( # noqa: R0913
75 width: np.ndarray,
76 left: np.ndarray,
77 height: np.ndarray,
78 top: np.ndarray,
79 left_top_val: np.ndarray,
80 right_top_val: np.ndarray,
81 left_bottom_val: np.ndarray,
82 right_bottom_val: np.ndarray,
83) -> np.ndarray:
84 delta_w = width - left
85 top_lerp = left_top_val + (right_top_val - left_top_val) * delta_w
86 bottom_lerp = (
87 left_bottom_val + (right_bottom_val - left_bottom_val) * delta_w
88 )
89 return top_lerp + ((bottom_lerp - top_lerp).T * (height - top)).T
92def resize_mask_vec( # noqa: R0914
93 src: np.ndarray,
94 src_shape: Tuple[int, int],
95 target_shape: Tuple[int, int],
96 threshold: float,
97) -> np.ndarray:
98 """Resize mask from original deepstream object into numpy array.
100 Args:
101 src: Mask array from deepstream object.
102 src_shape: Shape of the original mask in (height,width) format.
103 target_shape: Shape of the target mask in (height,width) format.
104 threshold: Threshold for the mask.
106 Returns:
107 A 2d binary mask of np.uint8 valued 0 and 255.
109 See Also:
110 * `extract_maskrcnn_mask` in this module for sample usage from
111 deepstream.
112 * `resizeMask` function at
113 `sample_apps/deepstream-mrcnn-app/deepstream_mrcnn_test.cpp`
115 """
117 original_height, original_width = src_shape
118 target_height, target_width = target_shape
120 height, width = _gen_ranges(
121 original_height,
122 original_width,
123 target_height,
124 target_width,
125 )
127 left, right, top, bottom = _gen_clips(
128 width,
129 original_width,
130 height,
131 original_height,
132 )
134 left_top_idx, right_top_idx, left_bottom_idx, right_bottom_idx = _gen_idxs(
135 original_width,
136 left,
137 right,
138 top,
139 bottom,
140 )
142 (
143 left_top_val,
144 right_top_val,
145 left_bottom_val,
146 right_bottom_val,
147 ) = _take_vals(
148 src,
149 left_top_idx,
150 right_top_idx,
151 left_bottom_idx,
152 right_bottom_idx,
153 )
155 lerp = _interpolate(
156 width,
157 left,
158 height,
159 top,
160 left_top_val,
161 right_top_val,
162 left_bottom_val,
163 right_bottom_val,
164 )
166 ret = np.zeros_like(lerp, dtype=np.uint8)
167 ret[lerp >= threshold] = 255
168 return ret
171def extract_maskrcnn_mask(obj_meta: pyds.NvDsObjectMeta) -> np.ndarray:
172 """Extract maskrcnn mask from deepstream object.
174 Args:
175 obj_meta: Deepstream object meta from detection.
177 Returns:
178 A 2d binary mask of np.uint8 valued 0 and 255.
180 Example:
181 >>> obj_meta = pyds.NvDsObjectMeta.cast(...)
182 >>> mask = extract_maskrcnn_mask(obj_meta)
183 >>> mask.shape, mask.dtype
184 ((300,100), dtype('uint8'))
186 See Also:
187 `resize_mask_vec` for the internal implementation.
189 """
190 rect_height = int(np.ceil(obj_meta.rect_params.height))
191 rect_width = int(np.ceil(obj_meta.rect_params.width))
192 return resize_mask_vec(
193 obj_meta.mask_params.data,
194 (obj_meta.mask_params.height, obj_meta.mask_params.width),
195 (rect_height, rect_width),
196 obj_meta.mask_params.threshold,
197 )
200class DsBBox(TypedDict):
201 """Deepstream-style Bounding box."""
203 top: float
204 height: float
205 width: float
206 left: float
209def polygon_to_bbox(
210 polygons: List[List[int]],
211 top: float,
212 left: float,
213) -> DsBBox:
214 """Generate bounding box from polygons.
216 The polygons are assumed to be relative to a mask.
218 Args:
219 polygons: collection of polygons, where each polygon is a
220 sequence of the form 'y0,x0,y1,x1,...,yn,xn'.
221 top: top offset from the mask.
222 left: left offset from the mask.
224 Returns:
225 Bounding box which circumscribes the mask.
227 """
228 x_max = 0
229 y_max = 0
230 x_min = 0
231 y_min = 0
232 for polygon in polygons:
233 coords = np.array(list(grouped(polygon, 2)))
234 y_max_, x_max_ = coords.max(0)
235 y_min_, x_min_ = coords.min(0)
236 x_max = max(x_max, x_max_)
237 y_max = max(y_max, y_max_)
238 x_min = min(x_min, x_min_)
239 y_min = min(y_min, y_min_)
240 return {
241 "top": top + y_min,
242 "left": left + x_min,
243 "height": x_max - x_min,
244 "width": y_max - y_min,
245 }