114 lines
4.1 KiB
Python
114 lines
4.1 KiB
Python
import functools
|
|
|
|
import torch
|
|
import torch.utils.model_zoo as model_zoo
|
|
|
|
from ._preprocessing import preprocess_input
|
|
# from .densenet import densenet_encoders
|
|
# from .dpn import dpn_encoders
|
|
# from .efficientnet import efficient_net_encoders
|
|
# from .inceptionresnetv2 import inceptionresnetv2_encoders
|
|
# from .inceptionv4 import inceptionv4_encoders
|
|
# from .mobilenet import mobilenet_encoders
|
|
from .resnet import resnet_encoders
|
|
# from .senet import senet_encoders
|
|
# from .timm_efficientnet import timm_efficientnet_encoders
|
|
# from .timm_gernet import timm_gernet_encoders
|
|
# from .timm_mobilenetv3 import timm_mobilenetv3_encoders
|
|
# from .timm_regnet import timm_regnet_encoders
|
|
# from .timm_res2net import timm_res2net_encoders
|
|
# from .timm_resnest import timm_resnest_encoders
|
|
# from .timm_sknet import timm_sknet_encoders
|
|
# from .timm_universal import TimmUniversalEncoder
|
|
# from .vgg import vgg_encoders
|
|
# from .xception import xception_encoders
|
|
# from .swin_transformer import swin_transformer_encoders
|
|
# from .mit_encoder import mit_encoders
|
|
# from .hrnet import hrnet_encoders
|
|
|
|
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
|
|
encoders = {}
|
|
encoders.update(resnet_encoders)
|
|
# encoders.update(dpn_encoders)
|
|
# encoders.update(vgg_encoders)
|
|
# encoders.update(senet_encoders)
|
|
# encoders.update(densenet_encoders)
|
|
# encoders.update(inceptionresnetv2_encoders)
|
|
# encoders.update(inceptionv4_encoders)
|
|
# encoders.update(efficient_net_encoders)
|
|
# encoders.update(mobilenet_encoders)
|
|
# encoders.update(xception_encoders)
|
|
# encoders.update(timm_efficientnet_encoders)
|
|
# encoders.update(timm_resnest_encoders)
|
|
# encoders.update(timm_res2net_encoders)
|
|
# encoders.update(timm_regnet_encoders)
|
|
# encoders.update(timm_sknet_encoders)
|
|
# encoders.update(timm_mobilenetv3_encoders)
|
|
# encoders.update(timm_gernet_encoders)
|
|
# encoders.update(swin_transformer_encoders)
|
|
# encoders.update(mit_encoders)
|
|
# encoders.update(hrnet_encoders)
|
|
|
|
|
|
def get_encoder(name, in_channels=3, depth=5, weights=None, output_stride=32, **kwargs):
|
|
|
|
if name.startswith("tu-"):
|
|
name = name[3:]
|
|
# encoder = TimmUniversalEncoder(
|
|
# name=name,
|
|
# in_channels=in_channels,
|
|
# depth=depth,
|
|
# output_stride=output_stride,
|
|
# pretrained=weights is not None,
|
|
# **kwargs
|
|
# )
|
|
# return encoder
|
|
|
|
try:
|
|
Encoder = encoders[name]["encoder"]
|
|
except KeyError:
|
|
raise KeyError("Wrong encoder name `{}`, supported encoders: {}".format(name, list(encoders.keys())))
|
|
|
|
params = encoders[name]["params"]
|
|
params.update(depth=depth)
|
|
encoder = Encoder(**params)
|
|
|
|
if weights is not None:
|
|
try:
|
|
settings = encoders[name]["pretrained_settings"][weights]
|
|
except KeyError:
|
|
raise KeyError("Wrong pretrained weights `{}` for encoder `{}`. Available options are: {}".format(
|
|
weights, name, list(encoders[name]["pretrained_settings"].keys()),
|
|
))
|
|
encoder.load_state_dict(model_zoo.load_url(settings["url"], map_location=torch.device(DEVICE)))
|
|
|
|
encoder.set_in_channels(in_channels, pretrained=weights is not None)
|
|
if output_stride != 32:
|
|
encoder.make_dilated(output_stride)
|
|
|
|
return encoder
|
|
|
|
|
|
def get_encoder_names():
|
|
return list(encoders.keys())
|
|
|
|
|
|
def get_preprocessing_params(encoder_name, pretrained="imagenet"):
|
|
settings = encoders[encoder_name]["pretrained_settings"]
|
|
|
|
if pretrained not in settings.keys():
|
|
raise ValueError("Available pretrained options {}".format(settings.keys()))
|
|
|
|
formatted_settings = {}
|
|
formatted_settings["input_space"] = settings[pretrained].get("input_space")
|
|
formatted_settings["input_range"] = settings[pretrained].get("input_range")
|
|
formatted_settings["mean"] = settings[pretrained].get("mean")
|
|
formatted_settings["std"] = settings[pretrained].get("std")
|
|
return formatted_settings
|
|
|
|
|
|
def get_preprocessing_fn(encoder_name, pretrained="imagenet"):
|
|
params = get_preprocessing_params(encoder_name, pretrained=pretrained)
|
|
return functools.partial(preprocess_input, **params)
|