merged build_sam methods

lang_sam/lang_sam.py

@@ -53,33 +53,33 @@ def __init__(self, sam_type="vit_h", ckpt_path=None):
  self.sam_type = sam_type
  self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  self.build_groundingdino()
- if ckpt_path == None:
- self.build_sam(sam_type)
+ self.build_sam(ckpt_path)
+
+ def build_sam(self, ckpt_path):
+ if self.sam_type is None or ckpt_path is None:
+ if self.sam_type is None:
+ print("No sam type indicated. Using vit_h by default.")
+ self.sam_type = "vit_h"
+ checkpoint_url = SAM_MODELS[self.sam_type]
+ try:
+ sam = sam_model_registry[self.sam_type]()
+ state_dict = torch.hub.load_state_dict_from_url(checkpoint_url)
+ sam.load_state_dict(state_dict, strict=True)
+ except:
+ raise ValueError(f"Problem loading SAM please make sure you have the right model type: {sam_type} \
+ and a working checkpoint: {checkpoint_url}. Recommend deleting the checkpoint and \
+ re-downloading it.")
+ sam.to(device=self.device)
+ self.sam = SamPredictor(sam)
  else:
- self.build_sam_with_ckpt(sam_type, ckpt_path)
-
- def build_sam_with_ckpt(self, sam_type, ckpt_path):
- try:
- sam = sam_model_registry[sam_type](ckpt_path)
- except:
- raise ValueError(f"Problem loading SAM. Your model type: {sam_type} \
+ try:
+ sam = sam_model_registry[self.sam_type](ckpt_path)
+ except:
+ raise ValueError(f"Problem loading SAM. Your model type: {sam_type} \
  should match your checkpoint path: {ckpt_path}. Recommend calling LangSAM \
  using matching model type AND checkpoint path")
- sam.to(device=self.device)
- self.sam = SamPredictor(sam)
-
- def build_sam(self, sam_type):
- checkpoint_url = SAM_MODELS[sam_type]
- try:
- sam = sam_model_registry[sam_type]()
- state_dict = torch.hub.load_state_dict_from_url(checkpoint_url)
- sam.load_state_dict(state_dict, strict=True)
- except:
- raise ValueError(f"Problem loading SAM please make sure you have the right model type: {sam_type} \
- and a working checkpoint: {checkpoint_url}. Recommend deleting the checkpoint and \
- re-downloading it.")
- sam.to(device=self.device)
- self.sam = SamPredictor(sam)
+ sam.to(device=self.device)
+ self.sam = SamPredictor(sam)
 
  def build_groundingdino(self):
  ckpt_repo_id = "ShilongLiu/GroundingDINO"