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

Add xxh128 version of hashes for model #284

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
1 change: 1 addition & 0 deletions environment_cpu.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ dependencies:
- napari
- pip
- pooch
- python-xxhash
- python-elf >=0.4.8
- pytorch
- segment-anything
Expand Down
1 change: 1 addition & 0 deletions environment_gpu.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ dependencies:
- napari
- pip
- pooch
- python-xxhash
- python-elf >=0.4.8
- pytorch
- pytorch-cuda>=11.7 # you may need to update the cuda version to match your system
Expand Down
57 changes: 44 additions & 13 deletions micro_sam/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,12 @@
from segment_anything import sam_model_registry, SamPredictor
VIT_T_SUPPORT = False

try:
import xxhash
HAS_XXH128 = hasattr(xxhash, 'xxh128')
except ImportError:
HAS_XXH128 = False

try:
from napari.utils import progress as tqdm
except ImportError:
Expand Down Expand Up @@ -82,22 +88,47 @@ def models():
so any user changes to the default micro-sam cache directory location
are respected.
"""

# Provide hashes in both xxh128 (fast, but not cryptographically secure),
# and as sha256 (as a fallback) to validate if the file has been correctly
# downloaded.
# https://github.com/computational-cell-analytics/micro-sam/issues/283
# To generate the xxh128 hash
#
# xxh128sum filename
#
# You may need to install xxhash with conda or your system package manager.
registry_sha256 = {
# the default segment anything models
"vit_h": "sha256:a7bf3b02f3ebf1267aba913ff637d9a2d5c33d3173bb679e46d9f338c26f262e",
"vit_l": "sha256:3adcc4315b642a4d2101128f611684e8734c41232a17c648ed1693702a49a622",
"vit_b": "sha256:ec2df62732614e57411cdcf32a23ffdf28910380d03139ee0f4fcbe91eb8c912",
# the model with vit tiny backend fom https://github.com/ChaoningZhang/MobileSAM
"vit_t": "sha256:6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f",
# first version of finetuned models on zenodo
"vit_h_lm": "sha256:9a65ee0cddc05a98d60469a12a058859c89dc3ea3ba39fed9b90d786253fbf26",
"vit_b_lm": "sha256:5a59cc4064092d54cd4d92cd967e39168f3760905431e868e474d60fe5464ecd",
"vit_h_em": "sha256:ae3798a0646c8df1d4db147998a2d37e402ff57d3aa4e571792fbb911d8a979c",
"vit_b_em": "sha256:c04a714a4e14a110f0eec055a65f7409d54e6bf733164d2933a0ce556f7d6f81",
}
registry_xxh128 = {
# the default segment anything models
"vit_h": "xxh128:97698fac30bd929c2e6d8d8cc15933c2",
"vit_l": "xxh128:a82beb3c660661e3dd38d999cc860e9a",
"vit_b": "xxh128:6923c33df3637b6a922d7682bfc9a86b",
# the model with vit tiny backend fom https://github.com/ChaoningZhang/MobileSAM
"vit_t": "xxh128:8eadbc88aeb9d8c7e0b4b60c3db48bd0",
# first version of finetuned models on zenodo
"vit_h_lm": "xxh128:e113adac6a0a21514bb2d73de16b921b",
"vit_b_lm": "xxh128:5fc0851abf8a209dcbed4e95634d9e27",
"vit_h_em": "xxh128:64b6eb2d32ac9c5d9b022b1ac57f1cc6",
"vit_b_em": "xxh128:f50d499db5bf54dc9849c3dbd271d5c9",
}

models = pooch.create(
path=os.path.join(microsam_cachedir(), "models"),
base_url="",
registry={
# the default segment anything models
"vit_h": "a7bf3b02f3ebf1267aba913ff637d9a2d5c33d3173bb679e46d9f338c26f262e",
"vit_l": "3adcc4315b642a4d2101128f611684e8734c41232a17c648ed1693702a49a622",
"vit_b": "ec2df62732614e57411cdcf32a23ffdf28910380d03139ee0f4fcbe91eb8c912",
# the model with vit tiny backend fom https://github.com/ChaoningZhang/MobileSAM
"vit_t": "6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f",
# first version of finetuned models on zenodo
"vit_h_lm": "9a65ee0cddc05a98d60469a12a058859c89dc3ea3ba39fed9b90d786253fbf26",
"vit_b_lm": "5a59cc4064092d54cd4d92cd967e39168f3760905431e868e474d60fe5464ecd",
"vit_h_em": "ae3798a0646c8df1d4db147998a2d37e402ff57d3aa4e571792fbb911d8a979c",
"vit_b_em": "c04a714a4e14a110f0eec055a65f7409d54e6bf733164d2933a0ce556f7d6f81",
},
registry=registry_xxh128 if HAS_XXH128 else registry_sha256,
# Now specify custom URLs for some of the files in the registry.
urls={
# the default segment anything models
Expand Down