def init_model(cfg):
model_cfg = edict()
model_cfg.crop_size = (320, 480)
model_cfg.input_normalization = {
'mean': [.485, .456, .406],
'std': [.229, .224, .225]
}
model_cfg.num_max_points = 10
model_cfg.input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(model_cfg.input_normalization['mean'],
model_cfg.input_normalization['std']),
])
model = get_hrnet_model(width=32,
ocr_width=128,
max_interactive_points=model_cfg.num_max_points,
with_aux_output=True)
model.to(cfg.device)
model.apply(initializer.XavierGluon(rnd_type='gaussian', magnitude=2.0))
model.feature_extractor.load_pretrained_weights(
cfg.IMAGENET_PRETRAINED_MODELS.HRNETV2_W32)
return model, model_cfg
def load_hrnet_is_model(state_dict,
device,
backbone='auto',
width=48,
ocr_width=256,
small=False,
cpu_dist_maps=False,
norm_radius=260):
if backbone == 'auto':
num_fe_weights = len(
[x for x in state_dict.keys() if 'feature_extractor.' in x])
small = num_fe_weights < 1800
ocr_f_down = [
v for k, v in state_dict.items()
if 'object_context_block.f_down.1.0.bias' in k
]
assert len(ocr_f_down) == 1
ocr_width = ocr_f_down[0].shape[0]
s2_conv1_w = [
v for k, v in state_dict.items()
if 'stage2.0.branches.0.0.conv1.weight' in k
]
assert len(s2_conv1_w) == 1
width = s2_conv1_w[0].shape[0]
model = get_hrnet_model(width=width,
ocr_width=ocr_width,
small=small,
with_aux_output=False,
cpu_dist_maps=cpu_dist_maps,
norm_radius=norm_radius)
model.load_state_dict(state_dict, strict=False)
for param in model.parameters():
param.requires_grad = False
model.to(device)
model.eval()
return model