| from models.segment_anything.build_sam import sam_model_registry | |
| from models.segment_anything_samihs.build_sam_ihs import samihs_model_registry | |
| def get_model(modelname="SAM", args=None, opt=None, ckpt_path=None): | |
| if modelname == "SAM": | |
| model = sam_model_registry['vit_b'](checkpoint=args.sam_ckpt) | |
| elif modelname == 'MedSAM': | |
| model = sam_model_registry['vit_b'](checkpoint=None) | |
| elif modelname == "SAMIHS": | |
| if ckpt_path: | |
| model = samihs_model_registry['vit_b'](args=args, checkpoint=ckpt_path) | |
| else: | |
| model = samihs_model_registry['vit_b'](args=args, checkpoint=args.sam_ckpt) | |
| else: | |
| raise RuntimeError("Could not find the model:", modelname) | |
| return model |