Skip to content

Commit 77908fe

Browse files
authored
Merge pull request roboflow#1012 from roboflow/feat/type-check-improvement
feat: 📝 static type improvement color constants and PolygonZone
2 parents fb3a78f + 63b3387 commit 77908fe

File tree

2 files changed

+12
-11
lines changed

2 files changed

+12
-11
lines changed

supervision/detection/tools/polygon_zone.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
import cv2
55
import numpy as np
6+
import numpy.typing as npt
67

78
from supervision import Detections
89
from supervision.detection.utils import clip_boxes, polygon_to_mask
@@ -39,7 +40,7 @@ class PolygonZone:
3940
)
4041
def __init__(
4142
self,
42-
polygon: np.ndarray,
43+
polygon: npt.NDArray[np.int64],
4344
frame_resolution_wh: Tuple[int, int],
4445
triggering_anchors: Iterable[Position] = (Position.BOTTOM_CENTER,),
4546
):
@@ -54,7 +55,7 @@ def __init__(
5455
polygon=polygon, resolution_wh=(width + 1, height + 1)
5556
)
5657

57-
def trigger(self, detections: Detections) -> np.ndarray:
58+
def trigger(self, detections: Detections) -> npt.NDArray[np.bool_]:
5859
"""
5960
Determines if the detections are within the polygon zone.
6061
@@ -78,13 +79,13 @@ def trigger(self, detections: Detections) -> np.ndarray:
7879
]
7980
)
8081

81-
is_in_zone = (
82+
is_in_zone: npt.NDArray[np.bool_] = (
8283
self.mask[all_clipped_anchors[:, :, 1], all_clipped_anchors[:, :, 0]]
8384
.transpose()
8485
.astype(bool)
8586
)
86-
is_in_zone = np.all(is_in_zone, axis=1)
8787

88+
is_in_zone: npt.NDArray[np.bool_] = np.all(is_in_zone, axis=1)
8889
self.current_count = int(np.sum(is_in_zone))
8990
return is_in_zone.astype(bool)
9091

supervision/draw/color.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -176,31 +176,31 @@ def as_bgr(self) -> Tuple[int, int, int]:
176176
return self.b, self.g, self.r
177177

178178
@classproperty
179-
def WHITE(cls):
179+
def WHITE(cls) -> Color:
180180
return Color.from_hex("#FFFFFF")
181181

182182
@classproperty
183-
def BLACK(cls):
183+
def BLACK(cls) -> Color:
184184
return Color.from_hex("#000000")
185185

186186
@classproperty
187-
def RED(cls):
187+
def RED(cls) -> Color:
188188
return Color.from_hex("#FF0000")
189189

190190
@classproperty
191-
def GREEN(cls):
191+
def GREEN(cls) -> Color:
192192
return Color.from_hex("#00FF00")
193193

194194
@classproperty
195-
def BLUE(cls):
195+
def BLUE(cls) -> Color:
196196
return Color.from_hex("#0000FF")
197197

198198
@classproperty
199-
def YELLOW(cls):
199+
def YELLOW(cls) -> Color:
200200
return Color.from_hex("#FFFF00")
201201

202202
@classproperty
203-
def ROBOFLOW(cls):
203+
def ROBOFLOW(cls) -> Color:
204204
return Color.from_hex("#A351FB")
205205

206206
@classmethod

0 commit comments

Comments
 (0)