Skip to content

Commit

Permalink
check model size
Browse files Browse the repository at this point in the history
  • Loading branch information
cospectrum committed Dec 14, 2024
1 parent c5d01be commit a5c19cd
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 2 deletions.
2 changes: 0 additions & 2 deletions src/microwink/seg.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,6 @@ class SegModel:
def from_path(
path: str | os.PathLike, *, providers: Sequence[str] | None = None
) -> "SegModel":
if providers is None:
providers = ["CPUExecutionProvider"]
return SegModel.from_session(ort.InferenceSession(path, providers=providers))

@staticmethod
Expand Down
7 changes: 7 additions & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,17 @@
import pytest
import os

from pathlib import Path
from microwink import SegModel

KiB = 2**10
MiB = (2**10) * KiB


@pytest.fixture
def seg_model() -> SegModel:
path = Path("./models/seg_model.onnx")
assert path.exists()
size = os.path.getsize(path)
assert 10 * MiB < size < 15 * MiB
return SegModel.from_path(path)

0 comments on commit a5c19cd

Please sign in to comment.