osam 0.1.0__tar.gz

osam-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,39 @@
+name: ci
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.8", "3.12"]
+
+    steps:
+    - uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+    - uses: actions/checkout@v3
+
+    - name: Install main
+      run: |
+        pip install .[test]
+
+    - name: Lint
+      run: |
+        make lint
+
+    - name: Test
+      run: |
+        make test
+
+    - name: Install from dist
+      run: |
+        make build
+        pip install dist/osam-*.tar.gz
+        pip install dist/osam-*.whl

osam-0.1.0/.github/workflows/release.yml

@@ -0,0 +1,117 @@
+name: release
+
+on:
+  push:
+    tags:
+    - 'v*'
+
+jobs:
+  build_binary:
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      matrix:
+        os: [macos-latest, ubuntu-20.04, windows-latest]
+
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+
+    - name: Setup Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.10'
+        architecture: 'x64'
+        cache: 'pip'
+        cache-dependency-path: |
+          pyproject.toml
+
+    - name: Install main
+      run: |
+        pip install .
+
+    - name: Build Executable
+      uses: Nuitka/Nuitka-Action@main
+      with:
+        nuitka-version: main
+        script-name: osam
+        onefile: true
+
+    - name: Test Executable on Linux & macOS
+      run: |
+        ./build/osam.bin --version
+      if: matrix.os != 'windows-latest'
+
+    - name: Test Executable on Windows
+      run: |
+        ./build/osam.exe --version
+      if: matrix.os == 'windows-latest'
+
+    - name: Rename Executable
+      run: |
+        mv build/osam.bin build/osam-${{ runner.os }}.bin
+      if: matrix.os != 'windows-latest'
+
+    - name: Upload Artifacts
+      uses: actions/upload-artifact@v4
+      with:
+        name: build-${{ runner.os }}
+        path: |
+          build/*.exe
+          build/*.bin
+
+  release:
+    needs: build_binary
+
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Build Changelog
+        id: github_release
+        uses: mikepenz/release-changelog-builder-action@v3
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        with:
+          configurationJson: |
+            {
+                "template": "#{{CHANGELOG}}\n\n<details>\n<summary>Uncategorized</summary>\n\n#{{UNCATEGORIZED}}\n</details>\n\n---\n\n- For my daily development update, check [my Twitter/X](https://twitter.com/wkentaro_).\n- If you'd like to support this project, check [my sponsoring page](https://github.com/sponsors/wkentaro).",
+                "pr_template": "- #{{TITLE}} ##{{NUMBER}}",
+                "categories": [
+                    {
+                        "title": "## 🚀 Features",
+                        "labels": ["feature"]
+                    },
+                    {
+                        "title": "## ✨ Enhancement",
+                        "labels": ["enhancement"]
+                    },
+                    {
+                        "title": "## 🐛 Fixes",
+                        "labels": ["fix"]
+                    },
+                    {
+                        "title": "## 💬 Other",
+                        "labels": ["other"]
+                    }
+                ]
+            }
+
+      - name: Create Changelog
+        run: |
+          echo '${{steps.github_release.outputs.changelog}}' > ${{ github.workspace }}-CHANGELOG.txt
+
+      - name: Download Artifact
+        uses: actions/download-artifact@v4
+        with:
+          path: build
+          pattern: build-*
+          merge-multiple: true
+
+      - name: Release
+        uses: softprops/action-gh-release@v1
+        with:
+          draft: true
+          body_path: ${{ github.workspace }}-CHANGELOG.txt
+          files: |
+            build/*

osam-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Kentaro Wada.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

osam-0.1.0/Makefile

@@ -0,0 +1,27 @@
+all:
+	@echo '## Make commands ##'
+	@echo
+	@$(MAKE) -pRrq -f $(lastword $(MAKEFILE_LIST)) : 2>/dev/null | awk -v RS= -F: '/^# File/,/^# Finished Make data base/ {if ($$1 !~ "^[#.]") {print $$1}}' | sort | egrep -v -e '^[^[:alnum:]]' -e '^$@$$' | xargs
+
+lint:
+	mypy --package osam
+	ruff format --check
+	ruff check
+
+format:
+	ruff format
+	ruff check --fix
+
+test:
+	python -m pytest -n auto -v tests
+
+clean:
+	rm -rf build dist *.egg-info
+
+build: clean
+	python -m build --sdist --wheel
+
+upload: build
+	python -m twine upload dist/osam-*
+
+publish: build upload

osam-0.1.0/PKG-INFO

@@ -0,0 +1,120 @@
+Metadata-Version: 2.1
+Name: osam
+Version: 0.1.0
+Summary: Get up and running vision foundational models locally.
+Project-URL: Homepage, https://github.com/wkentaro/osam
+Author-email: Kentaro Wada <www.kentaro.wada@gmail.com>
+License: MIT
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: Environment :: Console
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.8
+Requires-Dist: click
+Requires-Dist: fastapi
+Requires-Dist: gdown>=5.1.0
+Requires-Dist: loguru
+Requires-Dist: onnxruntime
+Requires-Dist: pillow
+Requires-Dist: uvicorn
+Provides-Extra: test
+Requires-Dist: build; extra == 'test'
+Requires-Dist: mypy; extra == 'test'
+Requires-Dist: pytest; extra == 'test'
+Requires-Dist: pytest-xdist; extra == 'test'
+Requires-Dist: ruff; extra == 'test'
+Requires-Dist: twine; extra == 'test'
+Requires-Dist: types-pillow; extra == 'test'
+Description-Content-Type: text/markdown
+
+<div align="center">
+  <img alt="logo" height="200px" src=".readme/icon.png" >
+  <h1>Osam</h1>
+  <p>
+    Get up and running vision foundational models locally.
+  </p>
+  <br>
+  <br>
+  <br>
+</div>
+
+*Osam* (/oʊˈsɑm/) is a tool to run open source vision foundational models locally,
+built inspired by [Ollama](https://github.com/ollama/ollama).
+
+*Osam* gives you:
+
+- **Visual foundational models** - Segment-Anything Model, Efficient-SAM, etc;
+- **Local APIs** - CLI & Python & HTTP interface;
+- **Customization** - Host custom vision models.
+
+
+## Installation
+
+```bash
+pip install osam
+```
+
+
+## Usage
+
+### CLI
+
+```bash
+# Run a model with an image
+osam run efficient-sam:25m --image examples/_images/dogs.jpg > output.jpg
+
+# Get a JSON output
+osam run efficient-sam:25m --image examples/_images/dogs.jpg --json
+# {"model": "efficient-sam:25m", "mask": "..."}
+
+# Give a prompt
+osam run efficient-sam:25m --image examples/_images/dogs.jpg \
+  --prompt '{"points": [[1439, 504], [1439, 1289]], "point_labels": [1, 1]}' > output.jpg
+```
+
+<img src="examples/_images/dogs.jpg" width="35%"> <img src=".readme/dogs_output.jpg" width="35%">  
+<i>Input and output images ('dogs.jpg', 'output.jpg').</i>
+
+### Python
+
+```python
+import osam.apis
+import osam.types
+
+request = osam.types.GenerateRequest(
+    model="efficient-sam:25m",
+    image=np.asarray(PIL.Image.open("examples/_images/dogs.jpg")),
+    prompt=osam.types.Prompt(points=[[1439, 504], [1439, 1289]], point_labels=[1, 1]),
+)
+response = osam.apis.generate(request=request)
+PIL.Image.fromarray(response.mask).save("mask.jpg")
+```
+<img src="examples/_images/dogs.jpg" width="35%"> <img src=".readme/dogs_mask.jpg" width="35%">  
+<i>Input and output images ('dogs.jpg', 'mask.jpg').</i>
+
+### HTTP
+
+```bash
+# Get up the server
+osam serve
+
+# POST request
+curl 127.0.0.1:11368/api/generate -X POST \
+  -H "Content-Type: application/json" \
+  -d "{\"model\": \"efficient-sam:25m\", \"image\": \"$(cat examples/_images/dogs.jpg | base64)\"}" \
+  | jq -r .mask | base64 --decode > mask.jpg
+```
+
+## License
+
+MIT

osam-0.1.0/README.md

@@ -0,0 +1,81 @@
+<div align="center">
+  <img alt="logo" height="200px" src=".readme/icon.png" >
+  <h1>Osam</h1>
+  <p>
+    Get up and running vision foundational models locally.
+  </p>
+  <br>
+  <br>
+  <br>
+</div>
+
+*Osam* (/oʊˈsɑm/) is a tool to run open source vision foundational models locally,
+built inspired by [Ollama](https://github.com/ollama/ollama).
+
+*Osam* gives you:
+
+- **Visual foundational models** - Segment-Anything Model, Efficient-SAM, etc;
+- **Local APIs** - CLI & Python & HTTP interface;
+- **Customization** - Host custom vision models.
+
+
+## Installation
+
+```bash
+pip install osam
+```
+
+
+## Usage
+
+### CLI
+
+```bash
+# Run a model with an image
+osam run efficient-sam:25m --image examples/_images/dogs.jpg > output.jpg
+
+# Get a JSON output
+osam run efficient-sam:25m --image examples/_images/dogs.jpg --json
+# {"model": "efficient-sam:25m", "mask": "..."}
+
+# Give a prompt
+osam run efficient-sam:25m --image examples/_images/dogs.jpg \
+  --prompt '{"points": [[1439, 504], [1439, 1289]], "point_labels": [1, 1]}' > output.jpg
+```
+
+<img src="examples/_images/dogs.jpg" width="35%"> <img src=".readme/dogs_output.jpg" width="35%">  
+<i>Input and output images ('dogs.jpg', 'output.jpg').</i>
+
+### Python
+
+```python
+import osam.apis
+import osam.types
+
+request = osam.types.GenerateRequest(
+    model="efficient-sam:25m",
+    image=np.asarray(PIL.Image.open("examples/_images/dogs.jpg")),
+    prompt=osam.types.Prompt(points=[[1439, 504], [1439, 1289]], point_labels=[1, 1]),
+)
+response = osam.apis.generate(request=request)
+PIL.Image.fromarray(response.mask).save("mask.jpg")
+```
+<img src="examples/_images/dogs.jpg" width="35%"> <img src=".readme/dogs_mask.jpg" width="35%">  
+<i>Input and output images ('dogs.jpg', 'mask.jpg').</i>
+
+### HTTP
+
+```bash
+# Get up the server
+osam serve
+
+# POST request
+curl 127.0.0.1:11368/api/generate -X POST \
+  -H "Content-Type: application/json" \
+  -d "{\"model\": \"efficient-sam:25m\", \"image\": \"$(cat examples/_images/dogs.jpg | base64)\"}" \
+  | jq -r .mask | base64 --decode > mask.jpg
+```
+
+## License
+
+MIT

osam-0.1.0/examples/bash_comparemodels/README.md

@@ -0,0 +1,7 @@
+# Bash Shell Example
+
+## Usage
+
+```
+./comparemodels.sh
+```

osam-0.1.0/examples/bash_comparemodels/comparemodels.sh

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+
+image_file='../_images/dogs.jpg'
+number_of_choices=2
+
+choices=$(osam list --all | awk '{print $1}' | sed '1d')
+
+echo "Select $number_of_choices models to compare:"
+select model in $choices; do
+    if [[ -n $model ]]; then
+        echo "You have selected $model"
+        selections+=("$model")
+        ((COUNT++))
+        if [[ $COUNT -eq $number_of_choices ]]; then
+            break
+        fi
+    else
+        echo "Invalid selection"
+    fi
+done
+
+output_files=()
+for model in "${selections[@]}"; do
+    echo "Running $model"
+    output_file=${model/:/-}.jpg
+    osam run $model --image $image_file > $output_file
+    echo "Exported $model.jpg"
+    output_files+=($output_file)
+done
+echo "Compare results: ${output_files[@]}"

osam-0.1.0/osam/__init__.py

@@ -0,0 +1,3 @@
+import importlib.metadata
+
+__version__ = importlib.metadata.version("osam")

osam-0.1.0/osam/__main__.py

@@ -0,0 +1,154 @@
+import datetime
+import io
+import json
+import os
+import sys
+
+import click
+import numpy as np
+import PIL.Image
+import uvicorn
+from loguru import logger
+
+from osam import __version__
+from osam import _humanize
+from osam import _models
+from osam import _tabulate
+from osam import apis
+from osam import types
+
+
+@click.group(context_settings=dict(help_option_names=["-h", "--help"]))
+@click.version_option(__version__)
+def cli():
+    logger.remove(0)
+    logger.add(
+        sys.stderr, level="INFO", colorize=True, format="<level>{message}</level>"
+    )
+    os.makedirs(os.path.expanduser("~/.cache/osam"), exist_ok=True)
+    logger.add(
+        os.path.expanduser("~/.cache/osam/osam.log"), colorize=True, level="DEBUG"
+    )
+
+
+@cli.command(help="Help about any command")
+@click.argument("subcommand", required=False, type=str)
+@click.pass_context
+def help(ctx, subcommand):
+    if subcommand is None:
+        click.echo(cli.get_help(ctx))
+        return
+
+    subcommand_obj = cli.get_command(ctx, subcommand)
+    if subcommand_obj is None:
+        logger.warning("Unknown subcommand {subcommand!r}", subcommand=subcommand)
+        click.echo(cli.get_help(ctx))
+    else:
+        click.echo(subcommand_obj.get_help(ctx))
+
+
+@cli.command(help="List models")
+@click.option("--all", "-a", "show_all", is_flag=True, help="show all models")
+def list(show_all):
+    rows = []
+    for model in _models.MODELS:
+        size = model.get_size()
+        modified_at = model.get_modified_at()
+
+        if size is None or modified_at is None:
+            if show_all:
+                size = "<not pulled>"
+                modified_at = "<not pulled>"
+            else:
+                continue
+        else:
+            size = _humanize.naturalsize(size)
+            modified_at = _humanize.naturaltime(
+                datetime.datetime.fromtimestamp(modified_at)
+            )
+
+        rows.append([model.name, model.get_id(), size, modified_at])
+    click.echo(_tabulate.tabulate(rows, headers=["NAME", "ID", "SIZE", "MODIFIED"]))
+
+
+@cli.command(help="Pull a model")
+@click.argument("model_name", metavar="model", type=str)
+def pull(model_name):
+    for cls in _models.MODELS:
+        if cls.name == model_name:
+            break
+    else:
+        logger.warning("Model {model_name!r} not found.", model_name=model_name)
+        sys.exit(1)
+
+    logger.info("Pulling {model_name!r}...", model_name=model_name)
+    cls.pull()
+    logger.info("Pulled {model_name!r}", model_name=model_name)
+
+
+@cli.command(help="Remove a model")
+@click.argument("model_name", metavar="model", type=str)
+def rm(model_name):
+    for cls in _models.MODELS:
+        if cls.name == model_name:
+            break
+    else:
+        logger.warning("Model {model_name} not found.", model_name=model_name)
+        sys.exit(1)
+
+    logger.info("Removing {model_name!r}...", model_name=model_name)
+    cls.remove()
+    logger.info("Removed {model_name!r}", model_name=model_name)
+
+
+@cli.command(help="Start server")
+@click.option("--reload", is_flag=True, help="reload server on file changes")
+def serve(reload):
+    logger.info("Starting server...")
+    uvicorn.run("osam._server:app", host="127.0.0.1", port=11368, reload=reload)
+
+
+@cli.command(help="Run a model")
+@click.argument("model_name", metavar="model", type=str)
+@click.option(
+    "--image",
+    "image_path",
+    type=click.Path(exists=True),
+    help="image path",
+    required=True,
+)
+@click.option("--prompt", type=json.loads, help="prompt")
+@click.option("--json", is_flag=True, help="json output")
+def run(model_name: str, image_path: str, prompt, json: bool) -> None:
+    try:
+        request: types.GenerateRequest = types.GenerateRequest(
+            model=model_name,
+            image=np.asarray(PIL.Image.open(image_path)),
+            prompt=prompt,
+        )
+        response: types.GenerateResponse = apis.generate(request=request)
+    except ValueError as e:
+        logger.error("{e}", e=e)
+        sys.exit(1)
+
+    if json:
+        click.echo(response.model_dump_json())
+    else:
+        visualization: np.ndarray = (
+            0.5 * request.image
+            + 0.5
+            * np.array([0, 255, 0])[None, None, :]
+            * (response.mask > 0)[:, :, None]
+        ).astype(np.uint8)
+        sys.stdout.buffer.write(_image_ndarray_to_data(visualization))
+
+
+def _image_ndarray_to_data(ndarray: np.ndarray) -> bytes:
+    pil = PIL.Image.fromarray(ndarray)
+    with io.BytesIO() as f:
+        pil.save(f, format="PNG")
+        return f.getvalue()
+
+
+if __name__ == "__main__":
+    cli()

osam-0.1.0/osam/_humanize.py

@@ -0,0 +1,28 @@
+import datetime
+import math
+
+
+def naturalsize(size: int) -> str:
+    """Convert byte size to human readable format."""
+    if size == 0:
+        return "0B"
+    base = 1024
+    prefixes = ["B", "KB", "MB", "GB", "TB", "PB"]
+    power = int(math.log(size, base))
+    prefix = prefixes[power]
+    size /= base**power
+    return f"{size:.2f} {prefix}"
+
+
+def naturaltime(dt: datetime.datetime) -> str:
+    """Convert datetime to human readable format."""
+    now = datetime.datetime.now()
+    delta = now - dt
+    if delta.days > 0:
+        return f"{delta.days} days ago"
+    elif delta.seconds < 60:
+        return f"{delta.seconds} seconds ago"
+    elif delta.seconds < 3600:
+        return f"{delta.seconds // 60} minutes ago"
+    else:
+        return f"{delta.seconds // 3600} hours ago"

osam-0.1.0/osam/_json.py

@@ -0,0 +1,19 @@
+import base64
+import io
+
+import numpy as np
+import PIL.Image
+
+
+def image_ndarray_to_b64data(ndarray):
+    pil = PIL.Image.fromarray(ndarray)
+    f = io.BytesIO()
+    pil.save(f, format="PNG")
+    data = f.getvalue()
+    return base64.b64encode(data).decode("utf-8")
+
+
+def image_b64data_to_ndarray(b64data):
+    data = base64.b64decode(b64data)
+    pil = PIL.Image.open(io.BytesIO(data))
+    return np.asarray(pil)

osam-0.1.0/osam/_models/__init__.py

@@ -0,0 +1,14 @@
+from ._base import ModelBase  # noqa: F401
+from ._efficient_sam import EfficientSam10m
+from ._efficient_sam import EfficientSam25m
+from ._sam import Sam91m
+from ._sam import Sam308m
+from ._sam import Sam636m
+
+MODELS = [
+    EfficientSam10m,
+    EfficientSam25m,
+    Sam91m,
+    Sam308m,
+    Sam636m,
+]

osam-0.1.0/osam/_models/_base.py

@@ -0,0 +1,102 @@
+import dataclasses
+import hashlib
+import os
+from typing import Dict
+from typing import Optional
+
+import gdown
+import numpy as np
+import onnxruntime
+from loguru import logger
+
+from osam import types
+
+
+@dataclasses.dataclass
+class ModelBlob:
+    url: str
+    hash: str
+
+    @property
+    def path(self):
+        return os.path.expanduser(f"~/.cache/osam/models/blobs/{self.hash}")
+
+    @property
+    def size(self):
+        if os.path.exists(self.path):
+            return os.stat(self.path).st_size
+        else:
+            return None
+
+    @property
+    def modified_at(self):
+        if os.path.exists(self.path):
+            return os.stat(self.path).st_mtime
+        else:
+            return None
+
+    def pull(self):
+        gdown.cached_download(url=self.url, path=self.path, hash=self.hash)
+
+    def remove(self):
+        if os.path.exists(self.path):
+            logger.debug("Removing model blob {path!r}", path=self.path)
+            os.remove(self.path)
+        else:
+            logger.warning("Model blob {path!r} not found", path=self.path)
+
+
+class ModelBase:
+    name: str
+
+    _blobs: Dict[str, ModelBlob]
+    _inference_sessions: Dict[str, onnxruntime.InferenceSession]
+
+    def __init__(self):
+        self.pull()
+        self._inference_sessions = {
+            key: onnxruntime.InferenceSession(blob.path)
+            for key, blob in self._blobs.items()
+        }
+
+    @classmethod
+    def pull(cls):
+        for blob in cls._blobs.values():
+            blob.pull()
+
+    @classmethod
+    def remove(cls):
+        for blob in cls._blobs.values():
+            blob.remove()
+
+    @classmethod
+    def get_id(cls) -> str:
+        return hashlib.md5(
+            "+".join(blob.hash for blob in cls._blobs.values()).encode()
+        ).hexdigest()[:12]
+
+    @classmethod
+    def get_size(cls) -> Optional[int]:
+        size = 0
+        for blob in cls._blobs.values():
+            if blob.size is None:
+                return None
+            size += blob.size
+        return size
+
+    @classmethod
+    def get_modified_at(cls) -> Optional[float]:
+        modified_at = 0
+        for blob in cls._blobs.values():
+            if blob.modified_at is None:
+                return None
+            modified_at = max(modified_at, blob.modified_at)
+        return modified_at
+
+    def encode_image(self, image: np.ndarray) -> types.ImageEmbedding:
+        raise NotImplementedError
+
+    def generate_mask(
+        self, image_embedding: types.ImageEmbedding, prompt: types.Prompt
+    ) -> np.ndarray:
+        raise NotImplementedError

osam-0.1.0/osam/_models/_efficient_sam.py

@@ -0,0 +1,87 @@
+import numpy as np
+
+from osam._models._base import ModelBase
+from osam._models._base import ModelBlob
+from osam.types import ImageEmbedding
+from osam.types import Prompt
+
+
+class EfficientSam(ModelBase):
+    def encode_image(self, image: np.ndarray) -> ImageEmbedding:
+        if image.ndim == 2:
+            raise ValueError("Grayscale images are not supported")
+        if image.ndim == 3 and image.shape[2] == 4:
+            raise ValueError("RGBA images are not supported")
+
+        batched_images = image.transpose(2, 0, 1)[None].astype(np.float32) / 255
+        image_embedding = self._inference_sessions["encoder"].run(
+            output_names=None,
+            input_feed={"batched_images": batched_images},
+        )[0][0]  # (embedding_dim, height, width)
+
+        return ImageEmbedding(
+            original_height=image.shape[0],
+            original_width=image.shape[1],
+            embedding=image_embedding,
+        )
+
+    def generate_mask(
+        self,
+        image_embedding: ImageEmbedding,
+        prompt: Prompt,
+    ) -> np.ndarray:
+        input_point = np.array(prompt.points, dtype=np.float32)
+        input_label = np.array(prompt.point_labels, dtype=np.float32)
+
+        # batch_size, embedding_dim, height, width
+        batched_image_embedding = image_embedding.embedding[None, :, :, :]
+        # batch_size, num_queries, num_points, 2
+        batched_point_coords = input_point[None, None, :, :]
+        # batch_size, num_queries, num_points
+        batched_point_labels = input_label[None, None, :]
+
+        decoder_inputs = {
+            "image_embeddings": batched_image_embedding,
+            "batched_point_coords": batched_point_coords,
+            "batched_point_labels": batched_point_labels,
+            "orig_im_size": np.array(
+                (image_embedding.original_height, image_embedding.original_width),
+                dtype=np.int64,
+            ),
+        }
+
+        masks, _, _ = self._inference_sessions["decoder"].run(None, decoder_inputs)
+        mask = masks[0, 0, 0, :, :]  # (1, 1, 3, H, W) -> (H, W)
+        mask = mask > 0.0
+
+        return mask
+
+
+class EfficientSam10m(EfficientSam):
+    name = "efficient-sam:10m"
+
+    _blobs = {
+        "encoder": ModelBlob(
+            url="https://github.com/labelmeai/efficient-sam/releases/download/onnx-models-20231225/efficient_sam_vitt_encoder.onnx",
+            hash="sha256:7a73ee65aa2c37237c89b4b18e73082f757ffb173899609c5d97a2bbd4ebb02d",
+        ),
+        "decoder": ModelBlob(
+            url="https://github.com/labelmeai/efficient-sam/releases/download/onnx-models-20231225/efficient_sam_vitt_decoder.onnx",
+            hash="sha256:e1afe46232c3bfa3470a6a81c7d3181836a94ea89528aff4e0f2d2c611989efd",
+        ),
+    }
+
+
+class EfficientSam25m(EfficientSam):
+    name = "efficient-sam:25m"
+
+    _blobs = {
+        "encoder": ModelBlob(
+            url="https://github.com/labelmeai/efficient-sam/releases/download/onnx-models-20231225/efficient_sam_vits_encoder.onnx",
+            hash="sha256:4cacbb23c6903b1acf87f1d77ed806b840800c5fcd4ac8f650cbffed474b8896",
+        ),
+        "decoder": ModelBlob(
+            url="https://github.com/labelmeai/efficient-sam/releases/download/onnx-models-20231225/efficient_sam_vits_decoder.onnx",
+            hash="sha256:4727baf23dacfb51d4c16795b2ac382c403505556d0284e84c6ff3d4e8e36f22",
+        ),
+    }

osam-0.1.0/osam/_models/_sam.py

@@ -0,0 +1,184 @@
+import numpy as np
+import onnxruntime
+import PIL.Image
+
+from osam._models._base import ModelBase
+from osam._models._base import ModelBlob
+from osam.types import ImageEmbedding
+from osam.types import Prompt
+
+
+class Sam(ModelBase):
+    _image_size: int = 1024
+
+    def encode_image(self, image: np.ndarray):
+        if image.ndim == 2:
+            raise ValueError("Grayscale images are not supported")
+        if image.ndim == 3 and image.shape[2] == 4:
+            raise ValueError("RGBA images are not supported")
+
+        image_embedding = _compute_image_embedding(
+            encoder_session=self._inference_sessions["encoder"],
+            image=image,
+            image_size=self._image_size,
+        )
+
+        return ImageEmbedding(
+            original_height=image.shape[0],
+            original_width=image.shape[1],
+            embedding=image_embedding,
+        )
+
+    def generate_mask(
+        self,
+        image_embedding: ImageEmbedding,
+        prompt: Prompt,
+    ) -> np.ndarray:
+        return _generate_mask(
+            decoder_session=self._inference_sessions["decoder"],
+            image_embedding=image_embedding,
+            prompt=prompt,
+            image_size=self._image_size,
+        )
+
+
+def _compute_scale_to_resize_image(height: int, width: int, image_size: int):
+    if width > height:
+        scale = image_size / width
+        new_height = int(round(height * scale))
+        new_width = image_size
+    else:
+        scale = image_size / height
+        new_height = image_size
+        new_width = int(round(width * scale))
+    return scale, new_height, new_width
+
+
+def _resize_image(image: np.ndarray, image_size: int):
+    scale, new_height, new_width = _compute_scale_to_resize_image(
+        height=image.shape[0],
+        width=image.shape[1],
+        image_size=image_size,
+    )
+    scaled_image = np.asarray(
+        PIL.Image.fromarray(image).resize(
+            (new_width, new_height), resample=PIL.Image.BILINEAR
+        )
+    ).astype(np.float32)
+    return scale, scaled_image
+
+
+def _compute_image_embedding(
+    encoder_session: onnxruntime.InferenceSession,
+    image: np.ndarray,
+    image_size: int,
+):
+    scale, x = _resize_image(image=image, image_size=image_size)
+    x = (x - np.array([123.675, 116.28, 103.53], dtype=np.float32)) / np.array(
+        [58.395, 57.12, 57.375], dtype=np.float32
+    )
+    x = np.pad(
+        x,
+        (
+            (0, image_size - x.shape[0]),
+            (0, image_size - x.shape[1]),
+            (0, 0),
+        ),
+    )
+    x = x.transpose(2, 0, 1)[None, :, :, :]
+
+    output = encoder_session.run(output_names=None, input_feed={"x": x})
+    image_embedding = output[0][0]  # (embedding_dim, height, width)
+    return image_embedding
+
+
+def _generate_mask(
+    decoder_session: onnxruntime.InferenceSession,
+    image_embedding: ImageEmbedding,
+    prompt: Prompt,
+    image_size: int,
+):
+    onnx_coord = np.concatenate([prompt.points, np.array([[0.0, 0.0]])], axis=0)[
+        None, :, :
+    ]
+    onnx_label = np.concatenate([prompt.point_labels, np.array([-1])], axis=0)[
+        None, :
+    ].astype(np.float32)
+
+    scale, new_height, new_width = _compute_scale_to_resize_image(
+        height=image_embedding.original_height,
+        width=image_embedding.original_width,
+        image_size=image_size,
+    )
+    onnx_coord = (
+        onnx_coord.astype(float)
+        * (
+            new_width / image_embedding.original_width,
+            new_height / image_embedding.original_height,
+        )
+    ).astype(np.float32)
+
+    onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)
+    onnx_has_mask_input = np.array([-1], dtype=np.float32)
+
+    decoder_inputs = {
+        "image_embeddings": image_embedding.embedding[None, :, :, :],
+        "point_coords": onnx_coord,
+        "point_labels": onnx_label,
+        "mask_input": onnx_mask_input,
+        "has_mask_input": onnx_has_mask_input,
+        "orig_im_size": np.array(
+            (image_embedding.original_height, image_embedding.original_width),
+            dtype=np.float32,
+        ),
+    }
+
+    masks, _, _ = decoder_session.run(None, decoder_inputs)
+    mask = masks[0, 0]  # (1, 1, H, W) -> (H, W)
+    mask = mask > 0.0
+    return mask
+
+
+class Sam91m(Sam):
+    name = "sam:91m"
+
+    _blobs = {
+        "encoder": ModelBlob(
+            url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_b_01ec64.quantized.encoder.onnx",
+            hash="sha256:3346b9cc551c9902fbf3b203935e933592b22e042365f58321c17fc12641fd6a",
+        ),
+        "decoder": ModelBlob(
+            url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_b_01ec64.quantized.decoder.onnx",
+            hash="sha256:edbcf1a0afaa55621fb0abe6b3db1516818b609ea9368f309746a3afc68f7613",
+        ),
+    }
+
+
+class Sam308m(Sam):
+    name = "sam:308m"
+
+    _blobs = {
+        "encoder": ModelBlob(
+            url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_l_0b3195.quantized.encoder.onnx",
+            hash="sha256:f7158a4a1fe7f670ef47ea2f7f852685425c1ed6caa40f5df86cbe2b0502034f",
+        ),
+        "decoder": ModelBlob(
+            url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_l_0b3195.quantized.decoder.onnx",
+            hash="sha256:552ebb23bf52c5e5b971ac710d1eb8dccfd88b36cc6aff881d1536d1662e6d7b",
+        ),
+    }
+
+
+class Sam636m(Sam):
+    name = "sam:636m"
+
+    _blobs = {
+        "encoder": ModelBlob(
+            url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_h_4b8939.quantized.encoder.onnx",
+            hash="sha256:a5c745fd4279efc5e5436b412200e983dafc2249ce172af6cc0002a71bb5f485",
+        ),
+        "decoder": ModelBlob(
+            url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_h_4b8939.quantized.decoder.onnx",
+            hash="sha256:020b385a45ffe51097e1acd10cd791075a86171153505f789a793bc382eef210",
+        ),
+    }

osam-0.1.0/osam/_server.py

@@ -0,0 +1,33 @@
+import logging
+import os
+
+import fastapi
+
+from osam import apis
+from osam import types
+
+app: fastapi.FastAPI = fastapi.FastAPI()
+
+
+@app.get("/")
+async def index():
+    return fastapi.Response("osam is running")
+
+
+@app.post("/api/generate")
+async def generate(request: types.GenerateRequest) -> types.GenerateResponse:
+    try:
+        return apis.generate(request=request)
+    except ValueError as e:
+        raise fastapi.HTTPException(400, str(e))
+
+
+# https://github.com/tiangolo/fastapi/issues/1508#issuecomment-638365277
+@app.on_event("startup")
+async def startup_event():
+    uvicorn_logger = logging.getLogger("uvicorn")
+    handler = logging.FileHandler(os.path.expanduser("~/.cache/osam/osam.log"))
+    handler.setLevel(logging.DEBUG)
+    fmt = "%(asctime)s [%(levelname)s] %(module)s:%(funcName)s:%(lineno)s - %(message)s"
+    handler.setFormatter(logging.Formatter(fmt=fmt))
+    uvicorn_logger.addHandler(handler)

osam-0.1.0/osam/_tabulate.py

@@ -0,0 +1,19 @@
+import collections
+
+
+def tabulate(rows, headers):
+    col_widths = collections.defaultdict(int)
+    for row in [headers] + rows:
+        for c, item in enumerate(row):
+            col_widths[c] = max(col_widths[c], len(item))
+
+    table = ""
+
+    for c, header in enumerate(headers):
+        table += header.ljust(col_widths[c] + 4)
+
+    for row in rows:
+        table += "\n"
+        for c, item in enumerate(row):
+            table += item.ljust(col_widths[c] + 4)
+    return table

osam-0.1.0/osam/apis.py

@@ -0,0 +1,40 @@
+from typing import Optional
+
+import numpy as np
+from loguru import logger
+
+from osam import _models
+from osam import types
+
+model: Optional[_models.ModelBase] = None
+
+
+def generate(request: types.GenerateRequest) -> types.GenerateResponse:
+    global model
+
+    if model is None or model.name != request.model:
+        for model_cls in _models.MODELS:
+            if model_cls.name == request.model:
+                model = model_cls()
+                break
+        else:
+            raise ValueError(f"Model not found: {request.model!r}")
+
+    image: np.ndarray = request.image
+
+    if request.prompt is None:
+        height, width = image.shape[:2]
+        request.prompt = types.Prompt(
+            points=np.array([[width / 2, height / 2]], dtype=np.float32),
+            point_labels=np.array([1], dtype=np.int32),
+        )
+        logger.warning(
+            "Prompt is not given, so using the center point as prompt: {prompt!r}",
+            prompt=request.prompt,
+        )
+
+    image_embedding: types.ImageEmbedding = model.encode_image(image=image)
+    mask: np.ndarray = model.generate_mask(
+        image_embedding=image_embedding, prompt=request.prompt
+    )
+    return types.GenerateResponse(model=request.model, mask=mask)

osam-0.1.0/osam/types.py

@@ -0,0 +1,86 @@
+from typing import List
+from typing import Optional
+from typing import Union
+
+import numpy as np
+import pydantic
+
+from osam import _json
+
+
+class ImageEmbedding(pydantic.BaseModel):
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    original_height: int
+    original_width: int
+    embedding: np.ndarray
+
+    @pydantic.validator("embedding")
+    def validate_embedding(cls, embedding):
+        if embedding.ndim != 3:
+            raise ValueError(
+                "embedding must be 3-dimensional: (embedding_dim, height, width)"
+            )
+        return embedding
+
+
+class Prompt(pydantic.BaseModel):
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    points: np.ndarray
+    point_labels: np.ndarray
+
+    @pydantic.validator("points", pre=True)
+    def validate_points(cls, points):
+        if isinstance(points, list):
+            points = np.array(points, dtype=float)
+        if points.ndim != 2:
+            raise ValueError("points must be 2-dimensional")
+        if points.shape[1] != 2:
+            raise ValueError("points must have 2 columns")
+        return points
+
+    @pydantic.field_serializer("points")
+    def serialize_points(self, points: np.ndarray) -> List[List[float]]:
+        return points.tolist()
+
+    @pydantic.validator("point_labels", pre=True)
+    def validate_point_labels(cls, point_labels, values):
+        if isinstance(point_labels, list):
+            point_labels = np.array(point_labels, dtype=int)
+        if point_labels.ndim != 1:
+            raise ValueError("point_labels must be 1-dimensional")
+        if "points" in values and point_labels.shape[0] != values["points"].shape[0]:
+            raise ValueError("point_labels must have the same number of rows as points")
+        if not set(np.unique(point_labels).tolist()).issubset({0, 1}):
+            raise ValueError("point_labels must contain only 0s and 1s")
+        return point_labels
+
+    @pydantic.field_serializer("point_labels")
+    def serialize_point_labels(self, point_labels: np.ndarray) -> List[int]:
+        return point_labels.tolist()
+
+
+class GenerateRequest(pydantic.BaseModel):
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    model: str
+    image: np.ndarray
+    prompt: Optional[Prompt] = pydantic.Field(default=None)
+
+    @pydantic.validator("image", pre=True)
+    def validate_image(cls, image: Union[str, np.ndarray]) -> np.ndarray:
+        if isinstance(image, str):
+            return _json.image_b64data_to_ndarray(b64data=image)
+        return image
+
+
+class GenerateResponse(pydantic.BaseModel):
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    model: str
+    mask: np.ndarray
+
+    @pydantic.field_serializer("mask")
+    def serialize_mask(self, mask: np.ndarray) -> str:
+        return _json.image_ndarray_to_b64data(ndarray=mask)

osam-0.1.0/pyproject.toml

@@ -0,0 +1,102 @@
+[build-system]
+requires = ["hatchling>=1.20.0", "hatch-vcs", "hatch-fancy-pypi-readme"]
+build-backend = "hatchling.build"
+
+[project]
+name = "osam"
+description = "Get up and running vision foundational models locally."
+license = { text = "MIT" }
+requires-python = ">=3.8"
+authors = [
+  { name = "Kentaro Wada", email = "www.kentaro.wada@gmail.com" },
+]
+keywords = []
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Environment :: Console",
+  "Intended Audience :: Developers",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: OS Independent",
+  "Programming Language :: Python",
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: 3.8",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Topic :: Software Development :: Libraries :: Python Modules",
+]
+dependencies = [
+  "click",
+  "gdown>=5.1.0",
+  "fastapi",
+  "loguru",
+  "onnxruntime",
+  "Pillow",
+  "uvicorn",
+]
+dynamic = ["readme", "version"]
+
+[project.optional-dependencies]
+test = [
+  "build",
+  "mypy",
+  "types-Pillow",
+  "pytest",
+  "pytest-xdist",
+  "ruff",
+  "twine",
+]
+
+[project.urls]
+Homepage = "https://github.com/wkentaro/osam"
+
+[tool.hatch.metadata.hooks.fancy-pypi-readme]
+content-type = "text/markdown"
+fragments = [
+  { path = "README.md" },
+]
+
+[tool.hatch.version]
+source = "vcs"
+
+[project.scripts]
+osam = "osam:__main__.cli"
+
+[tool.mypy]
+ignore_missing_imports = true
+
+[tool.ruff]
+exclude = [
+  ".conda",
+  ".git",
+  "src",
+]
+
+line-length = 88
+indent-width = 4
+
+[tool.ruff.lint]
+# Enable Pyflakes (`F`), pycodestyle (`E`), isort (`I`).
+select = ["E", "F", "I"]
+ignore = []
+
+# Allow fix for all enabled rules (when `--fix`) is provided.
+fixable = ["ALL"]
+unfixable = []
+
+[tool.ruff.format]
+# Like Black, use double quotes for strings.
+quote-style = "double"
+
+# Like Black, indent with spaces, rather than tabs.
+indent-style = "space"
+
+# Like Black, respect magic trailing commas.
+skip-magic-trailing-comma = false
+
+# Like Black, automatically detect the appropriate line ending.
+line-ending = "auto"
+
+[tool.ruff.isort]
+force-single-line = true

osam-0.1.0/tests/test___main__.py

@@ -0,0 +1,15 @@
+import io
+import os
+import shlex
+import subprocess
+
+import PIL.Image
+
+here = os.path.dirname(os.path.abspath(__file__))
+
+
+def test_run():
+    cmd = f"osam run efficient-sam:10m --image {here}/../examples/_images/dogs.jpg"
+    output = subprocess.check_output(shlex.split(cmd))
+    image = PIL.Image.open(io.BytesIO(output))
+    assert image.size == (2560, 1600)