def init_model(cfg): model_cfg = edict() model_cfg.crop_size = (512, 512) model_cfg.input_normalization = { 'mean': [.485, .456, .406], 'std': [.229, .224, .225] } model_cfg.input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(model_cfg.input_normalization['mean'], model_cfg.input_normalization['std']), ]) model = HRNetIHModel(BMCONFIGS['improved_dih512'], mask_fusion='', downsize_hrnet_input=True, small=False) model.to(cfg.device) model.apply(initializer.XavierGluon(rnd_type='gaussian', magnitude=2.0)) model.backbone.load_pretrained_weights( cfg.IMAGENET_PRETRAINED_MODELS.HRNETV2_W18) return model, model_cfg
def init_model(cfg): model_cfg = edict() model_cfg.crop_size = (512, 512) model_cfg.input_normalization = { 'mean': [.485, .456, .406], 'std': [.229, .224, .225] } model_cfg.input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(model_cfg.input_normalization['mean'], model_cfg.input_normalization['std']), ]) model = DeepImageHarmonization(depth=8, batchnorm_from=2, image_fusion=True) model.to(cfg.device) model.apply(initializer.XavierGluon(rnd_type='gaussian', magnitude=2.0)) return model, model_cfg
def init_model(cfg): model_cfg = edict() model_cfg.crop_size = (256, 256) model_cfg.input_normalization = { 'mean': [.485, .456, .406], 'std': [.229, .224, .225] } model_cfg.input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(model_cfg.input_normalization['mean'], model_cfg.input_normalization['std']), ]) model = DeepLabIHModel(BMCONFIGS['improved_dih256']) model.to(cfg.device) model.apply(initializer.XavierGluon(rnd_type='gaussian', magnitude=2.0)) model.backbone.load_pretrained_weights() return model, model_cfg
def init_model(cfg): model_cfg = edict() model_cfg.crop_size = (256, 256) model_cfg.input_normalization = { 'mean': [.485, .456, .406], 'std': [.229, .224, .225] } model_cfg.depth = 4 model_cfg.input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(model_cfg.input_normalization['mean'], model_cfg.input_normalization['std']), ]) model = ISEUNetV1(depth=4, batchnorm_from=2, attend_from=1, ch=64) model.to(cfg.device) model.apply(initializer.XavierGluon(rnd_type='gaussian', magnitude=1.0)) return model, model_cfg
def init_model(cfg): model_cfg = edict() model_cfg.crop_size = (256, 256) model_cfg.input_normalization = { 'mean': [.485, .456, .406], 'std': [.229, .224, .225] } model_cfg.input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(model_cfg.input_normalization['mean'], model_cfg.input_normalization['std']), ]) model = HRNetIHModel(BMCONFIGS['improved_dih256'], pyramid_channels=256) model.to(cfg.device) model.apply(initializer.XavierGluon(rnd_type='gaussian', magnitude=2.0)) model.backbone.load_pretrained_weights( cfg.IMAGENET_PRETRAINED_MODELS.HRNETV2_W18_SMALL) return model, model_cfg