Skip to content

Commit

Permalink
hooks: add hook for sam2
Browse files Browse the repository at this point in the history
Add hook for `sam2` (Segment Anything Model 2).
  • Loading branch information
rokm committed Jan 5, 2025
1 parent 80443a5 commit c3b5792
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 0 deletions.
30 changes: 30 additions & 0 deletions _pyinstaller_hooks_contrib/stdhooks/hook-sam2.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
# ------------------------------------------------------------------
# Copyright (c) 2025 PyInstaller Development Team.
#
# This file is distributed under the terms of the GNU General Public
# License (version 2.0 or later).
#
# The full license is available in LICENSE, distributed with
# this software.
#
# SPDX-License-Identifier: GPL-2.0-or-later
# ------------------------------------------------------------------

# Hook for Segment Anything Model 2 (SAM 2): https://pypi.org/project/sam2

from PyInstaller.utils.hooks import collect_data_files, collect_submodules

# Collect config .yaml files.
datas = collect_data_files('sam2')

# Ensure that all indirectly-imported modules are collected (e.g., `sam2.modeling.backbones`).
hiddenimports = collect_submodules('sam2')

# Due to use of `torch.script`, we need to collect source .py files for `sam2`. The `sam2/__init__.py` also seems to be
# required by `hydra`. Furthermore, the source-based introspection attempts to load the source of stdlib `enum` module.
# The module collection mode support and run-time discovery of source .py files for modules that are collected into
# `base_library.zip` archive was added in pyinstaller/pyinstaller#8971 (i.e., PyInstaller > 6.11.1).
module_collection_mode = {
'sam2': 'pyz+py',
'enum': 'pyz+py', # requires PyInstaller > 6.11.1; no-op in earlier versions
}
1 change: 1 addition & 0 deletions news/847.new.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add hook for ``sam2`` (Segment Anything Model 2).
21 changes: 21 additions & 0 deletions tests/test_deep_learning.py
Original file line number Diff line number Diff line change
Expand Up @@ -469,3 +469,24 @@ def forward(self, x, y):
assert (z == x + y).all()
""", app_name="model_test", app_args=[str(model_file)])


# Basic test for Segment Anything Model 2 (SAM 2)
@importorskip('torch')
@importorskip('sam2')
def test_sam2(pyi_builder):
pyi_builder.test_source("""
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
# No checkpoint data (= untrained model), since we would have to download it.
checkpoint = None
# Run on CPU to make test applicable to all OSes (default device is CUDA).
device = 'cpu'
model_cfg = "configs/sam2.1/sam2.1_hiera_t.yaml"
predictor = SAM2ImagePredictor(build_sam2(model_cfg, checkpoint, device=device))
print(predictor)
""")

0 comments on commit c3b5792

Please sign in to comment.