import copy import os import tempfile from collections import OrderedDict from functools import partial from types import SimpleNamespace from typing import Callable, Tuple import numpy as np import torch import torchvision.models as models from loguru import logger from ..models.dino.head import DINOHead from ..models.encoders.swin_transformer import swin_base, swin_small, swin_tiny from ..models.encoders.vision_transformer import vit_base, vit_small, vit_tiny from ..models.ibot.head import iBOTHead from ..models.mae.model import masked_vit_base, masked_vit_small, masked_vit_tiny from ..utils.utils import compare_models, set_requires_grad from .wrappers import ( MonetHuggingFaceWrapper, UNetWrapper, ViTHuggingFaceWrapper, ViTWrapper, Wrapper, ) class Embedder: base_path = "https://github.com/vm02-self-supervised-dermatology/self-supervised-models/raw/main" model_dict = { "vit_tiny": vit_tiny, "vit_small": vit_small, "vit_base": vit_base, "swin_tiny": swin_tiny, "swin_small": swin_small, "swin_base": swin_base, "masked_vit_tiny": masked_vit_tiny, "masked_vit_small": masked_vit_small, "masked_vit_base": masked_vit_base, } @staticmethod def load_pretrained( ssl: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: # get the model url model_url = Embedder.get_model_url(ssl) # download the model checkpoint with tempfile.NamedTemporaryFile() as tmp: try: if model_url != "": torch.hub.download_url_to_file(model_url, tmp.name, progress=debug) except Exception as e: logger.error(e) logger.info("Trying again.") torch.hub.download_url_to_file(model_url, tmp.name, progress=debug) # get the loader function loader_func = Embedder.get_model_func(ssl) # load the model load_ret = loader_func( ckp_path=tmp.name, return_info=return_info, debug=debug, **kwargs, ) return load_ret @staticmethod def get_model_url(ssl: str): model_dict = { "byol": f"{Embedder.base_path}/byol/checkpoint-epoch100.pth", "simclr": f"{Embedder.base_path}/simclr/checkpoint-epoch100.pth", "colorme": f"{Embedder.base_path}/colorme/checkpoint-epoch100.pth", "dino": f"{Embedder.base_path}/dino/checkpoint-epoch100.pth", "ibot": f"{Embedder.base_path}/ibot/model_best.pth", "ibot_isic": f"{Embedder.base_path}/ibot_isic/checkpoint-epoch100.pth", "ibot_vit_small": f"{Embedder.base_path}/ibot_vit_small/checkpoint-epoch200.pth", "resnet50_random": "", "vit_tiny_random": "", "imagenet": "", "imagenet_vit_tiny": "", "imagenet_vit_small": "", "imagenet_dino": f"{Embedder.base_path}/imagenet_dino/checkpoint-epoch100.pth", "simclr_imagenet": f"{Embedder.base_path}/simclr_imagenet/resnet50_imagenet_bs2k_epochs600.pth.tar", # Foundation Models "dino_derma": f"{Embedder.base_path}/fund/dino/checkpoint-epoch100.pth", "dino_qderma": f"{Embedder.base_path}/fund/dino_qderma/checkpoint-epoch100.pth", "inet_dino_qderma": f"{Embedder.base_path}/fund/inet_dino_qderma/checkpoint-epoch100.pth", "ibot_derma": f"{Embedder.base_path}/fund/ibot/checkpoint-epoch100.pth", "ibot_qderma": f"{Embedder.base_path}/fund/ibot_qderma/checkpoint-epoch100.pth", "mae_qderma": f"{Embedder.base_path}/fund/mae_qderma/checkpoint-epoch100.pth", # Open-source Models "dinov2_vits14": "", "dinov2_vitb14": "", "monet": "", # PanDerm Models "panderm_base": "", "panderm_large": "", } # get the model url model_url = model_dict.get(ssl, np.nan) if model_url is np.nan: raise ValueError("Unrecognized model name.") return model_url @staticmethod def get_model_func(ssl: str) -> Callable: model_dict_func = { "simclr": Embedder.load_simclr, "byol": Embedder.load_byol, "colorme": Embedder.load_colorme, "dino": Embedder.load_dino, "ibot": Embedder.load_ibot, "ibot_isic": Embedder.load_ibot, "ibot_vit_small": Embedder.load_ibot, "resnet50_random": Embedder.load_resnet50_rand, "vit_tiny_random": Embedder.load_vit_tiny_rand, "imagenet": Embedder.load_resnet50_imagenet, "imagenet_vit_tiny": partial( Embedder.load_vit_imagenet, hf_name="WinKawaks/vit-tiny-patch16-224", out_dim=768, ), "imagenet_vit_small": partial( Embedder.load_vit_imagenet, hf_name="WinKawaks/vit-small-patch16-224", out_dim=1_536, ), "imagenet_dino": Embedder.load_dino, "simclr_imagenet": Embedder.load_simclr_imagenet, # Foundation Models "dino_derma": Embedder.load_dino, "dino_qderma": Embedder.load_dino, "inet_dino_qderma": Embedder.load_inet_dino, "ibot_derma": Embedder.load_ibot, "ibot_qderma": Embedder.load_ibot, "mae_qderma": Embedder.load_mae, # Open-source Models "dinov2_vits14": partial(Embedder.load_dinov2, dino_model="dinov2_vits14"), "dinov2_vitb14": partial(Embedder.load_dinov2, dino_model="dinov2_vitb14"), "monet": partial( Embedder.load_monet, hf_name="suinleelab/monet", out_dim=4096, ), # PanDerm Models "panderm_base": partial(Embedder.load_panderm, variant="base"), "panderm_large": partial(Embedder.load_panderm, variant="large"), } model_func = model_dict_func.get(ssl, None) if model_func is None: raise ValueError("Unrecognized model name.") return model_func @staticmethod def load_resnet50_rand( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: # load a dummy model model = models.resnet50(weights=None, progress=debug) # ResNet Model without last layer model = torch.nn.Sequential(*list(model.children())[:-1]) model = Wrapper(model=model) set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ResNet" info.ssl_type = "ResNet50-Random" info.out_dim = 2048 return model, info, {} return model @staticmethod def load_vit_tiny_rand( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: # load a dummy model model = vit_tiny() # wrap the ViT with a helper model = ViTWrapper(model) out_dim = 256 # handle the number of layers in the head n_head_layers = kwargs.get("n_head_layers", None) if n_head_layers is not None: model, out_dim = Embedder.vit_handle_heads( model=model, n_head_layers=n_head_layers, ) set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ViT" info.ssl_type = "ViT-Tiny" info.out_dim = out_dim return model, info, {} return model @staticmethod def load_resnet50_imagenet( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: # load a dummy model model = models.resnet50(weights="IMAGENET1K_V1", progress=debug) # ResNet Model without last layer model = torch.nn.Sequential(*list(model.children())[:-1]) model = Wrapper(model=model) set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ResNet" info.ssl_type = "ImageNet" info.out_dim = 2048 return model, info, {} return model @staticmethod def load_vit_imagenet( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: # load the huggingface model model = ViTHuggingFaceWrapper( vit_huggingface_name=kwargs.get( "hf_name", "WinKawaks/vit-tiny-patch16-224" ), ) set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ViT" info.ssl_type = "ImageNet-ViT" info.out_dim = kwargs.get("out_dim", 768) return model, info, {} return model @staticmethod def load_monet( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: # load the huggingface model model = MonetHuggingFaceWrapper( vit_huggingface_name=kwargs.get("hf_name", "suinleelab/monet"), ) set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ViT" info.ssl_type = "ImageNet-ViT" info.out_dim = kwargs.get("out_dim", 4096) return model, info, {} return model @staticmethod def load_panderm( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: from .helper_panderm import panderm_base_patch16_224, panderm_large_patch16_224 # Determine model variant from kwargs or checkpoint name variant = kwargs.get("variant", "large") # Auto-detect variant from checkpoint filename if not specified if "panderm_bb" in ckp_path: variant = "base" elif "panderm_ll" in ckp_path: variant = "large" # Load appropriate model architecture if variant == "base": model = panderm_base_patch16_224() out_dim = 768 else: model = panderm_large_patch16_224() out_dim = 1024 # Check if ckp_path is a valid file path # If load_pretrained was called with an empty model_url, ckp_path will be a temp file # In that case, we need to get the actual checkpoint path from kwargs actual_ckp_path = kwargs.get("checkpoint_path", ckp_path) if not os.path.isfile(actual_ckp_path): raise FileNotFoundError( f"PanDerm checkpoint not found at {actual_ckp_path}. " f"Please download the checkpoint or provide a valid path using --checkpoint_path" ) # Load checkpoint checkpoint = torch.load(actual_ckp_path, map_location="cpu") # Handle different checkpoint structures # Large model has "encoder." prefix, base model doesn't if "encoder.cls_token" in checkpoint: # Large model format - need to strip "encoder." prefix state_dict = { k.replace("encoder.", ""): v for k, v in checkpoint.items() if k.startswith("encoder.") } elif "model" in checkpoint: # Wrapped in "model" key state_dict = checkpoint["model"] state_dict = {k.replace("encoder.", ""): v for k, v in state_dict.items()} else: # Base model format - direct state dict state_dict = checkpoint model.load_state_dict(state_dict, strict=False) model = Wrapper(model) set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ViT" info.ssl_type = f"PanDerm-{variant.capitalize()}" info.out_dim = kwargs.get("out_dim", out_dim) return model, info, {} return model @staticmethod def load_dinov2( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: # load the huggingface model model = torch.hub.load("facebookresearch/dinov2", kwargs.get("dino_model")) model = ViTWrapper(model) set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ViT" info.ssl_type = "ImageNet-ViT" info.out_dim = model(torch.rand(1, 3, 224, 224)).shape[-1] return model, info, {} return model @staticmethod def load_simclr_imagenet( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: # load a dummy model model = models.resnet50(weights=None, progress=debug) dummy_model = copy.deepcopy(model) # retreive the config file config = {} to_restore = {"config": config} # load the trained model Embedder.restart_from_checkpoint( ckp_path, state_dict=model, replace_ckp_str="convnet.", run_variables=to_restore, hide_logs=True, ) config = to_restore["config"] # check if the dummy model params and the loaded differ n_differs = compare_models(dummy_model, model) if n_differs == 0: raise ValueError( "Dummy model and loaded model are not different, " "checkpoint wasn't loaded correctly" ) # ResNet Model without last layer model = torch.nn.Sequential(*list(model.children())[:-1]) set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ResNet" info.ssl_type = "SimCLR" info.out_dim = 2048 return model, info, config return model @staticmethod def load_simclr( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: # load a dummy model model = models.resnet50(weights=None, progress=debug) # ResNet Model without last layer model = torch.nn.Sequential(*list(model.children())[:-1]) dummy_model = copy.deepcopy(model) # retreive the config file config = {} to_restore = {"config": config} # load the trained model Embedder.restart_from_checkpoint( ckp_path, state_dict=model, replace_ckp_str="model.", run_variables=to_restore, hide_logs=True, ) config = to_restore["config"] # check if the dummy model params and the loaded differ n_differs = compare_models(dummy_model, model) if n_differs == 0: raise ValueError( "Dummy model and loaded model are not different, " "checkpoint wasn't loaded correctly" ) set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ResNet" info.ssl_type = "SimCLR" info.out_dim = 2048 return model, info, config return model @staticmethod def load_byol( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: # load a dummy model model = models.resnet50(weights=None, progress=debug) dummy_model = copy.deepcopy(model) # retreive the config file config = {} to_restore = {"config": config} # load the trained model Embedder.restart_from_checkpoint( ckp_path, state_dict=model, replace_ckp_str="net.", run_variables=to_restore, hide_logs=True, ) config = to_restore["config"] # check if the dummy model params and the loaded differ n_differs = compare_models(dummy_model, model) if n_differs == 0: raise ValueError( "Dummy model and loaded model are not different, " "checkpoint wasn't loaded correctly" ) # ResNet Model without last layer model = torch.nn.Sequential(*list(model.children())[:-1]) set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ResNet" info.ssl_type = "BYOL" info.out_dim = 2048 return model, info, config return model @staticmethod def load_colorme( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: # load a dummy model import segmentation_models_pytorch as smp model = smp.Unet( encoder_name="resnet50", in_channels=1, classes=2, encoder_weights=None, ) model = model.encoder dummy_model = copy.deepcopy(model) # retreive the config file config = {} to_restore = {"config": config} # load the trained model Embedder.restart_from_checkpoint( ckp_path, state_dict=model, replace_ckp_str="enc_dec_model.encoder.", run_variables=to_restore, hide_logs=True, ) config = to_restore["config"] # check if the dummy model params and the loaded differ n_differs = compare_models(dummy_model, model) if n_differs == 0: raise ValueError( "Dummy model and loaded model are not different, " "checkpoint wasn't loaded correctly" ) # wrap the UNet encoder with a helper model = UNetWrapper(model) set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ResNet-1Channel" info.ssl_type = "ColorMe" info.out_dim = 2048 return model, info, config return model @staticmethod def load_dino( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: model, config = Embedder.load_vit( ckp_path, debug, model_load_dict={"teacher": None} ) head = DINOHead( model.embed_dim * config["model"]["eval"]["n_last_blocks"], config["model"]["out_dim"], use_bn=config["model"]["use_bn_in_head"], norm_last_layer=config["model"]["norm_last_layer"], ) Embedder.restart_from_checkpoint( ckp_path, teacher=head, replace_ckp_str="head.", hide_logs=True, ) # wrap the ViT with a helper emb_dim = model.embed_dim out_dim = 256 model = ViTWrapper(model, head) # handle the number of layers in the head n_head_layers = kwargs.get("n_head_layers", None) if n_head_layers is not None: model, out_dim = Embedder.vit_handle_heads( model=model, n_head_layers=n_head_layers, emb_dim=emb_dim, ) set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ViT" info.ssl_type = "DINO" info.out_dim = out_dim return model, info, config return model @staticmethod def load_vit( ckp_path: str, debug: bool = False, model_load_dict: dict = {}, ) -> Tuple[torch.nn.Module, dict]: # retreive the config file config = {} to_restore = {"config": config} # get the model architecture model, config = Embedder.get_base_model_from_config( ckp_path=ckp_path, to_restore=to_restore, debug=debug, ) dummy_model = copy.deepcopy(model) # load the trained model if len(model_load_dict.keys()) > 0: model_load_dict[list(model_load_dict.keys())[0]] = model else: model_load_dict["state_dict"] = model Embedder.restart_from_checkpoint( ckp_path, replace_ckp_str="backbone.", run_variables=to_restore, hide_logs=True, **model_load_dict, ) model.masked_im_modeling = False model.return_all_tokens = False # check if the dummy model params and the loaded differ n_differs = compare_models(dummy_model, model) if n_differs == 0: raise ValueError( "Dummy model and loaded model are not different, " "checkpoint wasn't loaded correctly" ) return model, config @staticmethod def get_base_model_from_config(ckp_path: str, to_restore: dict, debug: bool): # get the config of the saved model Embedder.restart_from_checkpoint( ckp_path, run_variables=to_restore, hide_logs=True, ) config = to_restore["config"] # get the model architecture model_arch = Embedder.model_dict.get(config["model"]["base_model"], None) if model_arch is None: raise ValueError( f"Invalid base model name: {config['model']['base_model']}" ) # load a dummy model if "teacher" in config["model"].keys(): model = model_arch(**config["model"]["teacher"]) elif "configs" in config["model"].keys(): model = model_arch(**config["model"]["configs"]) else: raise ValueError(f"Can't interpret the model config: {config['model']}") return model, config @staticmethod def load_ibot( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: model, config = Embedder.load_vit( ckp_path, debug, model_load_dict={"teacher": None} ) head = iBOTHead( in_dim=int( config["model"]["emb_dim"] * 4 ), # default head takes the last 4 layers out_dim=config["model"]["out_dim"], patch_out_dim=config["model"]["patch_out_dim"], use_bn=config["model"]["use_bn_in_head"], norm_last_layer=config["model"]["norm_last_layer"], shared_head=config["model"]["shared_head"], ) Embedder.restart_from_checkpoint( ckp_path, teacher=head, replace_ckp_str="head.", hide_logs=True, ) # wrap the ViT with a helper emb_dim = model.embed_dim out_dim = 256 model = ViTWrapper(model, head) # handle the number of layers in the head n_head_layers = kwargs.get("n_head_layers", None) if n_head_layers is not None: model, out_dim = Embedder.vit_handle_heads( model=model, n_head_layers=n_head_layers, emb_dim=emb_dim, ) # set the gradients correctly set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ViT" info.ssl_type = "iBOT" info.out_dim = out_dim return model, info, config return model @staticmethod def load_mae( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: model, config = Embedder.load_vit( ckp_path, debug, model_load_dict={"state_dict": None} ) model = ViTWrapper(model) out_dim = 256 # handle the number of layers in the head n_head_layers = kwargs.get("n_head_layers", None) if n_head_layers is not None: model, out_dim = Embedder.vit_handle_heads( model=model, n_head_layers=n_head_layers, ) set_requires_grad(model, True) if return_info: # information about the model info = SimpleNamespace() info.model_type = "ViT" info.ssl_type = "DINO" info.out_dim = out_dim return model, info, config return model @staticmethod def load_inet_dino( ckp_path: str, return_info: bool = False, debug: bool = False, **kwargs, ) -> torch.nn.Module: # retreive the config file config = {} to_restore = {"config": config} Embedder.restart_from_checkpoint( ckp_path, run_variables=to_restore, hide_logs=True, ) config = to_restore["config"] base_model_name = config["model"]["base_model"].replace("pretrained_", "") model, info, config = Embedder.load_pretrained( base_model_name, return_info=True, ) Embedder.restart_from_checkpoint( ckp_path, replace_ckp_str="backbone.", run_variables=to_restore, hide_logs=True, student=model, ) set_requires_grad(model, True) if return_info: return model, info, config return model @staticmethod def vit_handle_heads( model: torch.nn.Module, n_head_layers: int, emb_dim: int = 192, out_dim: int = 256, ): if n_head_layers == 0: model.head = torch.nn.Identity() out_dim = 4 * emb_dim elif n_head_layers == 1: model.head[1] = torch.nn.Identity() model.head[2] = torch.nn.Identity() model.head[3] = torch.nn.Identity() model.head[4] = torch.nn.Identity() out_dim = 2048 elif n_head_layers == 2: model.head[3] = torch.nn.Identity() model.head[4] = torch.nn.Identity() out_dim = 2048 return model, out_dim @staticmethod def restart_from_checkpoint( ckp_path, run_variables=None, replace_ckp_str="module.", hide_logs: bool = False, **kwargs, ): if not os.path.isfile(ckp_path): logger.info("Pre-trained weights not found. Training from scratch.") return if not hide_logs: logger.info("Found checkpoint at {}".format(ckp_path)) # open checkpoint file checkpoint = torch.load(ckp_path, map_location="cpu", weights_only=False) # key is what to look for in the checkpoint file # value is the object to load # example: {'state_dict': model} for key, value in kwargs.items(): if key in checkpoint and value is not None: try: msg = value.load_state_dict(checkpoint[key], strict=False) if msg is None or len(msg.missing_keys) > 0: k = next(iter(checkpoint[key])) if replace_ckp_str in k: logger.debug( f"=> Found `{replace_ckp_str}` in {key}, trying to transform." ) transf_state_dict = OrderedDict() for k, v in checkpoint[key].items(): # remove the module from the key # this is caused by the distributed training k = k.replace(replace_ckp_str, "") transf_state_dict[k] = v msg = value.load_state_dict(transf_state_dict, strict=False) logger.debug( "=> loaded '{}' from checkpoint '{}' with msg {}".format( key, ckp_path, msg ) ) except TypeError: try: msg = value.load_state_dict(checkpoint[key]) logger.debug( "=> loaded '{}' from checkpoint: '{}'".format(key, ckp_path) ) except ValueError: logger.error( "=> failed to load '{}' from checkpoint: '{}'".format( key, ckp_path ) ) else: logger.error( "=> key '{}' not found in checkpoint: '{}'".format(key, ckp_path) ) # reload variable important for the run if run_variables is not None: for var_name in run_variables: if var_name in checkpoint: run_variables[var_name] = checkpoint[var_name]