import torch
from typing import OrderedDict


def get_clean_state_dict(checkpoint_name):
  """
  Get the state dictionary for the backbone only. Indeed the checkpoints provided in 
  https://github.com/d-li14/involution#model-zoo contain backbone, head, ...

  :param checkpoint_name: path to the checkpoint file (.pth)

  :return state dictionary of the backbone (RedNetInvolution)
  """
  state_dict = torch.load(checkpoint_name)['state_dict']

  clean_state_dict = OrderedDict()
  for key in state_dict.keys():
    if 'backbone' in key:
      new_key = '.'.join(key.split('.')[1:])
      clean_state_dict[new_key] = state_dict[key]
  
  return clean_state_dict