From ef9d059eaa15d4d1bacb161426ed4bd15a3eaf38 Mon Sep 17 00:00:00 2001
From: Donghoon Kim <dhkim2810@gmail.com>
Date: Fri, 30 Jun 2023 13:54:48 +0900
Subject: [PATCH] Update README recording to updated model package

---
 README.md     | 18 ++++++++----------
 app/README.md |  5 +++--
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 3378c29..79218ee 100644
--- a/README.md
+++ b/README.md
@@ -104,19 +104,17 @@ cd MobileSAM; pip install -e .
 The MobileSAM can be loaded in the following ways:
 
 ```
-from mobile_encoder.setup_mobile_sam import setup_model
-checkpoint = torch.load('../weights/mobile_sam.pt')
-mobile_sam = setup_model()
-mobile_sam.load_state_dict(checkpoint,strict=True)
-```
+from mobile_sam import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
 
-Then the model can be easily used in just a few lines to get masks from a given prompt:
+model_type = "vit_t"
+sam_checkpoint = "./weights/mobile_sam.pt"
 
-```
-from segment_anything import SamPredictor
-device = "cuda"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+mobile_sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
 mobile_sam.to(device=device)
 mobile_sam.eval()
+
 predictor = SamPredictor(mobile_sam)
 predictor.set_image(<your_image>)
 masks, _, _ = predictor.predict(<input_prompts>)
@@ -125,7 +123,7 @@ masks, _, _ = predictor.predict(<input_prompts>)
 or generate masks for an entire image:
 
 ```
-from segment_anything import SamAutomaticMaskGenerator
+from mobile_sam import SamAutomaticMaskGenerator
 
 mask_generator = SamAutomaticMaskGenerator(mobile_sam)
 masks = mask_generator.generate(<your_image>)
diff --git a/app/README.md b/app/README.md
index 724b920..fa636dd 100755
--- a/app/README.md
+++ b/app/README.md
@@ -13,14 +13,15 @@ license: apache-2.0
 
 # Faster Segment Anything(MobileSAM)
 
-Official PyTorch Implementation of the <a href="https://github.com/ChaoningZhang/MobileSAM">.
+Demo of official PyTorch implementation of the <a href="https://github.com/ChaoningZhang/MobileSAM">.
 
 
 **MobileSAM** performs on par with the original SAM (at least visually) and keeps exactly the same pipeline as the original SAM except for a change on the image encoder.
 Specifically, we replace the original heavyweight ViT-H encoder (632M) with a much smaller Tiny-ViT (5M). On a single GPU, MobileSAM runs around 12ms per image: 8ms on the image encoder and 4ms on the mask decoder. 
 
 ## To run on local PC
-First, mobile_sam must be installed to run on pc. [Instructions](https://github.com/dhkim2810/MobileSAM/tree/master#installation)
+First, mobile_sam must be installed to run on pc. Refer to [Installation Instruction](https://github.com/dhkim2810/MobileSAM/tree/master#installation)
+
 Then run the following
 
 ```