| from transformers import SamConfig, SamModel, SamProcessor, SamImageProcessor | |
| from transformers.models.sam.convert_sam_original_to_hf_format import replace_keys | |
| from segment_anything import sam_model_registry # pip install git+https://github.com/facebookresearch/segment-anything.git | |
| # load the MedSAM ViT-B model | |
| checkpoint = 'medsam_vit_b.pth' # https://drive.google.com/drive/folders/1ETWmi4AiniJeWOt6HAsYgTjYv_fkgzoN?usp=drive_link | |
| pt_model = sam_model_registry['vit_b'](checkpoint) | |
| pt_state_dict = pt_model.state_dict() | |
| # tweak the model's weights to transformers design | |
| hf_state_dict = replace_keys(pt_state_dict) | |
| # save the model | |
| hf_model = SamModel(config=SamConfig()) | |
| hf_model.load_state_dict(hf_state_dict) | |
| hf_model.save_pretrained('./') | |
| # update the processor | |
| hf_processor = SamProcessor( | |
| image_processor=SamImageProcessor( | |
| do_normalize=False, | |
| image_mean=[0, 0, 0], | |
| image_std=[1, 1, 1], | |
| resample=3, | |
| ) | |
| ) | |
| # save the processor | |
| hf_processor.save_pretrained('./') | |