Add code snippets

README.md

@@ -1,20 +1,49 @@
 ---
 license: apache-2.0
+pipeline_tag: mask-generation
 library_name: sam2
 ---
 
-
 Repository for SAM 2: Segment Anything in Images and Videos, a foundation model towards solving promptable visual segmentation in images and videos from FAIR. See the [SAM 2 paper](https://arxiv.org/abs/2408.00714) for more information.
 
-
 The official code is publicly release in this [repo](https://github.com/facebookresearch/segment-anything-2/).
-To download the SAM 2 (Hiera-L) checkpoint:
 
+## Usage
+
+For image prediction:
+
+```python
+import torch
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
+
+with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+    predictor.set_image(<your_image>)
+    masks, _, _ = predictor.predict(<input_prompts>)
 ```
-from huggingface_hub import hf_hub_download
-hf_hub_download(repo_id = "facebook/sam2-hiera-large", filename="sam2_hiera_large.pt", local_dir = "./")
+
+For video prediction:
+
+```python
+import torch
+from sam2.sam2_video_predictor import SAM2VideoPredictor
+
+predictor = SAM2VideoPredictor.from_pretrained("facebook/sam2-hiera-large")
+
+with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+    state = predictor.init_state(<your_video>)
+
+    # add new prompts and instantly get the output on the same frame
+    frame_idx, object_ids, masks = predictor.add_new_points_or_box(state, <your_prompts>):
+
+    # propagate the prompts to get masklets throughout the video
+    for frame_idx, object_ids, masks in predictor.propagate_in_video(state):
+        ...
 ```
 
+Refer to the [demo notebooks](https://github.com/facebookresearch/segment-anything-2/tree/main/notebooks) for details.
+
 ### Citation
 
 To cite the paper, model, or software, please use the below: