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