Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

rm opencv dependency #3

Merged
merged 2 commits into from
Dec 15, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -10,3 +10,5 @@ wheels/
.venv

.DS_Store

draw.py
4 changes: 2 additions & 2 deletions assets/us_card.result.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
5 changes: 2 additions & 3 deletions pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,15 +1,14 @@
[project]
name = "microwink"
version = "0.0.1"
description = "Lightweight instance segmentation for card IDs"
version = "0.0.2"
description = "Lightweight instance segmentation of card IDs"
readme = "README.md"
license = { text = "Apache-2.0" }
authors = [{ name = "cospectrum", email = "severinalexeyv@gmail.com" }]
requires-python = ">=3.10"
dependencies = [
"numpy>=2.2.0",
"onnxruntime>=1.20.1",
"opencv-python>=4.10.0.84",
"pillow>=11.0.0",
]

Expand Down
2 changes: 1 addition & 1 deletion src/microwink/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ def draw_box(
box: Box,
*,
color: tuple[int, ...] | str | float = (255, 0, 0),
width: int = 4,
width: int = 3,
) -> PILImage:
image = image.copy()
draw = ImageDraw.Draw(image)
Expand Down
74 changes: 44 additions & 30 deletions src/microwink/seg.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
import math
import os
import math
import onnxruntime as ort # type: ignore # missing stubs
import numpy as np
import cv2 as cv

from typing import Sequence
from dataclasses import dataclass
from PIL import Image
from PIL.Image import Image as PILImage

from . import common
Expand Down Expand Up @@ -79,7 +79,9 @@ def from_session(session: ort.InferenceSession) -> "SegModel":
)

def apply(
self, image: PILImage, threshold: Threshold = Threshold.default()
self,
image: PILImage,
threshold: Threshold = Threshold.default(),
) -> list[SegResult]:
assert image.mode == "RGB"
tensor = self.preprocess(image)
Expand All @@ -104,7 +106,10 @@ def apply(
return out

def postprocess(
self, outs: list[np.ndarray], img_size: tuple[H, W], threshold: Threshold
self,
outs: list[np.ndarray],
img_size: tuple[H, W],
threshold: Threshold,
) -> Result | None:
NUM_MASKS = 32
box_out, mask_out = outs
Expand All @@ -130,7 +135,10 @@ def postprocess(
final_boxes = boxes[indexes]
final_scores = scores[indexes]
final_mask_maps = self.postprocess_mask(
mask_preds[indexes], mask_out, final_boxes, img_size
mask_preds[indexes],
mask_out,
final_boxes,
img_size,
)
assert len(final_boxes) == len(final_scores) == len(final_mask_maps)
return Result(
Expand Down Expand Up @@ -165,33 +173,30 @@ def postprocess_mask(
iw,
)
)
blur_size = (
int(iw / mask_width),
int(ih / mask_height),
)
for i in range(len(scaled_boxes)):
scaled_box: np.ndarray = scaled_boxes[i]
box: np.ndarray = boxes[i]
scale_x1 = int(math.floor(scaled_box[0]))
scale_y1 = int(math.floor(scaled_box[1]))
scale_x2 = int(math.ceil(scaled_box[2]))
scale_y2 = int(math.ceil(scaled_box[3]))

x1 = int(math.floor(box[0]))
y1 = int(math.floor(box[1]))
x2 = int(math.ceil(box[2]))
y2 = int(math.ceil(box[3]))

mask: np.ndarray = masks[i]
assert len(scaled_boxes) == len(masks)
assert len(scaled_boxes) == len(boxes)
for i, (box, scaled_box, mask) in enumerate(zip(boxes, scaled_boxes, masks)):
assert 2 == len(mask.shape)
final_mask = cv.resize(

scale_x1 = math.floor(scaled_box[0])
scale_y1 = math.floor(scaled_box[1])
scale_x2 = math.ceil(scaled_box[2])
scale_y2 = math.ceil(scaled_box[3])

x1 = math.floor(box[0])
y1 = math.floor(box[1])
x2 = math.ceil(box[2])
y2 = math.ceil(box[3])

ow, oh = (x2 - x1, y2 - y1)
assert ow >= 0
assert oh >= 0
resized_mask = resize(
mask[scale_y1:scale_y2, scale_x1:scale_x2],
(x2 - x1, y2 - y1),
interpolation=cv.INTER_CUBIC,
(ow, oh),
)
final_mask = cv.blur(final_mask, blur_size) # -> [-inf, +inf]
final_mask = common.sigmoid(final_mask).clip(0.0, 1.0)
mask_maps[i, y1:y2, x1:x2] = final_mask
assert resized_mask.shape == (oh, ow)
mask_maps[i, y1:y2, x1:x2] = common.sigmoid(resized_mask).clip(0.0, 1.0)

return mask_maps

Expand All @@ -204,7 +209,8 @@ def forward(self, tensor: np.ndarray) -> list[np.ndarray]:

def preprocess(self, image: PILImage) -> np.ndarray:
size = (self.input_shape.w, self.input_shape.h)
image = image.resize(size)
if image.size != size:
image = image.resize(size)
img = np.array(image).astype(np.float32)
assert len(img.shape) == 3
img /= 255.0
Expand Down Expand Up @@ -295,3 +301,11 @@ def compute_iou(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:

iou = intersection_area / union_area
return iou


def resize(buf: np.ndarray, size: tuple[W, H]) -> np.ndarray:
img = Image.fromarray(buf).resize(size)
out = np.array(img)
assert out.dtype == buf.dtype
assert len(out.shape) == len(buf.shape)
return out
4 changes: 2 additions & 2 deletions tests/truth/seg_model/mklovin.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
4 changes: 2 additions & 2 deletions tests/truth/seg_model/us_card.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
21 changes: 1 addition & 20 deletions uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading