Commit 7c524018 authored by Hanzi Mao's avatar Hanzi Mao
Browse files

update README: add model_type to sam registry

parent 9e1eb9fd
...@@ -44,7 +44,8 @@ First download a [model checkpoint](#model-checkpoints). Then the model can be u ...@@ -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: ...@@ -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/ ...@@ -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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment