We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Potential integer overflow
In the code below we see that `img` can be of type `uint8`, then `img` is multiplied by 2<sup>8</sup> and 2<sup>16</sup>, but `uint8` cannot store such large numbers.
https://github.com/openvinotoolkit/datumaro/blob/5209d42348b001b2f07a917d5dd7b76f3469c7ce/datumaro/util/mask_tools.py#L116-L119

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `datumaro/util/mask_tools.py`
Content:
```
1 
2 # Copyright (C) 2019-2020 Intel Corporation
3 #
4 # SPDX-License-Identifier: MIT
5 
6 from itertools import chain
7 import numpy as np
8 
9 from datumaro.util.image import lazy_image, load_image
10 
11 
12 def generate_colormap(length=256):
13     """
14     Generates colors using PASCAL VOC algorithm.
15 
16     Returns index -> (R, G, B) mapping.
17     """
18 
19     def get_bit(number, index):
20         return (number >> index) & 1
21 
22     colormap = np.zeros((length, 3), dtype=int)
23     indices = np.arange(length, dtype=int)
24 
25     for j in range(7, -1, -1):
26         for c in range(3):
27             colormap[:, c] |= get_bit(indices, c) << j
28         indices >>= 3
29 
30     return {
31         id: tuple(color) for id, color in enumerate(colormap)
32     }
33 
34 def invert_colormap(colormap):
35     return {
36         tuple(a): index for index, a in colormap.items()
37     }
38 
39 def check_is_mask(mask):
40     assert len(mask.shape) in {2, 3}
41     if len(mask.shape) == 3:
42         assert mask.shape[2] == 1
43 
44 _default_colormap = generate_colormap()
45 _default_unpaint_colormap = invert_colormap(_default_colormap)
46 
47 def unpaint_mask(painted_mask, inverse_colormap=None):
48     # Covert color mask to index mask
49 
50     # mask: HWC BGR [0; 255]
51     # colormap: (R, G, B) -> index
52     assert len(painted_mask.shape) == 3
53     if inverse_colormap is None:
54         inverse_colormap = _default_unpaint_colormap
55 
56     if callable(inverse_colormap):
57         map_fn = lambda a: inverse_colormap(
58                 (a >> 16) & 255, (a >> 8) & 255, a & 255
59             )
60     else:
61         map_fn = lambda a: inverse_colormap[(
62                 (a >> 16) & 255, (a >> 8) & 255, a & 255
63             )]
64 
65     painted_mask = painted_mask.astype(int)
66     painted_mask = painted_mask[:, :, 0] + \
67                    (painted_mask[:, :, 1] << 8) + \
68                    (painted_mask[:, :, 2] << 16)
69     uvals, unpainted_mask = np.unique(painted_mask, return_inverse=True)
70     palette = np.array([map_fn(v) for v in uvals],
71         dtype=np.min_scalar_type(len(uvals)))
72     unpainted_mask = palette[unpainted_mask].reshape(painted_mask.shape[:2])
73 
74     return unpainted_mask
75 
76 def paint_mask(mask, colormap=None):
77     """
78     Applies colormap to index mask
79 
80     mask: HW(C) [0; max_index] mask
81     colormap: index -> (R, G, B)
82     """
83     check_is_mask(mask)
84 
85     if colormap is None:
86         colormap = _default_colormap
87     if callable(colormap):
88         map_fn = colormap
89     else:
90         map_fn = lambda c: colormap.get(c, (-1, -1, -1))
91     palette = np.array([map_fn(c)[::-1] for c in range(256)], dtype=np.uint8)
92 
93     mask = mask.astype(np.uint8)
94     painted_mask = palette[mask].reshape((*mask.shape[:2], 3))
95     return painted_mask
96 
97 def remap_mask(mask, map_fn):
98     """
99     Changes mask elements from one colormap to another
100 
101     # mask: HW(C) [0; max_index] mask
102     """
103     check_is_mask(mask)
104 
105     return np.array([map_fn(c) for c in range(256)], dtype=np.uint8)[mask]
106 
107 def make_index_mask(binary_mask, index, dtype=None):
108     return binary_mask * np.array([index],
109         dtype=dtype or np.min_scalar_type(index))
110 
111 def make_binary_mask(mask):
112     if mask.dtype.kind == 'b':
113         return mask
114     return mask.astype(bool)
115 
116 def bgr2index(img):
117     if img.dtype.kind not in {'b', 'i', 'u'}:
118         img = img.astype(np.uint8)
119     return (img[..., 0] << 16) + (img[..., 1] << 8) + img[..., 2]
120 
121 def index2bgr(id_map):
122     return np.dstack((id_map >> 16, id_map >> 8, id_map)).astype(np.uint8)
123 
124 def load_mask(path, inverse_colormap=None):
125     mask = load_image(path, dtype=np.uint8)
126     if inverse_colormap is not None:
127         if len(mask.shape) == 3 and mask.shape[2] != 1:
128             mask = unpaint_mask(mask, inverse_colormap)
129     return mask
130 
131 def lazy_mask(path, inverse_colormap=None):
132     return lazy_image(path, lambda path: load_mask(path, inverse_colormap))
133 
134 def mask_to_rle(binary_mask):
135     # walk in row-major order as COCO format specifies
136     bounded = binary_mask.ravel(order='F')
137 
138     # add borders to sequence
139     # find boundary positions for sequences and compute their lengths
140     difs = np.diff(bounded, prepend=[1 - bounded[0]], append=[1 - bounded[-1]])
141     counts, = np.where(difs != 0)
142 
143     # start RLE encoding from 0 as COCO format specifies
144     if bounded[0] != 0:
145         counts = np.diff(counts, prepend=[0])
146     else:
147         counts = np.diff(counts)
148 
149     return {
150         'counts': counts,
151         'size': list(binary_mask.shape)
152     }
153 
154 def mask_to_polygons(mask, area_threshold=1):
155     """
156     Convert an instance mask to polygons
157 
158     Args:
159         mask: a 2d binary mask
160         tolerance: maximum distance from original points of
161             a polygon to the approximated ones
162         area_threshold: minimal area of generated polygons
163 
164     Returns:
165         A list of polygons like [[x1,y1, x2,y2 ...], [...]]
166     """
167     from pycocotools import mask as mask_utils
168     import cv2
169 
170     polygons = []
171 
172     contours, _ = cv2.findContours(mask.astype(np.uint8),
173         mode=cv2.RETR_TREE, method=cv2.CHAIN_APPROX_TC89_KCOS)
174 
175     for contour in contours:
176         if len(contour) <= 2:
177             continue
178 
179         contour = contour.reshape((-1, 2))
180 
181         if not np.array_equal(contour[0], contour[-1]):
182             contour = np.vstack((contour, contour[0])) # make polygon closed
183         contour = contour.flatten().clip(0) # [x0, y0, ...]
184 
185         # Check if the polygon is big enough
186         rle = mask_utils.frPyObjects([contour], mask.shape[0], mask.shape[1])
187         area = sum(mask_utils.area(rle))
188         if area_threshold <= area:
189             polygons.append(contour)
190     return polygons
191 
192 def crop_covered_segments(segments, width, height,
193         iou_threshold=0.0, ratio_tolerance=0.001, area_threshold=1,
194         return_masks=False):
195     """
196     Find all segments occluded by others and crop them to the visible part only.
197     Input segments are expected to be sorted from background to foreground.
198 
199     Args:
200         segments: 1d list of segment RLEs (in COCO format)
201         width: width of the image
202         height: height of the image
203         iou_threshold: IoU threshold for objects to be counted as intersected
204             By default is set to 0 to process any intersected objects
205         ratio_tolerance: an IoU "handicap" value for a situation
206             when an object is (almost) fully covered by another one and we
207             don't want make a "hole" in the background object
208         area_threshold: minimal area of included segments
209 
210     Returns:
211         A list of input segments' parts (in the same order as input):
212             [
213                 [[x1,y1, x2,y2 ...], ...], # input segment #0 parts
214                 mask1, # input segment #1 mask (if source segment is mask)
215                 [], # when source segment is too small
216                 ...
217             ]
218     """
219     from pycocotools import mask as mask_utils
220 
221     segments = [[s] for s in segments]
222     input_rles = [mask_utils.frPyObjects(s, height, width) for s in segments]
223 
224     for i, rle_bottom in enumerate(input_rles):
225         area_bottom = sum(mask_utils.area(rle_bottom))
226         if area_bottom < area_threshold:
227             segments[i] = [] if not return_masks else None
228             continue
229 
230         rles_top = []
231         for j in range(i + 1, len(input_rles)):
232             rle_top = input_rles[j]
233             iou = sum(mask_utils.iou(rle_bottom, rle_top, [0, 0]))[0]
234 
235             if iou <= iou_threshold:
236                 continue
237 
238             area_top = sum(mask_utils.area(rle_top))
239             area_ratio = area_top / area_bottom
240 
241             # If a segment is fully inside another one, skip this segment
242             if abs(area_ratio - iou) < ratio_tolerance:
243                 continue
244 
245             # Check if the bottom segment is fully covered by the top one.
246             # There is a mistake in the annotation, keep the background one
247             if abs(1 / area_ratio - iou) < ratio_tolerance:
248                 rles_top = []
249                 break
250 
251             rles_top += rle_top
252 
253         if not rles_top and not isinstance(segments[i][0], dict) \
254                 and not return_masks:
255             continue
256 
257         rle_bottom = rle_bottom[0]
258         bottom_mask = mask_utils.decode(rle_bottom).astype(np.uint8)
259 
260         if rles_top:
261             rle_top = mask_utils.merge(rles_top)
262             top_mask = mask_utils.decode(rle_top).astype(np.uint8)
263 
264             bottom_mask -= top_mask
265             bottom_mask[bottom_mask != 1] = 0
266 
267         if not return_masks and not isinstance(segments[i][0], dict):
268             segments[i] = mask_to_polygons(bottom_mask,
269                 area_threshold=area_threshold)
270         else:
271             segments[i] = bottom_mask
272 
273     return segments
274 
275 def rles_to_mask(rles, width, height):
276     from pycocotools import mask as mask_utils
277 
278     rles = mask_utils.frPyObjects(rles, height, width)
279     rles = mask_utils.merge(rles)
280     mask = mask_utils.decode(rles)
281     return mask
282 
283 def find_mask_bbox(mask):
284     cols = np.any(mask, axis=0)
285     rows = np.any(mask, axis=1)
286     x0, x1 = np.where(cols)[0][[0, -1]]
287     y0, y1 = np.where(rows)[0][[0, -1]]
288     return [x0, y0, x1 - x0, y1 - y0]
289 
290 def merge_masks(masks, start=None):
291     """
292         Merges masks into one, mask order is responsible for z order.
293         To avoid memory explosion on mask materialization, consider passing
294         a generator.
295 
296         Inputs: a sequence of index masks or (binary mask, index) pairs
297         Outputs: an index mask
298     """
299     if start is not None:
300         masks = chain([start], masks)
301 
302     it = iter(masks)
303 
304     try:
305         merged_mask = next(it)
306         if isinstance(merged_mask, tuple) and len(merged_mask) == 2:
307             merged_mask = merged_mask[0] * merged_mask[1]
308     except StopIteration:
309         return None
310 
311     for m in it:
312         if isinstance(m, tuple) and len(m) == 2:
313             merged_mask = np.where(m[0], m[1], merged_mask)
314         else:
315             merged_mask = np.where(m, m, merged_mask)
316 
317     return merged_mask
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/datumaro/util/mask_tools.py b/datumaro/util/mask_tools.py
--- a/datumaro/util/mask_tools.py
+++ b/datumaro/util/mask_tools.py
@@ -114,8 +114,8 @@
     return mask.astype(bool)
 
 def bgr2index(img):
-    if img.dtype.kind not in {'b', 'i', 'u'}:
-        img = img.astype(np.uint8)
+    if img.dtype.kind not in {'b', 'i', 'u'} or img.dtype.itemsize < 4:
+        img = img.astype(np.uint32)
     return (img[..., 0] << 16) + (img[..., 1] << 8) + img[..., 2]
 
 def index2bgr(id_map):

{"golden_diff": "diff --git a/datumaro/util/mask_tools.py b/datumaro/util/mask_tools.py\n--- a/datumaro/util/mask_tools.py\n+++ b/datumaro/util/mask_tools.py\n@@ -114,8 +114,8 @@\n     return mask.astype(bool)\n \n def bgr2index(img):\n-    if img.dtype.kind not in {'b', 'i', 'u'}:\n-        img = img.astype(np.uint8)\n+    if img.dtype.kind not in {'b', 'i', 'u'} or img.dtype.itemsize < 4:\n+        img = img.astype(np.uint32)\n     return (img[..., 0] << 16) + (img[..., 1] << 8) + img[..., 2]\n \n def index2bgr(id_map):\n", "issue": "Potential integer overflow\nIn the code below we see that `img` can be of type `uint8`, then `img` is multiplied by 2<sup>8</sup> and 2<sup>16</sup>, but `uint8` cannot store such large numbers.\r\nhttps://github.com/openvinotoolkit/datumaro/blob/5209d42348b001b2f07a917d5dd7b76f3469c7ce/datumaro/util/mask_tools.py#L116-L119\n", "before_files": [{"content": "\n# Copyright (C) 2019-2020 Intel Corporation\n#\n# SPDX-License-Identifier: MIT\n\nfrom itertools import chain\nimport numpy as np\n\nfrom datumaro.util.image import lazy_image, load_image\n\n\ndef generate_colormap(length=256):\n    \"\"\"\n    Generates colors using PASCAL VOC algorithm.\n\n    Returns index -> (R, G, B) mapping.\n    \"\"\"\n\n    def get_bit(number, index):\n        return (number >> index) & 1\n\n    colormap = np.zeros((length, 3), dtype=int)\n    indices = np.arange(length, dtype=int)\n\n    for j in range(7, -1, -1):\n        for c in range(3):\n            colormap[:, c] |= get_bit(indices, c) << j\n        indices >>= 3\n\n    return {\n        id: tuple(color) for id, color in enumerate(colormap)\n    }\n\ndef invert_colormap(colormap):\n    return {\n        tuple(a): index for index, a in colormap.items()\n    }\n\ndef check_is_mask(mask):\n    assert len(mask.shape) in {2, 3}\n    if len(mask.shape) == 3:\n        assert mask.shape[2] == 1\n\n_default_colormap = generate_colormap()\n_default_unpaint_colormap = invert_colormap(_default_colormap)\n\ndef unpaint_mask(painted_mask, inverse_colormap=None):\n    # Covert color mask to index mask\n\n    # mask: HWC BGR [0; 255]\n    # colormap: (R, G, B) -> index\n    assert len(painted_mask.shape) == 3\n    if inverse_colormap is None:\n        inverse_colormap = _default_unpaint_colormap\n\n    if callable(inverse_colormap):\n        map_fn = lambda a: inverse_colormap(\n                (a >> 16) & 255, (a >> 8) & 255, a & 255\n            )\n    else:\n        map_fn = lambda a: inverse_colormap[(\n                (a >> 16) & 255, (a >> 8) & 255, a & 255\n            )]\n\n    painted_mask = painted_mask.astype(int)\n    painted_mask = painted_mask[:, :, 0] + \\\n                   (painted_mask[:, :, 1] << 8) + \\\n                   (painted_mask[:, :, 2] << 16)\n    uvals, unpainted_mask = np.unique(painted_mask, return_inverse=True)\n    palette = np.array([map_fn(v) for v in uvals],\n        dtype=np.min_scalar_type(len(uvals)))\n    unpainted_mask = palette[unpainted_mask].reshape(painted_mask.shape[:2])\n\n    return unpainted_mask\n\ndef paint_mask(mask, colormap=None):\n    \"\"\"\n    Applies colormap to index mask\n\n    mask: HW(C) [0; max_index] mask\n    colormap: index -> (R, G, B)\n    \"\"\"\n    check_is_mask(mask)\n\n    if colormap is None:\n        colormap = _default_colormap\n    if callable(colormap):\n        map_fn = colormap\n    else:\n        map_fn = lambda c: colormap.get(c, (-1, -1, -1))\n    palette = np.array([map_fn(c)[::-1] for c in range(256)], dtype=np.uint8)\n\n    mask = mask.astype(np.uint8)\n    painted_mask = palette[mask].reshape((*mask.shape[:2], 3))\n    return painted_mask\n\ndef remap_mask(mask, map_fn):\n    \"\"\"\n    Changes mask elements from one colormap to another\n\n    # mask: HW(C) [0; max_index] mask\n    \"\"\"\n    check_is_mask(mask)\n\n    return np.array([map_fn(c) for c in range(256)], dtype=np.uint8)[mask]\n\ndef make_index_mask(binary_mask, index, dtype=None):\n    return binary_mask * np.array([index],\n        dtype=dtype or np.min_scalar_type(index))\n\ndef make_binary_mask(mask):\n    if mask.dtype.kind == 'b':\n        return mask\n    return mask.astype(bool)\n\ndef bgr2index(img):\n    if img.dtype.kind not in {'b', 'i', 'u'}:\n        img = img.astype(np.uint8)\n    return (img[..., 0] << 16) + (img[..., 1] << 8) + img[..., 2]\n\ndef index2bgr(id_map):\n    return np.dstack((id_map >> 16, id_map >> 8, id_map)).astype(np.uint8)\n\ndef load_mask(path, inverse_colormap=None):\n    mask = load_image(path, dtype=np.uint8)\n    if inverse_colormap is not None:\n        if len(mask.shape) == 3 and mask.shape[2] != 1:\n            mask = unpaint_mask(mask, inverse_colormap)\n    return mask\n\ndef lazy_mask(path, inverse_colormap=None):\n    return lazy_image(path, lambda path: load_mask(path, inverse_colormap))\n\ndef mask_to_rle(binary_mask):\n    # walk in row-major order as COCO format specifies\n    bounded = binary_mask.ravel(order='F')\n\n    # add borders to sequence\n    # find boundary positions for sequences and compute their lengths\n    difs = np.diff(bounded, prepend=[1 - bounded[0]], append=[1 - bounded[-1]])\n    counts, = np.where(difs != 0)\n\n    # start RLE encoding from 0 as COCO format specifies\n    if bounded[0] != 0:\n        counts = np.diff(counts, prepend=[0])\n    else:\n        counts = np.diff(counts)\n\n    return {\n        'counts': counts,\n        'size': list(binary_mask.shape)\n    }\n\ndef mask_to_polygons(mask, area_threshold=1):\n    \"\"\"\n    Convert an instance mask to polygons\n\n    Args:\n        mask: a 2d binary mask\n        tolerance: maximum distance from original points of\n            a polygon to the approximated ones\n        area_threshold: minimal area of generated polygons\n\n    Returns:\n        A list of polygons like [[x1,y1, x2,y2 ...], [...]]\n    \"\"\"\n    from pycocotools import mask as mask_utils\n    import cv2\n\n    polygons = []\n\n    contours, _ = cv2.findContours(mask.astype(np.uint8),\n        mode=cv2.RETR_TREE, method=cv2.CHAIN_APPROX_TC89_KCOS)\n\n    for contour in contours:\n        if len(contour) <= 2:\n            continue\n\n        contour = contour.reshape((-1, 2))\n\n        if not np.array_equal(contour[0], contour[-1]):\n            contour = np.vstack((contour, contour[0])) # make polygon closed\n        contour = contour.flatten().clip(0) # [x0, y0, ...]\n\n        # Check if the polygon is big enough\n        rle = mask_utils.frPyObjects([contour], mask.shape[0], mask.shape[1])\n        area = sum(mask_utils.area(rle))\n        if area_threshold <= area:\n            polygons.append(contour)\n    return polygons\n\ndef crop_covered_segments(segments, width, height,\n        iou_threshold=0.0, ratio_tolerance=0.001, area_threshold=1,\n        return_masks=False):\n    \"\"\"\n    Find all segments occluded by others and crop them to the visible part only.\n    Input segments are expected to be sorted from background to foreground.\n\n    Args:\n        segments: 1d list of segment RLEs (in COCO format)\n        width: width of the image\n        height: height of the image\n        iou_threshold: IoU threshold for objects to be counted as intersected\n            By default is set to 0 to process any intersected objects\n        ratio_tolerance: an IoU \"handicap\" value for a situation\n            when an object is (almost) fully covered by another one and we\n            don't want make a \"hole\" in the background object\n        area_threshold: minimal area of included segments\n\n    Returns:\n        A list of input segments' parts (in the same order as input):\n            [\n                [[x1,y1, x2,y2 ...], ...], # input segment #0 parts\n                mask1, # input segment #1 mask (if source segment is mask)\n                [], # when source segment is too small\n                ...\n            ]\n    \"\"\"\n    from pycocotools import mask as mask_utils\n\n    segments = [[s] for s in segments]\n    input_rles = [mask_utils.frPyObjects(s, height, width) for s in segments]\n\n    for i, rle_bottom in enumerate(input_rles):\n        area_bottom = sum(mask_utils.area(rle_bottom))\n        if area_bottom < area_threshold:\n            segments[i] = [] if not return_masks else None\n            continue\n\n        rles_top = []\n        for j in range(i + 1, len(input_rles)):\n            rle_top = input_rles[j]\n            iou = sum(mask_utils.iou(rle_bottom, rle_top, [0, 0]))[0]\n\n            if iou <= iou_threshold:\n                continue\n\n            area_top = sum(mask_utils.area(rle_top))\n            area_ratio = area_top / area_bottom\n\n            # If a segment is fully inside another one, skip this segment\n            if abs(area_ratio - iou) < ratio_tolerance:\n                continue\n\n            # Check if the bottom segment is fully covered by the top one.\n            # There is a mistake in the annotation, keep the background one\n            if abs(1 / area_ratio - iou) < ratio_tolerance:\n                rles_top = []\n                break\n\n            rles_top += rle_top\n\n        if not rles_top and not isinstance(segments[i][0], dict) \\\n                and not return_masks:\n            continue\n\n        rle_bottom = rle_bottom[0]\n        bottom_mask = mask_utils.decode(rle_bottom).astype(np.uint8)\n\n        if rles_top:\n            rle_top = mask_utils.merge(rles_top)\n            top_mask = mask_utils.decode(rle_top).astype(np.uint8)\n\n            bottom_mask -= top_mask\n            bottom_mask[bottom_mask != 1] = 0\n\n        if not return_masks and not isinstance(segments[i][0], dict):\n            segments[i] = mask_to_polygons(bottom_mask,\n                area_threshold=area_threshold)\n        else:\n            segments[i] = bottom_mask\n\n    return segments\n\ndef rles_to_mask(rles, width, height):\n    from pycocotools import mask as mask_utils\n\n    rles = mask_utils.frPyObjects(rles, height, width)\n    rles = mask_utils.merge(rles)\n    mask = mask_utils.decode(rles)\n    return mask\n\ndef find_mask_bbox(mask):\n    cols = np.any(mask, axis=0)\n    rows = np.any(mask, axis=1)\n    x0, x1 = np.where(cols)[0][[0, -1]]\n    y0, y1 = np.where(rows)[0][[0, -1]]\n    return [x0, y0, x1 - x0, y1 - y0]\n\ndef merge_masks(masks, start=None):\n    \"\"\"\n        Merges masks into one, mask order is responsible for z order.\n        To avoid memory explosion on mask materialization, consider passing\n        a generator.\n\n        Inputs: a sequence of index masks or (binary mask, index) pairs\n        Outputs: an index mask\n    \"\"\"\n    if start is not None:\n        masks = chain([start], masks)\n\n    it = iter(masks)\n\n    try:\n        merged_mask = next(it)\n        if isinstance(merged_mask, tuple) and len(merged_mask) == 2:\n            merged_mask = merged_mask[0] * merged_mask[1]\n    except StopIteration:\n        return None\n\n    for m in it:\n        if isinstance(m, tuple) and len(m) == 2:\n            merged_mask = np.where(m[0], m[1], merged_mask)\n        else:\n            merged_mask = np.where(m, m, merged_mask)\n\n    return merged_mask", "path": "datumaro/util/mask_tools.py"}], "after_files": [{"content": "\n# Copyright (C) 2019-2020 Intel Corporation\n#\n# SPDX-License-Identifier: MIT\n\nfrom itertools import chain\nimport numpy as np\n\nfrom datumaro.util.image import lazy_image, load_image\n\n\ndef generate_colormap(length=256):\n    \"\"\"\n    Generates colors using PASCAL VOC algorithm.\n\n    Returns index -> (R, G, B) mapping.\n    \"\"\"\n\n    def get_bit(number, index):\n        return (number >> index) & 1\n\n    colormap = np.zeros((length, 3), dtype=int)\n    indices = np.arange(length, dtype=int)\n\n    for j in range(7, -1, -1):\n        for c in range(3):\n            colormap[:, c] |= get_bit(indices, c) << j\n        indices >>= 3\n\n    return {\n        id: tuple(color) for id, color in enumerate(colormap)\n    }\n\ndef invert_colormap(colormap):\n    return {\n        tuple(a): index for index, a in colormap.items()\n    }\n\ndef check_is_mask(mask):\n    assert len(mask.shape) in {2, 3}\n    if len(mask.shape) == 3:\n        assert mask.shape[2] == 1\n\n_default_colormap = generate_colormap()\n_default_unpaint_colormap = invert_colormap(_default_colormap)\n\ndef unpaint_mask(painted_mask, inverse_colormap=None):\n    # Covert color mask to index mask\n\n    # mask: HWC BGR [0; 255]\n    # colormap: (R, G, B) -> index\n    assert len(painted_mask.shape) == 3\n    if inverse_colormap is None:\n        inverse_colormap = _default_unpaint_colormap\n\n    if callable(inverse_colormap):\n        map_fn = lambda a: inverse_colormap(\n                (a >> 16) & 255, (a >> 8) & 255, a & 255\n            )\n    else:\n        map_fn = lambda a: inverse_colormap[(\n                (a >> 16) & 255, (a >> 8) & 255, a & 255\n            )]\n\n    painted_mask = painted_mask.astype(int)\n    painted_mask = painted_mask[:, :, 0] + \\\n                   (painted_mask[:, :, 1] << 8) + \\\n                   (painted_mask[:, :, 2] << 16)\n    uvals, unpainted_mask = np.unique(painted_mask, return_inverse=True)\n    palette = np.array([map_fn(v) for v in uvals],\n        dtype=np.min_scalar_type(len(uvals)))\n    unpainted_mask = palette[unpainted_mask].reshape(painted_mask.shape[:2])\n\n    return unpainted_mask\n\ndef paint_mask(mask, colormap=None):\n    \"\"\"\n    Applies colormap to index mask\n\n    mask: HW(C) [0; max_index] mask\n    colormap: index -> (R, G, B)\n    \"\"\"\n    check_is_mask(mask)\n\n    if colormap is None:\n        colormap = _default_colormap\n    if callable(colormap):\n        map_fn = colormap\n    else:\n        map_fn = lambda c: colormap.get(c, (-1, -1, -1))\n    palette = np.array([map_fn(c)[::-1] for c in range(256)], dtype=np.uint8)\n\n    mask = mask.astype(np.uint8)\n    painted_mask = palette[mask].reshape((*mask.shape[:2], 3))\n    return painted_mask\n\ndef remap_mask(mask, map_fn):\n    \"\"\"\n    Changes mask elements from one colormap to another\n\n    # mask: HW(C) [0; max_index] mask\n    \"\"\"\n    check_is_mask(mask)\n\n    return np.array([map_fn(c) for c in range(256)], dtype=np.uint8)[mask]\n\ndef make_index_mask(binary_mask, index, dtype=None):\n    return binary_mask * np.array([index],\n        dtype=dtype or np.min_scalar_type(index))\n\ndef make_binary_mask(mask):\n    if mask.dtype.kind == 'b':\n        return mask\n    return mask.astype(bool)\n\ndef bgr2index(img):\n    if img.dtype.kind not in {'b', 'i', 'u'} or img.dtype.itemsize < 4:\n        img = img.astype(np.uint32)\n    return (img[..., 0] << 16) + (img[..., 1] << 8) + img[..., 2]\n\ndef index2bgr(id_map):\n    return np.dstack((id_map >> 16, id_map >> 8, id_map)).astype(np.uint8)\n\ndef load_mask(path, inverse_colormap=None):\n    mask = load_image(path, dtype=np.uint8)\n    if inverse_colormap is not None:\n        if len(mask.shape) == 3 and mask.shape[2] != 1:\n            mask = unpaint_mask(mask, inverse_colormap)\n    return mask\n\ndef lazy_mask(path, inverse_colormap=None):\n    return lazy_image(path, lambda path: load_mask(path, inverse_colormap))\n\ndef mask_to_rle(binary_mask):\n    # walk in row-major order as COCO format specifies\n    bounded = binary_mask.ravel(order='F')\n\n    # add borders to sequence\n    # find boundary positions for sequences and compute their lengths\n    difs = np.diff(bounded, prepend=[1 - bounded[0]], append=[1 - bounded[-1]])\n    counts, = np.where(difs != 0)\n\n    # start RLE encoding from 0 as COCO format specifies\n    if bounded[0] != 0:\n        counts = np.diff(counts, prepend=[0])\n    else:\n        counts = np.diff(counts)\n\n    return {\n        'counts': counts,\n        'size': list(binary_mask.shape)\n    }\n\ndef mask_to_polygons(mask, area_threshold=1):\n    \"\"\"\n    Convert an instance mask to polygons\n\n    Args:\n        mask: a 2d binary mask\n        tolerance: maximum distance from original points of\n            a polygon to the approximated ones\n        area_threshold: minimal area of generated polygons\n\n    Returns:\n        A list of polygons like [[x1,y1, x2,y2 ...], [...]]\n    \"\"\"\n    from pycocotools import mask as mask_utils\n    import cv2\n\n    polygons = []\n\n    contours, _ = cv2.findContours(mask.astype(np.uint8),\n        mode=cv2.RETR_TREE, method=cv2.CHAIN_APPROX_TC89_KCOS)\n\n    for contour in contours:\n        if len(contour) <= 2:\n            continue\n\n        contour = contour.reshape((-1, 2))\n\n        if not np.array_equal(contour[0], contour[-1]):\n            contour = np.vstack((contour, contour[0])) # make polygon closed\n        contour = contour.flatten().clip(0) # [x0, y0, ...]\n\n        # Check if the polygon is big enough\n        rle = mask_utils.frPyObjects([contour], mask.shape[0], mask.shape[1])\n        area = sum(mask_utils.area(rle))\n        if area_threshold <= area:\n            polygons.append(contour)\n    return polygons\n\ndef crop_covered_segments(segments, width, height,\n        iou_threshold=0.0, ratio_tolerance=0.001, area_threshold=1,\n        return_masks=False):\n    \"\"\"\n    Find all segments occluded by others and crop them to the visible part only.\n    Input segments are expected to be sorted from background to foreground.\n\n    Args:\n        segments: 1d list of segment RLEs (in COCO format)\n        width: width of the image\n        height: height of the image\n        iou_threshold: IoU threshold for objects to be counted as intersected\n            By default is set to 0 to process any intersected objects\n        ratio_tolerance: an IoU \"handicap\" value for a situation\n            when an object is (almost) fully covered by another one and we\n            don't want make a \"hole\" in the background object\n        area_threshold: minimal area of included segments\n\n    Returns:\n        A list of input segments' parts (in the same order as input):\n            [\n                [[x1,y1, x2,y2 ...], ...], # input segment #0 parts\n                mask1, # input segment #1 mask (if source segment is mask)\n                [], # when source segment is too small\n                ...\n            ]\n    \"\"\"\n    from pycocotools import mask as mask_utils\n\n    segments = [[s] for s in segments]\n    input_rles = [mask_utils.frPyObjects(s, height, width) for s in segments]\n\n    for i, rle_bottom in enumerate(input_rles):\n        area_bottom = sum(mask_utils.area(rle_bottom))\n        if area_bottom < area_threshold:\n            segments[i] = [] if not return_masks else None\n            continue\n\n        rles_top = []\n        for j in range(i + 1, len(input_rles)):\n            rle_top = input_rles[j]\n            iou = sum(mask_utils.iou(rle_bottom, rle_top, [0, 0]))[0]\n\n            if iou <= iou_threshold:\n                continue\n\n            area_top = sum(mask_utils.area(rle_top))\n            area_ratio = area_top / area_bottom\n\n            # If a segment is fully inside another one, skip this segment\n            if abs(area_ratio - iou) < ratio_tolerance:\n                continue\n\n            # Check if the bottom segment is fully covered by the top one.\n            # There is a mistake in the annotation, keep the background one\n            if abs(1 / area_ratio - iou) < ratio_tolerance:\n                rles_top = []\n                break\n\n            rles_top += rle_top\n\n        if not rles_top and not isinstance(segments[i][0], dict) \\\n                and not return_masks:\n            continue\n\n        rle_bottom = rle_bottom[0]\n        bottom_mask = mask_utils.decode(rle_bottom).astype(np.uint8)\n\n        if rles_top:\n            rle_top = mask_utils.merge(rles_top)\n            top_mask = mask_utils.decode(rle_top).astype(np.uint8)\n\n            bottom_mask -= top_mask\n            bottom_mask[bottom_mask != 1] = 0\n\n        if not return_masks and not isinstance(segments[i][0], dict):\n            segments[i] = mask_to_polygons(bottom_mask,\n                area_threshold=area_threshold)\n        else:\n            segments[i] = bottom_mask\n\n    return segments\n\ndef rles_to_mask(rles, width, height):\n    from pycocotools import mask as mask_utils\n\n    rles = mask_utils.frPyObjects(rles, height, width)\n    rles = mask_utils.merge(rles)\n    mask = mask_utils.decode(rles)\n    return mask\n\ndef find_mask_bbox(mask):\n    cols = np.any(mask, axis=0)\n    rows = np.any(mask, axis=1)\n    x0, x1 = np.where(cols)[0][[0, -1]]\n    y0, y1 = np.where(rows)[0][[0, -1]]\n    return [x0, y0, x1 - x0, y1 - y0]\n\ndef merge_masks(masks, start=None):\n    \"\"\"\n        Merges masks into one, mask order is responsible for z order.\n        To avoid memory explosion on mask materialization, consider passing\n        a generator.\n\n        Inputs: a sequence of index masks or (binary mask, index) pairs\n        Outputs: an index mask\n    \"\"\"\n    if start is not None:\n        masks = chain([start], masks)\n\n    it = iter(masks)\n\n    try:\n        merged_mask = next(it)\n        if isinstance(merged_mask, tuple) and len(merged_mask) == 2:\n            merged_mask = merged_mask[0] * merged_mask[1]\n    except StopIteration:\n        return None\n\n    for m in it:\n        if isinstance(m, tuple) and len(m) == 2:\n            merged_mask = np.where(m[0], m[1], merged_mask)\n        else:\n            merged_mask = np.where(m, m, merged_mask)\n\n    return merged_mask", "path": "datumaro/util/mask_tools.py"}]}