siglip2-toolkit 0.0.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (36) hide show
  1. siglip2_toolkit-0.0.1/.gitignore +20 -0
  2. siglip2_toolkit-0.0.1/LICENSE +21 -0
  3. siglip2_toolkit-0.0.1/PKG-INFO +153 -0
  4. siglip2_toolkit-0.0.1/README.md +114 -0
  5. siglip2_toolkit-0.0.1/pyproject.toml +82 -0
  6. siglip2_toolkit-0.0.1/src/siglip2_toolkit/__init__.py +71 -0
  7. siglip2_toolkit-0.0.1/src/siglip2_toolkit/__main__.py +6 -0
  8. siglip2_toolkit-0.0.1/src/siglip2_toolkit/classification.py +161 -0
  9. siglip2_toolkit-0.0.1/src/siglip2_toolkit/cli.py +557 -0
  10. siglip2_toolkit-0.0.1/src/siglip2_toolkit/defaults.py +3 -0
  11. siglip2_toolkit-0.0.1/src/siglip2_toolkit/embedder.py +324 -0
  12. siglip2_toolkit-0.0.1/src/siglip2_toolkit/exporter.py +272 -0
  13. siglip2_toolkit-0.0.1/src/siglip2_toolkit/filters.py +283 -0
  14. siglip2_toolkit-0.0.1/src/siglip2_toolkit/formats.py +293 -0
  15. siglip2_toolkit-0.0.1/src/siglip2_toolkit/models.py +421 -0
  16. siglip2_toolkit-0.0.1/src/siglip2_toolkit/retrieval.py +155 -0
  17. siglip2_toolkit-0.0.1/src/siglip2_toolkit/server.py +197 -0
  18. siglip2_toolkit-0.0.1/src/siglip2_toolkit/utils.py +245 -0
  19. siglip2_toolkit-0.0.1/tests/conftest.py +80 -0
  20. siglip2_toolkit-0.0.1/tests/data/SOURCES.md +22 -0
  21. siglip2_toolkit-0.0.1/tests/data/images/astronaut.jpg +0 -0
  22. siglip2_toolkit-0.0.1/tests/data/images/camera.jpg +0 -0
  23. siglip2_toolkit-0.0.1/tests/data/images/cat.jpg +0 -0
  24. siglip2_toolkit-0.0.1/tests/data/images/coffee.jpg +0 -0
  25. siglip2_toolkit-0.0.1/tests/data/images/grace_hopper.jpg +0 -0
  26. siglip2_toolkit-0.0.1/tests/data/images/horse.jpg +0 -0
  27. siglip2_toolkit-0.0.1/tests/data/images/hubble.jpg +0 -0
  28. siglip2_toolkit-0.0.1/tests/data/images/rocket.jpg +0 -0
  29. siglip2_toolkit-0.0.1/tests/fakes.py +66 -0
  30. siglip2_toolkit-0.0.1/tests/test_exporter.py +222 -0
  31. siglip2_toolkit-0.0.1/tests/test_filters.py +104 -0
  32. siglip2_toolkit-0.0.1/tests/test_formats.py +118 -0
  33. siglip2_toolkit-0.0.1/tests/test_models.py +192 -0
  34. siglip2_toolkit-0.0.1/tests/test_retrieval.py +71 -0
  35. siglip2_toolkit-0.0.1/tests/test_semantic.py +149 -0
  36. siglip2_toolkit-0.0.1/tests/test_utils.py +83 -0
