update README: add model_type to sam registry

pull/79/head^2
Hanzi Mao 2023-04-07 13:10:16 -07:00
parent 9e1eb9fdbc
commit 7c524018a6
1 changed files with 7 additions and 5 deletions

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
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>)
masks, _, _ = predictor.predict(<input_prompts>)
```
@ -53,14 +54,15 @@ or generate masks for an entire image:
```
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>)
```
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.
@ -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
```
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)**
* `vit_l`: [ViT-L SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth)