From 54ecfceac7c9a3710a642f593ddc79e5ff18f3b3 Mon Sep 17 00:00:00 2001 From: Alex Severin Date: Sun, 5 Jan 2025 00:01:30 +0300 Subject: [PATCH 1/4] add preprocess proptest --- src/microwink/seg.py | 28 +++++++++++++++------------- tests/test_props.py | 21 ++++++++++++++++++--- 2 files changed, 33 insertions(+), 16 deletions(-) diff --git a/src/microwink/seg.py b/src/microwink/seg.py index fd8b5e7..3fb4e36 100644 --- a/src/microwink/seg.py +++ b/src/microwink/seg.py @@ -16,7 +16,6 @@ H = NewType("H", int) W = NewType("W", int) -RgbBuf = NewType("RgbBuf", np.ndarray) @dataclass @@ -84,8 +83,8 @@ def apply( ) -> list[SegResult]: CLASS_ID = 0 assert image.mode == "RGB" - buf = RgbBuf(np.array(image)) - raw = self._run(buf, threshold.confidence, threshold.iou) + + raw = self._run(image, threshold.confidence, threshold.iou) if raw is None: return [] @@ -105,17 +104,16 @@ def apply( return results def _run( - self, img: RgbBuf, conf_threshold: float, iou_threshold: float + self, image: PILImage, conf_threshold: float, iou_threshold: float ) -> RawResult | None: NM = 32 - ih, iw, _ = img.shape - - blob, ratio, (pad_w, pad_h) = self.preprocess(img) + assert image.mode == "RGB" + blob, ratio, (pad_w, pad_h) = self.preprocess(image) assert blob.ndim == 4 preds = self.session.run(None, {self.input_.name: blob}) return self.postprocess( preds, - img_size=(ih, iw), + img_size=(H(image.height), W(image.width)), ratio=ratio, pad_w=pad_w, pad_h=pad_h, @@ -125,15 +123,20 @@ def _run( ) def preprocess( - self, img_buf: RgbBuf + self, image: PILImage ) -> tuple[np.ndarray, float, tuple[float, float]]: BORDER_COLOR = (114, 114, 114) EPS = 0.1 - img = np.array(img_buf) + + assert image.mode == "RGB" + img = np.array(image) + ih, iw, _ = img.shape oh, ow = self.model_height, self.model_width r = min(oh / ih, ow / iw) rw, rh = round(iw * r), round(ih * r) + rw = max(1, rw) + rh = max(1, rh) pad_w, pad_h = [ (ow - rw) / 2, @@ -167,14 +170,13 @@ def postprocess( B = 1 NM, MH, MW = (nm, 160, 160) NUM_CLASSES = 1 - C = 4 + NUM_CLASSES + NM x, protos = preds assert len(x) == len(protos) == B protos = protos[0] x = x[0].T assert protos.shape == (NM, MH, MW), protos.shape - assert x.shape == (len(x), C) + assert x.shape == (len(x), 4 + NUM_CLASSES + NM) likely = x[:, 4 : 4 + NUM_CLASSES].max(axis=1) > conf_threshold x = x[likely] @@ -335,7 +337,7 @@ def with_border( bottom: int, left: int, right: int, - color: tuple[int, int, int], + color: Color, ) -> np.ndarray: import cv2 diff --git a/tests/test_props.py b/tests/test_props.py index b51d110..2826879 100644 --- a/tests/test_props.py +++ b/tests/test_props.py @@ -10,11 +10,26 @@ @settings( - deadline=1 * 1000, - max_examples=50, + max_examples=200, +) +@given(img=arb_img((1, 2000), (1, 2000))) +def test_preprocess(img: PILImage, seg_model: SegModel) -> None: + B = 1 + CH = 3 + H = seg_model.model_height + W = seg_model.model_width + + blob, *_ = seg_model.preprocess(img) + assert blob.shape == (B, CH, H, W) + assert blob.min() >= 0.0 + assert blob.max() <= 1.0 + + +@settings( + deadline=2 * 1000, ) @given( - img=arb_img((1, 1000), (1, 1000)), + img=arb_img((1, 2000), (1, 2000)), iou=st.none() | st.floats(0.01, 1.0), score=st.none() | st.floats(0.01, 1.0), ) From d7c0aedfa714263626d0a45e7e1d425581e1a58e Mon Sep 17 00:00:00 2001 From: Alex Severin Date: Sun, 5 Jan 2025 00:15:18 +0300 Subject: [PATCH 2/4] rm opencv img padding --- pyproject.toml | 1 - src/microwink/seg.py | 12 +++++++----- uv.lock | 19 ------------------- 3 files changed, 7 insertions(+), 25 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9a5cd44..fc01be7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,6 @@ requires-python = ">=3.10" dependencies = [ "numpy>=2.2.0", "onnxruntime>=1.20.1", - "opencv-python>=4.10.0.84", "pillow>=11.0.0", ] diff --git a/src/microwink/seg.py b/src/microwink/seg.py index 3fb4e36..8b19cef 100644 --- a/src/microwink/seg.py +++ b/src/microwink/seg.py @@ -179,6 +179,7 @@ def postprocess( assert x.shape == (len(x), 4 + NUM_CLASSES + NM) likely = x[:, 4 : 4 + NUM_CLASSES].max(axis=1) > conf_threshold + assert likely.ndim == 1 x = x[likely] scores = x[:, 4 : 4 + NUM_CLASSES].max(axis=1) @@ -339,12 +340,13 @@ def with_border( right: int, color: Color, ) -> np.ndarray: - import cv2 - assert img.ndim == 3 - return cv2.copyMakeBorder( - img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color - ) + pil_img = Image.fromarray(img) + ow = pil_img.width + left + right + oh = pil_img.height + top + bottom + out = Image.new("RGB", (ow, oh), color) + out.paste(pil_img, (left, top)) + return np.array(out).astype(img.dtype) def resize(buf: np.ndarray, size: tuple[W, H]) -> np.ndarray: diff --git a/uv.lock b/uv.lock index 2821f45..af8a76d 100644 --- a/uv.lock +++ b/uv.lock @@ -107,7 +107,6 @@ source = { editable = "." } dependencies = [ { name = "numpy" }, { name = "onnxruntime" }, - { name = "opencv-python" }, { name = "pillow" }, ] @@ -123,7 +122,6 @@ dev = [ requires-dist = [ { name = "numpy", specifier = ">=2.2.0" }, { name = "onnxruntime", specifier = ">=1.20.1" }, - { name = "opencv-python", specifier = ">=4.10.0.84" }, { name = "pillow", specifier = ">=11.0.0" }, ] @@ -285,23 +283,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/14/56/fd990ca222cef4f9f4a9400567b9a15b220dee2eafffb16b2adbc55c8281/onnxruntime-1.20.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0df6f2df83d61f46e842dbcde610ede27218947c33e994545a22333491e72a3b", size = 13337040 }, ] -[[package]] -name = "opencv-python" -version = "4.10.0.84" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/4a/e7/b70a2d9ab205110d715906fc8ec83fbb00404aeb3a37a0654fdb68eb0c8c/opencv-python-4.10.0.84.tar.gz", hash = "sha256:72d234e4582e9658ffea8e9cae5b63d488ad06994ef12d81dc303b17472f3526", size = 95103981 } -wheels = [ - { url = "https://files.pythonhosted.org/packages/66/82/564168a349148298aca281e342551404ef5521f33fba17b388ead0a84dc5/opencv_python-4.10.0.84-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:fc182f8f4cda51b45f01c64e4cbedfc2f00aff799debebc305d8d0210c43f251", size = 54835524 }, - { url = "https://files.pythonhosted.org/packages/64/4a/016cda9ad7cf18c58ba074628a4eaae8aa55f3fd06a266398cef8831a5b9/opencv_python-4.10.0.84-cp37-abi3-macosx_12_0_x86_64.whl", hash = "sha256:71e575744f1d23f79741450254660442785f45a0797212852ee5199ef12eed98", size = 56475426 }, - { url = "https://files.pythonhosted.org/packages/81/e4/7a987ebecfe5ceaf32db413b67ff18eb3092c598408862fff4d7cc3fd19b/opencv_python-4.10.0.84-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:09a332b50488e2dda866a6c5573ee192fe3583239fb26ff2f7f9ceb0bc119ea6", size = 41746971 }, - { url = "https://files.pythonhosted.org/packages/3f/a4/d2537f47fd7fcfba966bd806e3ec18e7ee1681056d4b0a9c8d983983e4d5/opencv_python-4.10.0.84-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9ace140fc6d647fbe1c692bcb2abce768973491222c067c131d80957c595b71f", size = 62548253 }, - { url = "https://files.pythonhosted.org/packages/1e/39/bbf57e7b9dab623e8773f6ff36385456b7ae7fa9357a5e53db732c347eac/opencv_python-4.10.0.84-cp37-abi3-win32.whl", hash = "sha256:2db02bb7e50b703f0a2d50c50ced72e95c574e1e5a0bb35a8a86d0b35c98c236", size = 28737688 }, - { url = "https://files.pythonhosted.org/packages/ec/6c/fab8113424af5049f85717e8e527ca3773299a3c6b02506e66436e19874f/opencv_python-4.10.0.84-cp37-abi3-win_amd64.whl", hash = "sha256:32dbbd94c26f611dc5cc6979e6b7aa1f55a64d6b463cc1dcd3c95505a63e48fe", size = 38842521 }, -] - [[package]] name = "packaging" version = "24.2" From 35b97bce6723fdceb82db7e3e82d5e78ac7998e3 Mon Sep 17 00:00:00 2001 From: Alex Severin Date: Sun, 5 Jan 2025 00:20:21 +0300 Subject: [PATCH 3/4] add empty check for input image --- src/microwink/seg.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/microwink/seg.py b/src/microwink/seg.py index 8b19cef..40a3717 100644 --- a/src/microwink/seg.py +++ b/src/microwink/seg.py @@ -83,6 +83,8 @@ def apply( ) -> list[SegResult]: CLASS_ID = 0 assert image.mode == "RGB" + assert image.width > 0 + assert image.height > 0 raw = self._run(image, threshold.confidence, threshold.iou) if raw is None: From 68544a1de709bcc939d33605f551fc3088a58e0b Mon Sep 17 00:00:00 2001 From: Alex Severin Date: Sun, 5 Jan 2025 00:26:12 +0300 Subject: [PATCH 4/4] add threshold asserts --- src/microwink/seg.py | 5 ++++- tests/test_samples.py | 3 +++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/microwink/seg.py b/src/microwink/seg.py index 40a3717..8115900 100644 --- a/src/microwink/seg.py +++ b/src/microwink/seg.py @@ -86,7 +86,10 @@ def apply( assert image.width > 0 assert image.height > 0 - raw = self._run(image, threshold.confidence, threshold.iou) + assert 0.0 <= threshold.iou <= 1.0 + assert 0.0 <= threshold.confidence <= 1.0 + + raw = self._run(image, threshold.confidence, iou_threshold=threshold.iou) if raw is None: return [] diff --git a/tests/test_samples.py b/tests/test_samples.py index ef5cfd8..acfe6b1 100644 --- a/tests/test_samples.py +++ b/tests/test_samples.py @@ -49,7 +49,10 @@ def test_samples( actual = img.copy() for card, box in zip(cards, boxes): assert 0.0 < card.score < 1.0 + assert card.mask.min() >= 0.0 + assert card.mask.max() <= 1.0 assert round_box(card.box) == box + actual = draw_box(actual, card.box) actual = draw_mask(actual, card.mask > BIN_THRESHOLD) assert truth == actual