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]
# 'mean': [.838, .855, .770],
# 'std': [.281, .210, .300]
}
model_cfg.num_max_points = 12
model_cfg.input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(model_cfg.input_normalization['mean'],
model_cfg.input_normalization['std']),
])
model = get_deeplab_model(backbone='resnet50',
deeplab_ch=128,
aspp_dropout=0.20,
norm_radius=180)
model.to(cfg.device)
model.apply(initializer.XavierGluon(rnd_type='gaussian', magnitude=2.0))
model.feature_extractor.load_pretrained_weights()
return model, model_cfg
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
