Merge pull request #10 from cospectrum/dev

remove opencv

Author: cospectrum

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",
 ]
 

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,13 @@ def apply(
     ) -> list[SegResult]:
         CLASS_ID = 0
         assert image.mode == "RGB"
-        buf = RgbBuf(np.array(image))
-        raw = self._run(buf, threshold.confidence, threshold.iou)
+        assert image.width > 0
+        assert image.height > 0
+
+        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 []
 
@@ -105,17 +109,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 +128,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,16 +175,16 @@ 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
+        assert likely.ndim == 1
         x = x[likely]
 
         scores = x[:, 4 : 4 + NUM_CLASSES].max(axis=1)
@@ -335,14 +343,15 @@ def with_border(
         bottom: int,
         left: int,
         right: int,
-        color: tuple[int, int, 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:

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),
 )

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