update README: add model_type to sam registry

This commit is contained in:
Hanzi Mao 2023-04-07 13:10:16 -07:00
parent 9e1eb9fdbc
commit 7c524018a6

View File

@ -44,7 +44,8 @@ First download a [model checkpoint](#model-checkpoints). Then the model can be u
``` ```
from segment_anything import build_sam, SamPredictor from segment_anything import build_sam, SamPredictor
predictor = SamPredictor(build_sam(checkpoint="</path/to/model.pth>")) sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
predictor = SamPredictor(sam)
predictor.set_image(<your_image>) predictor.set_image(<your_image>)
masks, _, _ = predictor.predict(<input_prompts>) masks, _, _ = predictor.predict(<input_prompts>)
``` ```
@ -53,14 +54,15 @@ or generate masks for an entire image:
``` ```
from segment_anything import build_sam, SamAutomaticMaskGenerator from segment_anything import build_sam, SamAutomaticMaskGenerator
mask_generator = SamAutomaticMaskGenerator(build_sam(checkpoint="</path/to/model.pth>")) sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
mask_generator = SamAutomaticMaskGenerator(sam)
masks = mask_generator.generate(<your_image>) masks = mask_generator.generate(<your_image>)
``` ```
Additionally, masks can be generated for images from the command line: Additionally, masks can be generated for images from the command line:
``` ```
python scripts/amg.py --checkpoint <path/to/sam/checkpoint> --input <image_or_folder> --output <output_directory> python scripts/amg.py --checkpoint <path/to/checkpoint> --input <image_or_folder> --output <path/to/output>
``` ```
See the examples notebooks on [using SAM with prompts](/notebooks/predictor_example.ipynb) and [automatically generating masks](/notebooks/automatic_mask_generator_example.ipynb) for more details. See the examples notebooks on [using SAM with prompts](/notebooks/predictor_example.ipynb) and [automatically generating masks](/notebooks/automatic_mask_generator_example.ipynb) for more details.
@ -85,9 +87,9 @@ See the [example notebook](https://github.com/facebookresearch/segment-anything/
Three model versions of the model are available with different backbone sizes. These models can be instantiated by running Three model versions of the model are available with different backbone sizes. These models can be instantiated by running
``` ```
from segment_anything import sam_model_registry from segment_anything import sam_model_registry
sam = sam_model_registry["<name>"](checkpoint="<path/to/checkpoint>") sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
``` ```
Click the links below to download the checkpoint for the corresponding model name. The default model in bold can also be instantiated with `build_sam`, as in the examples in [Getting Started](#getting-started). Click the links below to download the checkpoint for the corresponding model type. The default model in bold can also be instantiated with `build_sam`, as in the examples in [Getting Started](#getting-started).
* **`default` or `vit_h`: [ViT-H SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth)** * **`default` or `vit_h`: [ViT-H SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth)**
* `vit_l`: [ViT-L SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth) * `vit_l`: [ViT-L SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth)