Spaces:
Sleeping
Sleeping
from pathlib import Path | |
import torch | |
import torch.nn as nn | |
from climategan.deeplab.deeplab_v2 import DeepLabV2Decoder | |
from climategan.deeplab.deeplab_v3 import DeepLabV3Decoder | |
from climategan.deeplab.mobilenet_v3 import MobileNetV2 | |
from climategan.deeplab.resnet101_v3 import ResNet101 | |
from climategan.deeplab.resnetmulti_v2 import ResNetMulti | |
def create_encoder(opts, no_init=False, verbose=0): | |
if opts.gen.encoder.architecture == "deeplabv2": | |
if verbose > 0: | |
print(" - Add Deeplabv2 Encoder") | |
return DeeplabV2Encoder(opts, no_init, verbose) | |
elif opts.gen.encoder.architecture == "deeplabv3": | |
if verbose > 0: | |
backone = opts.gen.deeplabv3.backbone | |
print(" - Add Deeplabv3 ({}) Encoder".format(backone)) | |
return build_v3_backbone(opts, no_init) | |
else: | |
raise NotImplementedError( | |
"Unknown encoder: {}".format(opts.gen.encoder.architecture) | |
) | |
def create_segmentation_decoder(opts, no_init=False, verbose=0): | |
if opts.gen.s.architecture == "deeplabv2": | |
if verbose > 0: | |
print(" - Add DeepLabV2Decoder") | |
return DeepLabV2Decoder(opts) | |
elif opts.gen.s.architecture == "deeplabv3": | |
if verbose > 0: | |
print(" - Add DeepLabV3Decoder") | |
return DeepLabV3Decoder(opts, no_init) | |
else: | |
raise NotImplementedError( | |
"Unknown Segmentation architecture: {}".format(opts.gen.s.architecture) | |
) | |
def build_v3_backbone(opts, no_init, verbose=0): | |
backbone = opts.gen.deeplabv3.backbone | |
output_stride = opts.gen.deeplabv3.output_stride | |
if backbone == "resnet": | |
resnet = ResNet101( | |
output_stride=output_stride, | |
BatchNorm=nn.BatchNorm2d, | |
verbose=verbose, | |
no_init=no_init, | |
) | |
if not no_init: | |
if opts.gen.deeplabv3.backbone == "resnet": | |
assert Path(opts.gen.deeplabv3.pretrained_model.resnet).exists() | |
std = torch.load(opts.gen.deeplabv3.pretrained_model.resnet) | |
resnet.load_state_dict( | |
{ | |
k.replace("backbone.", ""): v | |
for k, v in std.items() | |
if k.startswith("backbone.") | |
} | |
) | |
print( | |
" - Loaded pre-trained DeepLabv3+ Resnet101 Backbone as Encoder" | |
) | |
return resnet | |
elif opts.gen.deeplabv3.backbone == "mobilenet": | |
assert Path(opts.gen.deeplabv3.pretrained_model.mobilenet).exists() | |
mobilenet = MobileNetV2( | |
no_init=no_init, | |
pretrained_path=opts.gen.deeplabv3.pretrained_model.mobilenet, | |
) | |
print(" - Loaded pre-trained DeepLabv3+ MobileNetV2 Backbone as Encoder") | |
return mobilenet | |
else: | |
raise NotImplementedError("Unknown backbone in " + str(opts.gen.deeplabv3)) | |
class DeeplabV2Encoder(nn.Module): | |
def __init__(self, opts, no_init=False, verbose=0): | |
"""Deeplab architecture encoder""" | |
super().__init__() | |
self.model = ResNetMulti(opts.gen.deeplabv2.nblocks, opts.gen.encoder.n_res) | |
if opts.gen.deeplabv2.use_pretrained and not no_init: | |
saved_state_dict = torch.load(opts.gen.deeplabv2.pretrained_model) | |
new_params = self.model.state_dict().copy() | |
for i in saved_state_dict: | |
i_parts = i.split(".") | |
if not i_parts[1] in ["layer5", "resblock"]: | |
new_params[".".join(i_parts[1:])] = saved_state_dict[i] | |
self.model.load_state_dict(new_params) | |
if verbose > 0: | |
print(" - Loaded pretrained weights") | |
def forward(self, x): | |
return self.model(x) | |