@@ -0,0 +1,20 @@
1
+ # Python
2
+ __pycache__/
3
+ *.py[cod]
4
+ *.egg-info/
5
+ dist/
6
+ build/
7
+ .venv/
8
+ venv/
9
+
10
+ # Tooling
11
+ .pytest_cache/
12
+ .ruff_cache/
13
+ .mypy_cache/
14
+ .coverage
15
+ coverage.xml
16
+
17
+ # Editors / OS
18
+ .idea/
19
+ .vscode/
20
+ .DS_Store
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 GIA-UH (Universidad de La Habana) and Syalia S.R.L.
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,153 @@
1
+ Metadata-Version: 2.5
2
+ Name: siglip2-toolkit
3
+ Version: 0.0.1
4
+ Summary: Local-first toolkit around Google's SigLIP 2 models: text/image embeddings, zero-shot and fine-tuned classification, incremental corpus export and retrieval math. Library-first, with CLI and REST interfaces.
5
+ Project-URL: Homepage, https://github.com/gia-uh/siglip2-toolkit
6
+ Project-URL: Repository, https://github.com/gia-uh/siglip2-toolkit
7
+ Project-URL: Issues, https://github.com/gia-uh/siglip2-toolkit/issues
8
+ Author-email: Yudivián Almeida Cruz <yudivian@gmail.com>
9
+ License-Expression: MIT
10
+ License-File: LICENSE
11
+ Keywords: embeddings,image-classification,image-retrieval,image-search,local-first,siglip2
12
+ Classifier: Development Status :: 2 - Pre-Alpha
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Operating System :: OS Independent
16
+ Classifier: Programming Language :: Python :: 3
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Classifier: Topic :: Multimedia :: Graphics
20
+ Classifier: Topic :: Scientific/Engineering :: Image Recognition
21
+ Requires-Python: >=3.12
22
+ Requires-Dist: huggingface-hub>=0.23
23
+ Requires-Dist: numpy>=1.26
24
+ Requires-Dist: pillow-heif>=0.16
25
+ Requires-Dist: pillow>=10.0
26
+ Requires-Dist: pyyaml>=6.0
27
+ Requires-Dist: torch>=2.2
28
+ Requires-Dist: tqdm>=4.60
29
+ Requires-Dist: transformers>=4.51
30
+ Provides-Extra: dev
31
+ Requires-Dist: pytest-cov>=5; extra == 'dev'
32
+ Requires-Dist: pytest>=8; extra == 'dev'
33
+ Requires-Dist: ruff>=0.6; extra == 'dev'
34
+ Provides-Extra: server
35
+ Requires-Dist: fastapi>=0.115; extra == 'server'
36
+ Requires-Dist: python-multipart>=0.0.9; extra == 'server'
37
+ Requires-Dist: uvicorn[standard]>=0.30; extra == 'server'
38
+ Description-Content-Type: text/markdown
39
+
40
+ # siglip2-toolkit
41
+
42
+ Local-first toolkit around Google's **SigLIP 2** models: text/image embeddings, zero-shot and fine-tuned classification, incremental corpus export, and retrieval math.
43
+
44
+ It is **not a vector database**: it produces embeddings and portable exports, and stays agnostic about where you store or search your vectors.
45
+
46
+ Co-developed by [GIA-UH](https://github.com/gia-uh) and [Syalia S.R.L.](https://github.com/syalia-srl).
47
+
48
+ ## Installation
49
+
50
+ ```bash
51
+ pip install siglip2-toolkit # library + CLI
52
+ pip install "siglip2-toolkit[server]" # + REST server (FastAPI)
53
+ ```
54
+
55
+ ## Model management (the toolkit's own cache)
56
+
57
+ Models live in the toolkit's **own cache**, independent from HuggingFace's
58
+ cache. The only command that uses the network is `models download`;
59
+ everything else works with what you already have locally.
60
+
61
+ ```bash
62
+ siglip2 models # your cache (★ model in use, origin, size)
63
+ siglip2 models --hf # SigLIP-2 catalog on HuggingFace: ✓ have it / ☁ downloadable
64
+ siglip2 models current # the model in use (+ path, origin; --json)
65
+ siglip2 models use <id> # set the model in use (persistent)
66
+ siglip2 models download <org/name> [--revision REV] # into the toolkit cache (network)
67
+ siglip2 models import <path> [--name ALIAS] # copy a local model into the cache
68
+ siglip2 models delete <id> # delete from the cache (refuses the model in use)
69
+ ```
70
+
71
+ Configuration lives in `<home>/config.json` (default home
72
+ `~/.local/share/siglip2-toolkit/`, created on first use):
73
+
74
+ ```json
75
+ {
76
+ "cache_dir": "~/.local/share/siglip2-toolkit/models",
77
+ "current": "google/siglip2-base-patch16-naflex",
78
+ "device": "auto"
79
+ }
80
+ ```
81
+
82
+ - **`cache_dir`**: where the cache lives; change it to relocate it (e.g. another disk).
83
+ - **`current`**: the model in use.
84
+ - **`device`**: `auto` (CUDA → MPS → CPU), or `cuda`, `cpu`, `mps`, `cuda:1`, …
85
+ Any command's `--device` overrides it for that run.
86
+
87
+ There are **no environment variables**; for non-default installations pass
88
+ `--home DIR` (global option, before the command).
89
+
90
+ `--model` accepts either an identifier from your cache or a path to a local
91
+ model directory. The implicit default is
92
+ `google/siglip2-base-patch16-naflex`: on a machine that does not have it yet,
93
+ the first command tells you the exact `models download` to run. The toolkit
94
+ never picks a model for you.
95
+
96
+ Incremental exports refuse to run when the existing export was produced by a
97
+ different model: mixing embedding spaces would silently corrupt the vectors.
98
+
99
+ ## Three ways to use it
100
+
101
+ **1. As a library (primary interface)**
102
+
103
+ ```python
104
+ from siglip2_toolkit import Embedder, models
105
+
106
+ # Resolve from your cache (download/import first) or pass a local path.
107
+ resolved = models.resolve_model("google/siglip2-base-patch16-naflex")
108
+ embedder = Embedder(model=resolved.path, model_id=resolved.identifier, device="auto")
109
+
110
+ vectors = embedder.embed_texts(["a wedding at sunset"])
111
+ results = embedder.classify(["photo.jpg"], labels=["beach", "mountain", "wedding"])
112
+
113
+ from siglip2_toolkit.filters import field, matches
114
+
115
+ matches({"width": 2400}, field("width") >= 2000) # -> True
116
+ ```
117
+
118
+ **2. As a CLI**
119
+
120
+ ```bash
121
+ siglip2 models download google/siglip2-base-patch16-naflex
122
+ siglip2 models use google/siglip2-base-patch16-naflex
123
+ siglip2 export -i /path/to/photos -o embeddings.jsonl --incremental
124
+ siglip2 classify --image photo.jpg --labels "beach,mountain,wedding"
125
+ ```
126
+
127
+ **3. As a REST service (optional extra)**
128
+
129
+ ```bash
130
+ pip install "siglip2-toolkit[server]"
131
+ siglip2 serve --host 0.0.0.0 --port 8000
132
+ ```
133
+
134
+ Endpoints: `/v1/embed/text`, `/v1/embed/image`, `/v1/classify`, `/v1/health`, `/v1/info`.
135
+
136
+ ## Highlights
137
+
138
+ - **Library-first design**: the CLI and the REST server are thin adapters over the same public API.
139
+ - **Zero-shot classification**: prompt-based labels using the model's calibrated SigLIP logits, plus fine-tuned classification heads (`ImageClassifier`).
140
+ - **Own model cache**: downloads and imports live in the toolkit cache (independent from HuggingFace's); only `models download` uses the network.
141
+ - **GPU-ready**: `device: "auto"` picks CUDA when available (FP16; FP32 on CPU); override per command with `--device`.
142
+ - **Incremental export**: new, modified and deleted files are detected (modification detection requires metadata).
143
+ - **Resumable runs**: interrupted exports continue where they left off.
144
+ - **Optional metadata**: rich file metadata (EXIF, dimensions, timestamps, thumbnails) — or ids and vectors only, your choice.
145
+ - **Versioned export format** (`jsonl` / `pkl`) with an embedding-space manifest.
146
+ - **Retrieval math**: top-k scoring, RRF fusion, similarity-drop filtering and serializable metadata filters, composable with any store.
147
+
148
+ > **Status:** first preview (0.0.1, October 2026). The public API and the
149
+ export format may still change before 1.0; pin the version you depend on.
150
+
151
+ ## License
152
+
153
+ MIT — see [LICENSE](LICENSE).
@@ -0,0 +1,114 @@
1
+ # siglip2-toolkit
2
+
3
+ Local-first toolkit around Google's **SigLIP 2** models: text/image embeddings, zero-shot and fine-tuned classification, incremental corpus export, and retrieval math.
4
+
5
+ It is **not a vector database**: it produces embeddings and portable exports, and stays agnostic about where you store or search your vectors.
6
+
7
+ Co-developed by [GIA-UH](https://github.com/gia-uh) and [Syalia S.R.L.](https://github.com/syalia-srl).
8
+
9
+ ## Installation
10
+
11
+ ```bash
12
+ pip install siglip2-toolkit # library + CLI
13
+ pip install "siglip2-toolkit[server]" # + REST server (FastAPI)
14
+ ```
15
+
16
+ ## Model management (the toolkit's own cache)
17
+
18
+ Models live in the toolkit's **own cache**, independent from HuggingFace's
19
+ cache. The only command that uses the network is `models download`;
20
+ everything else works with what you already have locally.
21
+
22
+ ```bash
23
+ siglip2 models # your cache (★ model in use, origin, size)
24
+ siglip2 models --hf # SigLIP-2 catalog on HuggingFace: ✓ have it / ☁ downloadable
25
+ siglip2 models current # the model in use (+ path, origin; --json)
26
+ siglip2 models use <id> # set the model in use (persistent)
27
+ siglip2 models download <org/name> [--revision REV] # into the toolkit cache (network)
28
+ siglip2 models import <path> [--name ALIAS] # copy a local model into the cache
29
+ siglip2 models delete <id> # delete from the cache (refuses the model in use)
30
+ ```
31
+
32
+ Configuration lives in `<home>/config.json` (default home
33
+ `~/.local/share/siglip2-toolkit/`, created on first use):
34
+
35
+ ```json
36
+ {
37
+ "cache_dir": "~/.local/share/siglip2-toolkit/models",
38
+ "current": "google/siglip2-base-patch16-naflex",
39
+ "device": "auto"
40
+ }
41
+ ```
42
+
43
+ - **`cache_dir`**: where the cache lives; change it to relocate it (e.g. another disk).
44
+ - **`current`**: the model in use.
45
+ - **`device`**: `auto` (CUDA → MPS → CPU), or `cuda`, `cpu`, `mps`, `cuda:1`, …
46
+ Any command's `--device` overrides it for that run.
47
+
48
+ There are **no environment variables**; for non-default installations pass
49
+ `--home DIR` (global option, before the command).
50
+
51
+ `--model` accepts either an identifier from your cache or a path to a local
52
+ model directory. The implicit default is
53
+ `google/siglip2-base-patch16-naflex`: on a machine that does not have it yet,
54
+ the first command tells you the exact `models download` to run. The toolkit
55
+ never picks a model for you.
56
+
57
+ Incremental exports refuse to run when the existing export was produced by a
58
+ different model: mixing embedding spaces would silently corrupt the vectors.
59
+
60
+ ## Three ways to use it
61
+
62
+ **1. As a library (primary interface)**
63
+
64
+ ```python
65
+ from siglip2_toolkit import Embedder, models
66
+
67
+ # Resolve from your cache (download/import first) or pass a local path.
68
+ resolved = models.resolve_model("google/siglip2-base-patch16-naflex")
69
+ embedder = Embedder(model=resolved.path, model_id=resolved.identifier, device="auto")
70
+
71
+ vectors = embedder.embed_texts(["a wedding at sunset"])
72
+ results = embedder.classify(["photo.jpg"], labels=["beach", "mountain", "wedding"])
73
+
74
+ from siglip2_toolkit.filters import field, matches
75
+
76
+ matches({"width": 2400}, field("width") >= 2000) # -> True
77
+ ```
78
+
79
+ **2. As a CLI**
80
+
81
+ ```bash
82
+ siglip2 models download google/siglip2-base-patch16-naflex
83
+ siglip2 models use google/siglip2-base-patch16-naflex
84
+ siglip2 export -i /path/to/photos -o embeddings.jsonl --incremental
85
+ siglip2 classify --image photo.jpg --labels "beach,mountain,wedding"
86
+ ```
87
+
88
+ **3. As a REST service (optional extra)**
89
+
90
+ ```bash
91
+ pip install "siglip2-toolkit[server]"
92
+ siglip2 serve --host 0.0.0.0 --port 8000
93
+ ```
94
+
95
+ Endpoints: `/v1/embed/text`, `/v1/embed/image`, `/v1/classify`, `/v1/health`, `/v1/info`.
96
+
97
+ ## Highlights
98
+
99
+ - **Library-first design**: the CLI and the REST server are thin adapters over the same public API.
100
+ - **Zero-shot classification**: prompt-based labels using the model's calibrated SigLIP logits, plus fine-tuned classification heads (`ImageClassifier`).
101
+ - **Own model cache**: downloads and imports live in the toolkit cache (independent from HuggingFace's); only `models download` uses the network.
102
+ - **GPU-ready**: `device: "auto"` picks CUDA when available (FP16; FP32 on CPU); override per command with `--device`.
103
+ - **Incremental export**: new, modified and deleted files are detected (modification detection requires metadata).
104
+ - **Resumable runs**: interrupted exports continue where they left off.
105
+ - **Optional metadata**: rich file metadata (EXIF, dimensions, timestamps, thumbnails) — or ids and vectors only, your choice.
106
+ - **Versioned export format** (`jsonl` / `pkl`) with an embedding-space manifest.
107
+ - **Retrieval math**: top-k scoring, RRF fusion, similarity-drop filtering and serializable metadata filters, composable with any store.
108
+
109
+ > **Status:** first preview (0.0.1, October 2026). The public API and the
110
+ export format may still change before 1.0; pin the version you depend on.
111
+
112
+ ## License
113
+
114
+ MIT — see [LICENSE](LICENSE).
@@ -0,0 +1,82 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "siglip2-toolkit"
7
+ version = "0.0.1"
8
+ description = "Local-first toolkit around Google's SigLIP 2 models: text/image embeddings, zero-shot and fine-tuned classification, incremental corpus export and retrieval math. Library-first, with CLI and REST interfaces."
9
+ readme = "README.md"
10
+ requires-python = ">=3.12"
11
+ license = "MIT"
12
+ license-files = ["LICENSE"]
13
+ authors = [
14
+ { name = "Yudivián Almeida Cruz", email = "yudivian@gmail.com" },
15
+ ]
16
+ keywords = [
17
+ "siglip2",
18
+ "embeddings",
19
+ "image-search",
20
+ "image-classification",
21
+ "image-retrieval",
22
+ "local-first",
23
+ ]
24
+ classifiers = [
25
+ "Development Status :: 2 - Pre-Alpha",
26
+ "Intended Audience :: Developers",
27
+ "License :: OSI Approved :: MIT License",
28
+ "Operating System :: OS Independent",
29
+ "Programming Language :: Python :: 3",
30
+ "Programming Language :: Python :: 3.12",
31
+ "Programming Language :: Python :: 3.13",
32
+ "Topic :: Multimedia :: Graphics",
33
+ "Topic :: Scientific/Engineering :: Image Recognition",
34
+ ]
35
+
36
+ dependencies = [
37
+ "numpy>=1.26",
38
+ "torch>=2.2",
39
+ "transformers>=4.51",
40
+ "huggingface_hub>=0.23",
41
+ "Pillow>=10.0",
42
+ "pillow-heif>=0.16",
43
+ "PyYAML>=6.0",
44
+ "tqdm>=4.60",
45
+ ]
46
+
47
+ [project.optional-dependencies]
48
+ server = [
49
+ "fastapi>=0.115",
50
+ "uvicorn[standard]>=0.30",
51
+ "python-multipart>=0.0.9",
52
+ ]
53
+ dev = [
54
+ "pytest>=8",
55
+ "pytest-cov>=5",
56
+ "ruff>=0.6",
57
+ ]
58
+
59
+ [project.scripts]
60
+ siglip2 = "siglip2_toolkit.cli:main"
61
+
62
+ [project.urls]
63
+ Homepage = "https://github.com/gia-uh/siglip2-toolkit"
64
+ Repository = "https://github.com/gia-uh/siglip2-toolkit"
65
+ Issues = "https://github.com/gia-uh/siglip2-toolkit/issues"
66
+
67
+ [tool.hatch.build.targets.wheel]
68
+ packages = ["src/siglip2_toolkit"]
69
+
70
+ [tool.ruff]
71
+ line-length = 100
72
+ target-version = "py312"
73
+
74
+ [tool.ruff.lint]
75
+ select = ["E", "F", "I", "W", "UP"]
76
+
77
+ [tool.pytest.ini_options]
78
+ testpaths = ["tests"]
79
+ pythonpath = ["src"]
80
+ markers = [
81
+ "integration: requires torch and a cached SigLIP 2 model; enabled with SIGLIP2_INTEGRATION=1",
82
+ ]
@@ -0,0 +1,71 @@
1
+ """siglip2-toolkit — local-first toolkit around Google's SigLIP 2 models.
2
+
3
+ The package is deliberately lazy: importing :mod:`siglip2_toolkit` does not
4
+ import torch, transformers or FastAPI. Heavy modules are loaded on first
5
+ attribute access, so tooling and metadata inspection stay fast::
6
+
7
+ from siglip2_toolkit import Embedder
8
+
9
+ embedder = Embedder(model="google/siglip2-base-patch16-naflex", device="auto")
10
+ vectors = embedder.embed_texts(["a wedding at sunset"])
11
+ """
12
+
13
+ from importlib import import_module
14
+
15
+ try:
16
+ from importlib.metadata import version as _dist_version
17
+
18
+ __version__ = _dist_version("siglip2-toolkit")
19
+ except Exception: # pragma: no cover - source checkout without installed metadata
20
+ __version__ = "0.0.1"
21
+
22
+ #: Lazy attributes: name -> (relative module, attribute or None for the module itself)
23
+ _LAZY_ATTRS: dict[str, tuple[str, str | None]] = {
24
+ # Public API
25
+ "Embedder": (".embedder", "Embedder"),
26
+ "Exporter": (".exporter", "Exporter"),
27
+ "ExportConfig": (".exporter", "ExportConfig"),
28
+ "ExportResult": (".exporter", "ExportResult"),
29
+ "EmbeddingSpace": (".formats", "EmbeddingSpace"),
30
+ "ImageClassifier": (".classification", "ImageClassifier"),
31
+ "Filter": (".filters", "Filter"),
32
+ "field": (".filters", "field"),
33
+ # Submodules
34
+ "formats": (".formats", None),
35
+ "retrieval": (".retrieval", None),
36
+ "filters": (".filters", None),
37
+ "models": (".models", None),
38
+ "utils": (".utils", None),
39
+ "embedder": (".embedder", None),
40
+ "exporter": (".exporter", None),
41
+ "classification": (".classification", None),
42
+ "cli": (".cli", None),
43
+ "server": (".server", None),
44
+ }
45
+
46
+ __all__ = [
47
+ "Embedder",
48
+ "Exporter",
49
+ "ExportConfig",
50
+ "ExportResult",
51
+ "EmbeddingSpace",
52
+ "ImageClassifier",
53
+ "Filter",
54
+ "field",
55
+ "__version__",
56
+ ]
57
+
58
+
59
+ def __getattr__(name: str):
60
+ try:
61
+ module_name, attr = _LAZY_ATTRS[name]
62
+ except KeyError:
63
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
64
+ module = import_module(module_name, __name__)
65
+ value = module if attr is None else getattr(module, attr)
66
+ globals()[name] = value # cache for subsequent lookups
67
+ return value
68
+
69
+
70
+ def __dir__() -> list[str]:
71
+ return sorted(set(globals()) | set(_LAZY_ATTRS))
@@ -0,0 +1,6 @@
1
+ """Allow ``python -m siglip2_toolkit`` as an alias for the ``siglip2`` CLI."""
2
+
3
+ from .cli import main
4
+
5
+ if __name__ == "__main__":
6
+ raise SystemExit(main())
@@ -0,0 +1,161 @@
1
+ """Fine-tuned image classification with SigLIP-family checkpoints.
2
+
3
+ Zero-shot classification (prompt-based) lives on the embedder:
4
+ :meth:`siglip2_toolkit.embedder.Embedder.classify`. This module covers
5
+ checkpoints that ship a trained classification head, consumed through
6
+ ``AutoModelForImageClassification`` (e.g. ``SiglipForImageClassification``).
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import logging
12
+ import threading
13
+ from pathlib import Path
14
+ from typing import Any, Optional, Sequence
15
+
16
+ import numpy as np
17
+
18
+ from . import utils
19
+ from .embedder import ImageInput, softmax_rows
20
+
21
+ logger = logging.getLogger(__name__)
22
+
23
+
24
+ class ImageClassifier:
25
+ """Classifies images with a fine-tuned image-classification model.
26
+
27
+ The model and its image processor are loaded lazily on first use. Labels
28
+ come from the checkpoint's ``config.id2label``; ``num_labels`` overrides
29
+ the head size when the checkpoint requires it.
30
+ """
31
+
32
+ def __init__(
33
+ self,
34
+ model: str | Path,
35
+ device: str = "auto",
36
+ dtype: Optional[str] = None,
37
+ num_labels: Optional[int] = None,
38
+ trust_remote_code: bool = False,
39
+ ) -> None:
40
+ self.model_path = str(Path(model).expanduser())
41
+ self.model_name = self.model_path
42
+ self.requested_device = device
43
+ self.requested_dtype = dtype
44
+ self.num_labels = num_labels
45
+ self.trust_remote_code = trust_remote_code
46
+
47
+ self._model: Any = None
48
+ self._processor: Any = None
49
+ self._torch: Any = None
50
+ self._device: Any = None
51
+ self._id2label: dict[int, str] = {}
52
+ self._lock = threading.Lock()
53
+
54
+ # ------------------------------------------------------------------ #
55
+
56
+ @property
57
+ def is_loaded(self) -> bool:
58
+ return self._model is not None
59
+
60
+ @property
61
+ def labels(self) -> dict[int, str]:
62
+ """``id -> label`` mapping (empty until the model is loaded)."""
63
+ return dict(self._id2label)
64
+
65
+ @property
66
+ def device(self) -> Optional[str]:
67
+ return str(self._device) if self._device is not None else None
68
+
69
+ def load(self) -> "ImageClassifier":
70
+ """Load the model and its image processor from a local directory (idempotent)."""
71
+ if self._model is not None:
72
+ return self
73
+
74
+ path = Path(self.model_path)
75
+ if not path.is_dir():
76
+ raise FileNotFoundError(
77
+ f"ImageClassifier expects a local model directory, got: "
78
+ f"'{self.model_path}'. Get one with `siglip2 models download` or "
79
+ "`siglip2 models import`, or pass the path to one."
80
+ )
81
+
82
+ import torch
83
+ from transformers import AutoImageProcessor, AutoModelForImageClassification
84
+
85
+ device, dtype = utils.resolve_torch_device(
86
+ torch, self.requested_device, self.requested_dtype
87
+ )
88
+ logger.info("Loading classifier '%s' on %s (%s)...", self.model_name, device, dtype)
89
+
90
+ extra: dict[str, Any] = {}
91
+ if self.num_labels is not None:
92
+ extra["num_labels"] = self.num_labels
93
+
94
+ try:
95
+ model = AutoModelForImageClassification.from_pretrained(
96
+ self.model_path,
97
+ trust_remote_code=self.trust_remote_code,
98
+ torch_dtype=dtype,
99
+ **extra,
100
+ )
101
+ processor = AutoImageProcessor.from_pretrained(
102
+ self.model_path,
103
+ trust_remote_code=self.trust_remote_code,
104
+ )
105
+ except Exception as exc:
106
+ raise OSError(
107
+ f"Could not load a classifier from '{self.model_path}': {exc}"
108
+ ) from exc
109
+ model = model.to(device).eval()
110
+
111
+ self._torch = torch
112
+ self._model = model
113
+ self._processor = processor
114
+ self._device = device
115
+ self._id2label = {
116
+ int(key): str(value)
117
+ for key, value in (getattr(model.config, "id2label", None) or {}).items()
118
+ }
119
+ logger.info("Classifier ready: %d label(s) ✅", len(self._id2label))
120
+ return self
121
+
122
+ # ------------------------------------------------------------------ #
123
+
124
+ def classify(
125
+ self,
126
+ images: Sequence[ImageInput],
127
+ top_k: Optional[int] = None,
128
+ batch_size: int = 16,
129
+ ) -> list[list[tuple[str, float]]]:
130
+ """Classify images; returns (label, probability) pairs sorted by score."""
131
+ self.load()
132
+ images = list(images)
133
+ results: list[list[tuple[str, float]]] = []
134
+
135
+ for start in range(0, len(images), batch_size):
136
+ batch = [utils.load_image(item) for item in images[start : start + batch_size]]
137
+ if not batch:
138
+ continue
139
+ inputs = self._processor(images=batch, return_tensors="pt").to(self._device)
140
+ with self._lock, self._torch.no_grad():
141
+ logits = self._model(**inputs).logits
142
+ probabilities = softmax_rows(
143
+ logits.detach().cpu().to(self._torch.float32).numpy()
144
+ )
145
+ for row in probabilities:
146
+ order = np.argsort(row)[::-1]
147
+ if top_k is not None:
148
+ order = order[: max(0, int(top_k))]
149
+ results.append(
150
+ [
151
+ (self._id2label.get(int(index), f"LABEL_{int(index)}"), float(row[index]))
152
+ for index in order
153
+ ]
154
+ )
155
+ return results
156
+
157
+ def __repr__(self) -> str:
158
+ return (
159
+ f"ImageClassifier(model={self.model_name!r}, "
160
+ f"device={self.requested_device!r}, loaded={self.is_loaded})"
161
+ )