✨ Feat: Mask by Poly

sphinx/example_images.py

@@ -235,10 +235,15 @@ def example_vert():
 # vision.poly2mask.png
 def example_poly2mask():
     poly = [[10, 10], [20, 10], [30, 40], [20, 60], [10, 20]]
-    mask = zz.vision.poly2mask(poly, (70, 100))
-    cv2.imwrite(
-        f"{EXAMPLE_PATH}/vision.poly2mask.png",
-        mask.astype(np.uint8) * 255,
+    mask1 = zz.vision.poly2mask(poly, (70, 100))
+    poly = np.array(poly)
+    mask2 = zz.vision.poly2mask([poly, poly - 10, poly + 20], (70, 100))
+    zz.vision.vert(
+        [
+            mask1.astype(np.uint8) * 255,
+            np.transpose(mask2, (1, 2, 0)).astype(np.uint8) * 255,
+        ],
+        file_name=f"{EXAMPLE_PATH}/vision.vert",
     )
 
 
@@ -306,14 +311,20 @@ def example_mask():
     mks = mks.astype(bool)
     res1 = zz.vision.mask(IMAGE, mks)
     cls = [i for i in range(cnt)]
-    class_list = [cls[random.randint(0, 2)] for _ in range(cnt)]
+    class_list = [cls[random.randint(0, 5)] for _ in range(cnt)]
     class_color = {}
     for c in cls:
         class_color[c] = [random.randint(0, 255) for _ in range(3)]
     res2 = zz.vision.mask(IMAGE, mks, class_list=class_list, class_color=class_color)
     poly = np.array([[100, 400], [400, 400], [800, 900], [400, 1100], [100, 800]])
     res3 = zz.vision.mask(IMAGE, poly=poly)
-    zz.vision.vert([res1, res2, res3], file_name=f"{EXAMPLE_PATH}/vision.mask")
+    poly = zz.vision.xyxy2poly(
+        zz.vision.poly2xyxy((np.random.rand(cnt, 4, 2) * (W, H)))
+    )
+    res4 = zz.vision.mask(
+        IMAGE, poly=poly, class_list=class_list, class_color=class_color
+    )
+    zz.vision.vert([res1, res2, res3, res4], file_name=f"{EXAMPLE_PATH}/vision.mask")
 
 
 # vision.text.png

test/test_vision.py

@@ -224,6 +224,21 @@ def test_mask_bgra():
     assert "MASK_BGRA.png" in os.listdir()
 
 
+def test_mask_poly():
+    img = cv2.imread(f"{data}/test.jpg")
+    H, W, _ = img.shape
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
+    cnt = 30
+    poly = zz.vision.xyxy2poly(
+        zz.vision.poly2xyxy((np.random.rand(cnt, 4, 2) * (W, H)))
+    )
+    BGRA = zz.vision.mask(
+        img, poly=poly, color=[random.randint(0, 255) for _ in range(3)]
+    )
+    cv2.imwrite("MASK_POLY.png", BGRA)
+    assert "MASK_POLY.png" in os.listdir()
+
+
 def test_mask_gray_int():
     img = cv2.imread(f"{data}/test.jpg")
     H, W, _ = img.shape
@@ -237,7 +252,7 @@ def test_mask_gray_int():
     mks = mks.astype(bool)
     img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
     cls = [i for i in range(cnt)]
-    class_list = [cls[random.randint(0, 2)] for _ in range(cnt)]
+    class_list = [cls[random.randint(0, 5)] for _ in range(cnt)]
     class_color = {}
     for c in cls:
         class_color[c] = [random.randint(0, 255) for _ in range(3)]

zerohertzLib/vision/convert.py

@@ -283,38 +283,57 @@ def poly2xyxy(
     return _poly2xyxy(box)
 
 
+def _poly2mask(poly: NDArray[DTypeLike], shape: Tuple[int]) -> NDArray[bool]:
+    poly = Path(poly)
+    pts_x, pts_y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
+    pts_x, pts_y = pts_x.flatten(), pts_y.flatten()
+    points = np.vstack((pts_x, pts_y)).T
+    grid = poly.contains_points(points)
+    mask = grid.reshape(shape)
+    return mask
+
+
 def poly2mask(
-    poly: Union[List[Union[int, float]], NDArray[DTypeLike]], shape: Tuple[int]
+    poly: Union[List[Union[int, float]], NDArray[DTypeLike], List[NDArray[DTypeLike]]],
+    shape: Tuple[int],
 ) -> NDArray[bool]:
     """다각형 좌표를 입력받아 mask로 변환
 
     Args:
-        poly (``Union[List[Union[int, float]], NDArray[DTypeLike]]``): Mask의 꼭짓점 좌표 (``[N, 2]``)
+        poly (``Union[List[Union[int, float]], NDArray[DTypeLike], List[NDArray[DTypeLike]]]``): Mask의 꼭짓점 좌표 (``[M, 2]`` or ``[N, M, 2]``)
         shape (``Tuple[int]``): 출력될 mask의 shape ``(H, W)``
 
     Returns:
-        ``NDArray[bool]``: 시각화 결과 (``[H, W, C]``)
+        ``NDArray[bool]``: 변환된 mask (``[H, W]`` or ``[N, H, W]``)
 
     Examples:
         >>> poly = [[10, 10], [20, 10], [30, 40], [20, 60], [10, 20]]
-        >>> mask = zz.vision.poly2mask(poly, (70, 100))
-        >>> mask.shape
+        >>> mask1 = zz.vision.poly2mask(poly, (70, 100))
+        >>> mask1.shape
         (70, 100)
-        >>> mask.dtype
+        >>> mask1.dtype
+        dtype('bool')
+        >>> poly = np.array(poly)
+        >>> mask2 = zz.vision.poly2mask([poly, poly - 10, poly + 20], (70, 100))
+        >>> mask2.shape
+        (3, 70, 100)
+        >>> mask2.dtype
         dtype('bool')
 
         .. image:: _static/examples/dynamic/vision.poly2mask.png
             :align: center
             :width: 300px
     """
-    poly = _list2np(poly)
-    poly = Path(poly)
-    pts_x, pts_y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
-    pts_x, pts_y = pts_x.flatten(), pts_y.flatten()
-    points = np.vstack((pts_x, pts_y)).T
-    grid = poly.contains_points(points)
-    mask = grid.reshape(shape)
-    return mask
+    if (isinstance(poly, list) and isinstance(poly[0], np.ndarray)) or (
+        isinstance(poly, np.ndarray) and len(poly.shape) == 3
+    ):
+        mks = []
+        for _poly in poly:
+            mks.append(_poly2mask(_poly, shape))
+        mks = np.array(mks)
+    else:
+        mks = _poly2mask(_list2np(poly), shape)
+    return mks
 
 
 def poly2area(poly: Union[List[Union[int, float]], NDArray[DTypeLike]]) -> float:

zerohertzLib/vision/visual.py

@@ -119,7 +119,9 @@ def bbox(
 def mask(
     img: NDArray[np.uint8],
     mks: Optional[NDArray[bool]] = None,
-    poly: Optional[NDArray[DTypeLike]] = None,
+    poly: Optional[
+        Union[List[Union[int, float]], NDArray[DTypeLike], List[NDArray[DTypeLike]]]
+    ] = None,
     color: Optional[Tuple[int]] = (0, 0, 255),
     class_list: Optional[List[Union[int, str]]] = None,
     class_color: Optional[Dict[Union[int, str], Tuple[int]]] = None,
@@ -131,7 +133,7 @@ def mask(
     Args:
         img (``NDArray[np.uint8]``): 입력 image (``[H, W, C]``)
         mks (``Optional[NDArray[bool]]``): 입력 image 위에 병합할 mask (``[H, W]`` or ``[N, H, W]``)
-        poly (``Optional[NDArray[DTypeLike]]``): 입력 image 위에 병합할 mask (``[N, 2]``)
+        poly (``Optional[Union[List[Union[int, float]], NDArray[DTypeLike], List[NDArray[DTypeLike]]]]``): 입력 image 위에 병합할 mask (``[M, 2]`` or ``[N, M, 2]``)
         color (``Optional[Tuple[int]]``): Mask의 색
         class_list (``Optional[List[Union[int, str]]]``): ``mks`` 의 index에 따른 class
         class_color (``Optional[Dict[Union[int, str], Tuple[int]]]``): Class에 따른 색 (``color`` 무시)
@@ -156,20 +158,25 @@ def mask(
 
         Mask (with class):
             >>> cls = [i for i in range(cnt)]
-            >>> class_list = [cls[random.randint(0, 2)] for _ in range(cnt)]
+            >>> class_list = [cls[random.randint(0, 5)] for _ in range(cnt)]
             >>> class_color = {}
             >>> for c in cls:
             >>>     class_color[c] = [random.randint(0, 255) for _ in range(3)]
             >>> res2 = zz.vision.mask(img, mks, class_list=class_list, class_color=class_color)
 
-        Poly:
+        Poly (without class):
             >>> poly = np.array([[100, 400], [400, 400], [800, 900], [400, 1100], [100, 800]])
             >>> res3 = zz.vision.mask(img, poly=poly)
 
+        Poly (with class):
+            >>> poly = zz.vision.xyxy2poly(zz.vision.poly2xyxy((np.random.rand(cnt, 4, 2) * (W, H))))
+            >>> res4 = zz.vision.mask(img, poly=poly, class_list=class_list, class_color=class_color)
+
         .. image:: _static/examples/dynamic/vision.mask.png
             :align: center
             :width: 600px
     """
+    assert (mks is None) ^ (poly is None)
     shape = img.shape
     if len(shape) == 2:
         img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)