class ModelNetShow():
def __init__(self):
self.backbone_factory = BackboneFactory()
self.model_factory = ModelFactory()
self.show_process = ModelShow()
def model_show(self, model_path):
input_x = torch.randn(1, 3, 224, 224)
self.show_process.set_input(input_x)
model = self.model_factory.get_model(model_path)
self.show_process.show_from_model(model)
def backbone_show(self, backbone_path):
input_x = torch.randn(1, 3, 224, 224)
self.show_process.set_input(input_x)
model = self.backbone_factory.get_base_model(backbone_path)
self.show_process.show_from_model(model)
def onnx_show(self, onnx_path):
input_x = torch.randn(1, 3, 640, 352)
self.show_process.set_input(input_x)
self.show_process.show_from_onnx(onnx_path)
def pc_model_show(self, backbone_path):
input_x = torch.randn(1, 3, 1024)
self.show_process.set_input(input_x)
model = self.model_factory.get_model(backbone_path)
self.show_process.show_from_model(model)
def pc_backbone_show(self, backbone_path):
input_x = torch.randn(1, 3, 1024)
self.show_process.set_input(input_x)
model = self.backbone_factory.get_base_model(backbone_path)
self.show_process.show_from_model(model)
def __init__(self):
super().__init__()
self.set_name(ModelName.ComplexYOLO)
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
def __init__(self, modelDefine, cfg_dir):
super().__init__()
self.backbone_factory = BackboneFactory()
self.createTaskList = CreateModuleList()
self.modelDefine = modelDefine
self.cfg_dir = cfg_dir
self.create_block_list()
def __init__(self, data_channel=3, class_num=19):
super().__init__()
self.set_name(ModelName.ICNet)
self.data_channel = data_channel
self.class_number = class_num
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
def __init__(self, class_num=2):
super().__init__()
self.set_name(ModelName.FCNSeg)
self.class_number = class_num
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
def __init__(self, data_channel=3, class_num=150):
super().__init__()
self.set_name(ModelName.EncNetSeg)
self.data_channel = data_channel
self.class_number = class_num
self.is_jpu = True
self.lateral = False
self.is_se_loss = True
self.is_aux = True
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
class Inceptionv4Cls(BaseModel):
def __init__(self, data_channel=3, class_num=100):
super().__init__()
self.set_name(ModelName.Inceptionv4Cls)
self.data_channel = data_channel
self.class_number = class_num
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
def create_block_list(self):
self.clear_list()
backbone = self.factory.get_base_model(BackboneName.InceptionV4,
self.data_channel)
base_out_channels = backbone.get_outchannel_list()
self.add_block_list(BlockType.BaseNet, backbone, base_out_channels[-1])
avgpool = nn.AvgPool2d(7)
# avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.add_block_list(LayerType.GlobalAvgPool, avgpool,
base_out_channels[-1])
dropout = nn.Dropout2d(1 - 0.8)
self.add_block_list(LayerType.Dropout, dropout, base_out_channels[-1])
layer = FcLayer(base_out_channels[-1], self.class_number)
self.add_block_list(layer.get_name(), layer, self.class_number)
self.create_loss()
def create_loss(self, input_dict=None):
self.lossList = []
loss = CrossEntropy2d(ignore_index=250)
self.add_block_list(LossType.CrossEntropy2d, loss,
self.block_out_channels[-1])
self.lossList.append(loss)
def forward(self, x):
base_outputs = []
layer_outputs = []
output = []
for key, block in self._modules.items():
if BlockType.BaseNet in key:
base_outputs = block(x)
x = base_outputs[-1]
elif LayerType.RouteLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.ShortcutLayer in key:
x = block(layer_outputs)
elif LossType.YoloLoss in key:
output.append(x)
elif LossType.CrossEntropy2d in key:
output.append(x)
else:
x = block(x)
layer_outputs.append(x)
return output
class ModelConverter():
def __init__(self, input_size=(352, 640)):
self.backbone_factory = BackboneFactory()
self.model_factory = ModelFactory()
self.converter = TorchConvertOnnx()
self.input_size = input_size # w * h
def model_convert(self, model_path, weight_path, save_dir):
input_x = torch.randn(1, 3, self.input_size[1], self.input_size[0])
self.converter.set_input(input_x)
self.converter.set_save_dir(save_dir)
model = self.model_factory.get_model(model_path)
self.converter.torch2onnx(model, weight_path)
def base_model_convert(self, base_model_path, weight_path, save_dir):
input_x = torch.randn(1, 3, self.input_size[1], self.input_size[0])
self.converter.set_input(input_x)
self.converter.set_save_dir(save_dir)
model = self.backbone_factory.get_base_model(base_model_path)
self.converter.torch2onnx(model, weight_path)
class BiSeNet(BaseModel):
def __init__(self, data_channel=3, class_num=2):
super().__init__()
self.set_name(ModelName.FastSCNN)
self.data_channel = data_channel
self.class_number = class_num
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
def create_block_list(self):
self.clear_list()
spatial_path = SpatialPath(self.data_channel,
128,
bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(spatial_path.get_name(), spatial_path, 128)
backbone = self.factory.get_base_model(BackboneName.ResNet18)
base_out_channels = backbone.get_outchannel_list()
self.add_block_list(BlockType.BaseNet, backbone, base_out_channels[-1])
inter_channels = 128
global_context = GlobalAvgPooling(512,
inter_channels,
bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(global_context.get_name(), global_context,
base_out_channels[-1])
layer1 = RouteLayer('9')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in layer1.layers
])
self.add_block_list(layer1.get_name(), layer1, output_channel)
context_path1 = ContextPath(512,
inter_channels,
bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(context_path1.get_name(), context_path1,
inter_channels)
layer2 = RouteLayer('7')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in layer2.layers
])
self.add_block_list(layer2.get_name(), layer2, output_channel)
context_path2 = ContextPath(256,
inter_channels,
bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(context_path2.get_name(), context_path2,
inter_channels)
ffm = FeatureFusionBlock(256,
256,
4,
bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(ffm.get_name(), ffm, 256)
conv1 = ConvBNActivationBlock(in_channels=256,
out_channels=64,
kernel_size=3,
padding=1,
bias=False,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(conv1.get_name(), conv1, 64)
dropout = nn.Dropout(0.1)
self.add_block_list(LayerType.Dropout, dropout,
self.block_out_channels[-1])
conv2 = nn.Conv2d(64, self.class_number, 1)
self.add_block_list(LayerType.Convolutional, conv2, self.class_number)
layer = Upsample(scale_factor=8, mode='bilinear')
self.add_block_list(layer.get_name(), layer,
self.block_out_channels[-1])
self.create_loss()
def create_loss(self, input_dict=None):
self.lossList = []
loss = CrossEntropy2d(ignore_index=250)
self.add_block_list(LossType.CrossEntropy2d, loss,
self.block_out_channels[-1])
self.lossList.append(loss)
def forward(self, x):
base_outputs = []
layer_outputs = []
output = []
for key, block in self._modules.items():
if BiSeNetBlockName.SpatialPath in key:
spatial_out = block(x)
layer_outputs.append(spatial_out)
else:
if BlockType.BaseNet in key:
base_outputs = block(x)
x = base_outputs[-1]
elif LayerType.RouteLayer in key:
x = block(layer_outputs, base_outputs)
elif BiSeNetBlockName.ContextPath in key:
x = block(layer_outputs[-2], layer_outputs[-1])
elif BiSeNetBlockName.FeatureFusionBlock in key:
x = block(layer_outputs[0], layer_outputs[-1])
elif LossType.CrossEntropy2d in key:
output.append(x)
else:
x = block(x)
layer_outputs.append(x)
print(key, layer_outputs[-1].shape)
return output
class MyModel(BaseModel):
def __init__(self, modelDefine, cfg_dir):
super().__init__()
self.backbone_factory = BackboneFactory()
self.createTaskList = CreateModuleList()
self.modelDefine = modelDefine
self.cfg_dir = cfg_dir
self.create_block_list()
def create_block_list(self):
basicModel = self.creatBaseModel()
self.add_module(BlockType.BaseNet, basicModel)
taskBlockDict = self.createTask(basicModel)
for key, block in taskBlockDict.items():
self.add_module(key, block)
self.create_loss(input_dict=taskBlockDict)
def create_loss(self, input_dict=None):
self.lossList = []
for key, block in input_dict.items():
if LossType.CrossEntropy2d in key:
self.lossList.append(block)
elif LossType.OhemCrossEntropy2d in key:
self.lossList.append(block)
elif LossType.BinaryCrossEntropy2d in key:
self.lossList.append(block)
elif LossType.YoloLoss in key:
self.lossList.append(block)
def creatBaseModel(self):
result = None
base_net = self.modelDefine[0]
if base_net["type"] == BlockType.BaseNet:
base_net_name = base_net["name"]
self.modelDefine.pop(0)
input_name = base_net_name.strip()
if input_name.endswith("cfg"):
input_name = os.path.join(self.cfg_dir, input_name)
result = self.backbone_factory.get_base_model(input_name)
return result
def createTask(self, basicModel):
blockDict = OrderedDict()
outChannels = basicModel.get_outchannel_list()
if basicModel:
self.createTaskList.createOrderedDict(self.modelDefine,
outChannels)
blockDict = self.createTaskList.getBlockList()
return blockDict
def forward(self, x):
base_outputs = []
layer_outputs = []
output = []
for key, block in self._modules.items():
if BlockType.BaseNet in key:
base_outputs = block(x)
x = base_outputs[-1]
elif LayerType.MultiplyLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.AddLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.RouteLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.ShortRouteLayer in key:
x = block(layer_outputs)
elif LayerType.ShortcutLayer in key:
x = block(layer_outputs)
elif LossType.YoloLoss in key:
output.append(x)
elif LossType.CrossEntropy2d in key:
output.append(x)
elif LossType.OhemCrossEntropy2d in key:
output.append(x)
elif LossType.BinaryCrossEntropy2d in key:
output.append(x)
else:
x = block(x)
layer_outputs.append(x)
return output
def __init__(self):
self.backbone_factory = BackboneFactory()
self.model_factory = ModelFactory()
self.show_process = ModelShow()
class ComplexYOLO(BaseModel):
def __init__(self):
super().__init__()
self.set_name(ModelName.ComplexYOLO)
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
def create_block_list(self):
self.block_out_channels = []
self.index = 0
basic_model = self.factory.get_base_model("./cfg/det3d/complex-darknet19.cfg")
base_out_channels = basic_model.get_outchannel_list()
self.add_block_list(BlockType.BaseNet, basic_model, base_out_channels[-1])
input_channel = self.block_out_channels[-1]
output_channel = 1024
conv1 = ConvBNActivationBlock(input_channel, output_channel,
kernel_size=3, stride=1, padding=1,
bnName=self.bn_name, activationName=self.activation_name)
self.add_block_list(conv1.get_name(), conv1, output_channel)
input_channel = self.block_out_channels[-1]
output_channel = 1024
conv2 = ConvBNActivationBlock(input_channel, output_channel,
kernel_size=3, stride=1, padding=1,
bnName=self.bn_name, activationName=self.activation_name)
self.add_block_list(conv2.get_name(), conv2, output_channel)
layer1 = RouteLayer('12')
output_channel = sum([base_out_channels[i] if i >= 0 else self.block_out_channels[i] for i in layer1.layers])
self.add_block_list(layer1.get_name(), layer1, output_channel)
reorg = ReorgBlock()
output_channel = reorg.stride * reorg.stride * self.out_channels[-1]
self.add_block_list(reorg.get_name(), reorg, output_channel)
layer2 = RouteLayer('-3,-1')
output_channel = sum([base_out_channels[i] if i >= 0 else self.block_out_channels[i] for i in layer2.layers])
self.add_block_list(layer2.get_name(), layer2, output_channel)
input_channel = self.block_out_channels[-1]
output_channel = 1024
conv3 = ConvBNActivationBlock(input_channel, output_channel,
kernel_size=3, stride=1, padding=1,
bnName=self.bn_name, activationName=self.activation_name)
self.add_block_list(conv3.get_name(), conv3, output_channel)
input_channel = self.block_out_channels[-1]
output_channel = 75
conv4 = nn.Conv2d(in_channels=input_channel,
out_channels=output_channel,
kernel_size=1,
stride=1,
padding=0,
bias=True)
self.add_block_list(LayerType.Convolutional, conv4, output_channel)
def create_loss(self, input_dict=None):
self.lossList = []
def forward(self, x):
base_outputs = []
layer_outputs = []
output = []
for key, block in self._modules.items():
if BlockType.BaseNet in key:
base_outputs = block(x)
x = base_outputs[-1]
elif LayerType.RouteLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.ShortcutLayer in key:
x = block(layer_outputs)
else:
x = block(x)
layer_outputs.append(x)
return layer_outputs
class FCN8sSeg(BaseModel):
def __init__(self, class_num=2):
super().__init__()
self.set_name(ModelName.FCNSeg)
self.class_number = class_num
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
def create_block_list(self):
self.block_out_channels = []
self.index = 0
backbone = self.factory.get_base_model(BackboneName.Vgg16)
base_out_channels = backbone.get_outchannel_list()
self.add_block_list(BlockType.BaseNet, backbone, base_out_channels[-1])
layer1 = ConvBNActivationBlock(in_channels=base_out_channels[-1],
out_channels=4096,
kernel_size=7,
padding=3,
bias=True,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(layer1.get_name(), layer1, 4096)
dropout1 = nn.Dropout()
self.add_block_list(LayerType.Dropout, dropout1, 4096)
layer2 = ConvBNActivationBlock(in_channels=4096,
out_channels=4096,
kernel_size=1,
padding=0,
bias=True,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(layer2.get_name(), layer2, 4096)
dropout2 = nn.Dropout()
self.add_block_list(LayerType.Dropout, dropout2, 4096)
layer3 = nn.Conv2d(4096, self.class_number, kernel_size=1)
self.add_block_list(LayerType.Convolutional, layer3, self.class_number)
upsample1 = Upsample(scale_factor=2, mode='bilinear')
self.add_block_list(LayerType.Upsample, upsample1,
self.block_out_channels[-1])
layer4 = RouteLayer('13')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in layer4.layers
])
self.add_block_list(layer4.get_name(), layer4, output_channel)
layer5 = nn.Conv2d(512, self.class_number, kernel_size=1)
self.add_block_list(LayerType.Convolutional, layer5, self.class_number)
layer6 = AddLayer('-1,-3')
index = layer6.layers[0]
output_channel = base_out_channels[
index] if index >= 0 else self.block_out_channels[index]
self.add_block_list(layer6.get_name(), layer6, output_channel)
layer7 = RouteLayer('9')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in layer7.layers
])
self.add_block_list(layer7.get_name(), layer7, output_channel)
layer8 = nn.Conv2d(256, self.class_number, kernel_size=1)
self.add_block_list(LayerType.Convolutional, layer8, self.class_number)
layer9 = RouteLayer('-3')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in layer7.layers
])
self.add_block_list(layer9.get_name(), layer9, output_channel)
upsample2 = Upsample(scale_factor=2, mode='bilinear')
self.add_block_list(LayerType.Upsample, upsample2,
self.block_out_channels[-1])
layer10 = AddLayer('-1,-3')
index = layer10.layers[0]
output_channel = base_out_channels[
index] if index >= 0 else self.block_out_channels[index]
self.add_block_list(layer10.get_name(), layer10, output_channel)
upsample3 = Upsample(scale_factor=8, mode='bilinear')
self.add_block_list(LayerType.Upsample, upsample3,
self.block_out_channels[-1])
self.create_loss()
def create_loss(self, input_dict=None):
self.lossList = []
loss = CrossEntropy2d(ignore_index=250)
self.add_block_list(LossType.CrossEntropy2d, loss,
self.block_out_channels[-1])
self.lossList.append(loss)
def forward(self, x):
base_outputs = []
layer_outputs = []
output = []
for key, block in self._modules.items():
if BlockType.BaseNet in key:
base_outputs = block(x)
x = base_outputs[-1]
elif LayerType.MultiplyLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.AddLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.RouteLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.ShortcutLayer in key:
x = block(layer_outputs)
elif LossType.CrossEntropy2d in key:
output.append(x)
else:
x = block(x)
layer_outputs.append(x)
return output
class DeepLabV3Plus(BaseModel):
def __init__(self, data_channel=3, class_num=2):
super().__init__()
self.set_name(ModelName.DeepLabV3Plus)
self.data_channel = data_channel
self.class_number = class_num
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
def create_block_list(self):
self.clear_list()
c1_channels = 256
c4_channels = 2048
backbone = self.factory.get_base_model(BackboneName.Xception65)
base_out_channels = backbone.get_outchannel_list()
self.add_block_list(BlockType.BaseNet, backbone, base_out_channels[-1])
aspp = ASPP(c4_channels,
256,
bn_name=self.bn_name,
activation_name=self.activation_name)
out_channel = 256
self.add_block_list(aspp.get_name(), aspp, out_channel)
up1 = Upsample(scale_factor=4, mode='bilinear')
self.add_block_list(up1.get_name(), up1, self.block_out_channels[-1])
layer1 = RouteLayer('3')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in layer1.layers
])
self.add_block_list(layer1.get_name(), layer1, output_channel)
conv1 = ConvBNActivationBlock(in_channels=c1_channels,
out_channels=48,
kernel_size=1,
bias=False,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(conv1.get_name(), conv1, 48)
layer2 = RouteLayer('-3, -1')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in layer2.layers
])
self.add_block_list(layer2.get_name(), layer2, output_channel)
block1 = SeparableConv2dBNActivation(
output_channel,
256,
3,
relu_first=False,
bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(block1.get_name(), block1, 256)
block2 = SeparableConv2dBNActivation(
256,
256,
3,
relu_first=False,
bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(block2.get_name(), block2, 256)
conv2 = nn.Conv2d(256, self.class_number, 1)
self.add_block_list(LayerType.Convolutional, conv2, self.class_number)
layer3 = Upsample(scale_factor=4, mode='bilinear')
self.add_block_list(layer3.get_name(), layer3,
self.block_out_channels[-1])
self.create_loss()
def create_loss(self, input_dict=None):
self.lossList = []
loss = CrossEntropy2d(ignore_index=250)
self.add_block_list(LossType.CrossEntropy2d, loss,
self.block_out_channels[-1])
self.lossList.append(loss)
def forward(self, x):
base_outputs = []
layer_outputs = []
output = []
for key, block in self._modules.items():
if BlockType.BaseNet in key:
base_outputs = block(x)
x = base_outputs[-1]
elif LayerType.MultiplyLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.AddLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.RouteLayer in key:
x = block(layer_outputs, base_outputs)
elif LossType.CrossEntropy2d in key:
output.append(x)
else:
x = block(x)
layer_outputs.append(x)
print(key, x.shape)
return output
class ICNet(BaseModel):
def __init__(self, data_channel=3, class_num=19):
super().__init__()
self.set_name(ModelName.ICNet)
self.data_channel = data_channel
self.class_number = class_num
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
def create_block_list(self):
self.clear_list()
block1 = InputDownSample()
self.add_block_list(block1.get_name(), block1, self.data_channel)
conv1 = ConvBNActivationBlock(in_channels=self.data_channel,
out_channels=32,
kernel_size=3,
padding=1,
stride=2,
bias=False,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(conv1.get_name(), conv1, 32)
conv2 = ConvBNActivationBlock(in_channels=32,
out_channels=32,
kernel_size=3,
padding=1,
stride=2,
bias=False,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(conv2.get_name(), conv2, 32)
conv3 = ConvBNActivationBlock(in_channels=32,
out_channels=64,
kernel_size=3,
padding=1,
stride=2,
bias=False,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(conv3.get_name(), conv3, 64)
backbone1 = self.factory.get_base_model(BackboneName.ResNet50, data_channel=self.data_channel)
base_out_channels1 = backbone1.get_outchannel_list()
self.add_block_list(BlockType.BaseNet, backbone1, base_out_channels1[-1])
layer1 = RouteLayer('8')
output_channel = sum([base_out_channels1[i] if i >= 0
else self.block_out_channels[i] for i in layer1.layers])
self.add_block_list(layer1.get_name(), layer1, output_channel)
backbone2 = self.factory.get_base_model(BackboneName.ResNet50, data_channel=self.data_channel)
base_out_channels2 = backbone2.get_outchannel_list()
self.add_block_list(BlockType.BaseNet, backbone2, base_out_channels2[-1])
layer2 = RouteLayer('17')
output_channel = sum([base_out_channels2[i] if i >= 0
else self.block_out_channels[i] for i in layer2.layers])
self.add_block_list(layer2.get_name(), layer2, output_channel)
self.create_head()
self.create_loss()
def create_head(self):
cff_24 = CascadeFeatureFusion(2048, 512, 128, 7, bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(cff_24.get_name(), cff_24, 128)
cff_12 = CascadeFeatureFusion(128, 64, 128, 2, bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(cff_12.get_name(), cff_12, 128)
up1 = Upsample(scale_factor=2, mode='bilinear')
self.add_block_list(up1.get_name(), up1, 128)
conv_cls = nn.Conv2d(128, self.class_number, 1, bias=False)
self.add_block_list(LayerType.Convolutional, conv_cls, self.class_number)
up2 = Upsample(scale_factor=4, mode='bilinear')
self.add_block_list(up2.get_name(), up2, self.class_number)
def create_loss(self, input_dict=None):
self.lossList = []
loss = CrossEntropy2d(ignore_index=250)
self.add_block_list(LossType.CrossEntropy2d, loss, self.block_out_channels[-1])
self.lossList.append(loss)
def forward(self, x):
base_outputs = []
layer_outputs = []
output = []
input_datas = []
data_index = 0
index = 3
for key, block in self._modules.items():
if ICNetBlockName.InputDownSample in key:
input_datas = block(x)
elif BlockType.BaseNet in key:
base_outputs = block(input_datas[data_index])
data_index += 1
else:
if LayerType.RouteLayer in key:
x = block(layer_outputs, base_outputs)
elif ICNetBlockName.CascadeFeatureFusion in key:
x = block(layer_outputs[-1], layer_outputs[index])
index -= 1
elif LossType.CrossEntropy2d in key:
output.append(x)
else:
x = block(x)
layer_outputs.append(x)
print(key, layer_outputs[-1].shape)
return output
class EncNetSeg(BaseModel):
def __init__(self, data_channel=3, class_num=150):
super().__init__()
self.set_name(ModelName.EncNetSeg)
self.data_channel = data_channel
self.class_number = class_num
self.is_jpu = True
self.lateral = False
self.is_se_loss = True
self.is_aux = True
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
def create_block_list(self):
self.clear_list()
backbone = self.factory.get_base_model(BackboneName.ResNet50)
base_out_channels = backbone.get_outchannel_list()
self.add_block_list(BlockType.BaseNet, backbone, base_out_channels[-1])
if self.is_jpu:
jup = JPUBlock(layers='4,8,14,17',
in_planes=(512, 1024, 2048),
width=512,
bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(jup.get_name(), jup, 512 + 512 + 512 + 512)
self.enc_head(2048, base_out_channels)
self.create_loss()
if self.is_aux:
route = RouteLayer('14')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in route.layers
])
self.add_block_list(route.get_name(), route, output_channel)
fcn_head = FCNHeadBlock(1024,
self.class_number,
16,
bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(fcn_head.get_name(), fcn_head,
self.class_number)
def enc_head(self, in_channels, base_out_channels):
if self.is_jpu:
conv1 = ConvBNActivationBlock(in_channels=in_channels,
out_channels=512,
kernel_size=1,
bias=False,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(conv1.get_name(), conv1, 512)
else:
conv1 = ConvBNActivationBlock(in_channels=in_channels,
out_channels=512,
kernel_size=3,
padding=1,
bias=False,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(conv1.get_name(), conv1, 512)
if self.lateral:
route1 = RouteLayer('8')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in route1.layers
])
self.add_block_list(route1.get_name(), route1, output_channel)
connect1 = ConvBNActivationBlock(
in_channels=output_channel,
out_channels=512,
kernel_size=1,
bias=False,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(connect1.get_name(), connect1, 512)
route2 = RouteLayer('14')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in route2.layers
])
self.add_block_list(route2.get_name(), route2, output_channel)
connect2 = ConvBNActivationBlock(
in_channels=output_channel,
out_channels=512,
kernel_size=1,
bias=False,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(connect2.get_name(), connect2, 512)
route3 = RouteLayer('-5,-3,-1')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in route2.layers
])
self.add_block_list(route3.get_name(), route3, output_channel)
fusion = ConvBNActivationBlock(in_channels=output_channel,
out_channels=512,
kernel_size=3,
padding=1,
bias=False,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(fusion.get_name(), fusion, 512)
encmodule = EncBlock(in_channels=512,
nclass=self.class_number,
se_loss=self.is_se_loss,
bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(encmodule.get_name(), encmodule, 512)
dropout = nn.Dropout2d(0.1, False)
self.add_block_list(LayerType.Dropout, dropout,
self.block_out_channels[-1])
conv2 = nn.Conv2d(self.block_out_channels[-1], self.class_number, 1)
self.add_block_list(LayerType.Convolutional, conv2, self.class_number)
up = Upsample(scale_factor=8, mode='bilinear')
self.add_block_list(up.get_name(), up, self.class_number)
def create_loss(self, input_dict=None):
self.lossList = []
loss = EncNetLoss(self.class_number,
se_loss=self.is_se_loss,
aux=self.is_aux,
ignore_index=250)
self.add_block_list(LossType.EncNetLoss, loss,
self.block_out_channels[-1])
self.lossList.append(loss)
def forward(self, x):
base_outputs = []
layer_outputs = []
output = []
se_loss = None
aux_loss = None
for key, block in self._modules.items():
if BlockType.BaseNet in key:
base_outputs = block(x)
x = base_outputs[-1]
elif LayerType.RouteLayer in key:
x = block(layer_outputs, base_outputs)
elif EncNetBlockName.JPUBlock in key:
x = block(layer_outputs, base_outputs)
elif EncNetBlockName.EncBlock in key:
x, se_loss = block(x)
elif LossType.EncNetLoss in key:
output.append(x)
elif EncNetBlockName.FCNHeadBlock in key:
x = block(x)
aux_loss = x
else:
x = block(x)
layer_outputs.append(x)
print(key, x.shape)
output.append(aux_loss)
output.append(se_loss)
return output
class RefineNetSeg(BaseModel):
def __init__(self, data_channel=3, class_num=2):
super().__init__()
self.set_name(ModelName.RefineNetSeg)
self.data_channel = data_channel
self.class_number = class_num
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
def create_block_list(self):
self.block_out_channels = []
self.index = 0
backbone = self.factory.get_base_model(BackboneName.ResNet101)
base_out_channels = backbone.get_outchannel_list()
self.add_block_list(BlockType.BaseNet, backbone, base_out_channels[-1])
dropout1 = nn.Dropout(p=0.5)
self.add_block_list(LayerType.Dropout, dropout1,
self.block_out_channels[-1])
layer1 = RouteLayer('31')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in layer1.layers
])
self.add_block_list(layer1.get_name(), layer1, output_channel)
dropout2 = nn.Dropout(p=0.5)
self.add_block_list(LayerType.Dropout, dropout2,
self.block_out_channels[-1])
layer2 = RouteLayer('-3')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in layer1.layers
])
self.add_block_list(layer2.get_name(), layer2, output_channel)
conv1 = nn.Conv2d(2048, 512, 1, bias=False)
self.add_block_list(LayerType.Convolutional, conv1, 512)
layer3 = RouteLayer('-3')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in layer1.layers
])
self.add_block_list(layer3.get_name(), layer3, output_channel)
conv2 = nn.Conv2d(1024, 256, 1, bias=False)
self.add_block_list(LayerType.Convolutional, conv2, 256)
refinenet1 = RefineNetBlock(512,
512,
4,
activation_name=self.activation_name)
output_channel = 256
self.add_block_list(refinenet1.get_name(), refinenet1, output_channel)
layer4 = RouteLayer('8')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in layer1.layers
])
self.add_block_list(layer4.get_name(), layer4, output_channel)
conv3 = nn.Conv2d(512, 256, 1, bias=False)
self.add_block_list(LayerType.Convolutional, conv3, 256)
refinenet2 = RefineNetBlock(256,
256,
4,
activation_name=self.activation_name)
output_channel = 256
self.add_block_list(refinenet2.get_name(), refinenet2, output_channel)
layer5 = RouteLayer('4')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in layer1.layers
])
self.add_block_list(layer5.get_name(), layer5, output_channel)
conv4 = nn.Conv2d(256, 256, 1, bias=False)
self.add_block_list(LayerType.Convolutional, conv4, 256)
refinenet3 = RefineNetBlock(256,
256,
4,
activation_name=self.activation_name)
output_channel = 256
self.add_block_list(refinenet3.get_name(), refinenet3, output_channel)
crp = CRPBlock(256, 256, 4)
output_channel = 256
self.add_block_list(crp.get_name(), crp, output_channel)
conv5 = nn.Conv2d(256,
self.class_number,
kernel_size=3,
stride=1,
padding=1,
bias=True)
self.add_block_list(LayerType.Convolutional, conv5, self.class_number)
layer6 = Upsample(scale_factor=4, mode='bilinear')
self.add_block_list(layer6.get_name(), layer6,
self.block_out_channels[-1])
self.create_loss()
def create_loss(self, input_dict=None):
self.lossList = []
loss = CrossEntropy2d(ignore_index=250)
self.add_block_list(LossType.CrossEntropy2d, loss,
self.block_out_channels[-1])
self.lossList.append(loss)
def forward(self, x):
base_outputs = []
layer_outputs = []
output = []
for key, block in self._modules.items():
if BlockType.BaseNet in key:
base_outputs = block(x)
x = base_outputs[-1]
elif LayerType.MultiplyLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.AddLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.RouteLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.ShortcutLayer in key:
x = block(layer_outputs)
elif RefineNetBlockName.RefineNetBlock in key:
x = block(layer_outputs[-3], layer_outputs[-1])
elif LossType.CrossEntropy2d in key:
output.append(x)
else:
x = block(x)
layer_outputs.append(x)
return output
class DeepLabV3(BaseModel):
def __init__(self, data_channel=3, class_num=19):
super().__init__()
self.set_name(ModelName.DeepLabV3)
self.data_channel = data_channel
self.class_number = class_num
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
self.loss_flag = -1
def create_block_list(self):
self.clear_list()
self.lossList = []
backbone = self.factory.get_base_model(BackboneName.ResNet50)
base_out_channels = backbone.get_outchannel_list()
self.add_block_list(BlockType.BaseNet, backbone, base_out_channels[-1])
head = ASPPBlock(2048,
bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(head.get_name(), head, 512)
out_conv = nn.Conv2d(512,
self.class_number,
kernel_size=1,
stride=1,
padding=0,
bias=True)
self.add_block_list(LayerType.Convolutional, out_conv,
self.class_number)
up = Upsample(scale_factor=32, mode='bilinear')
self.add_block_list(up.get_name(), up, self.block_out_channels[-1])
self.loss_flag = 0
self.clear_list()
self.dsn_head(base_out_channels)
self.loss_flag = 1
self.clear_list()
def dsn_head(self, base_out_channels):
route = RouteLayer('14')
output_channel = sum([
base_out_channels[i] if i >= 0 else self.block_out_channels[i]
for i in route.layers
])
self.add_block_list(route.get_name(), route, output_channel)
conv1 = ConvBNActivationBlock(in_channels=output_channel,
out_channels=512,
kernel_size=3,
stride=1,
padding=1,
bias=True,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(conv1.get_name(), conv1, 512)
dropout = nn.Dropout2d(0.1)
self.add_block_list(LayerType.Dropout, dropout,
self.block_out_channels[-1])
out_conv = nn.Conv2d(512,
self.class_number,
kernel_size=1,
stride=1,
padding=0,
bias=True)
self.add_block_list(LayerType.Convolutional, out_conv,
self.class_number)
up1 = Upsample(scale_factor=16, mode='bilinear')
self.add_block_list(up1.get_name(), up1, self.block_out_channels[-1])
def create_loss(self, input_dict=None):
if self.loss_flag == 0:
loss = CrossEntropy2d(ignore_index=250)
self.add_block_list(LossType.CrossEntropy2d, loss,
self.block_out_channels[-1])
self.lossList.append(loss)
elif self.loss_flag == 1:
loss = CrossEntropy2d(ignore_index=250)
self.add_block_list(LossType.CrossEntropy2d, loss,
self.block_out_channels[-1])
self.lossList.append(loss)
def forward(self, x):
base_outputs = []
layer_outputs = []
output = []
for key, block in self._modules.items():
if BlockType.BaseNet in key:
base_outputs = block(x)
x = base_outputs[-1]
elif LayerType.RouteLayer in key:
x = block(layer_outputs, base_outputs)
elif LossType.CrossEntropy2d in key:
output.append(x)
else:
x = block(x)
layer_outputs.append(x)
print(key, x.shape)
return output
def backbone_model_print(model_name):
backbone_factory = BackboneFactory()
input_x = torch.randn(1, 3, 32, 32)
backbone = backbone_factory.get_base_model(model_name)
backbone.print_block_name()
def __init__(self, input_size=(352, 640)):
self.backbone_factory = BackboneFactory()
self.model_factory = ModelFactory()
self.converter = TorchConvertOnnx()
self.input_size = input_size # w * h
class GhostNetCls(BaseModel):
def __init__(self, data_channel=3, class_num=100):
super().__init__()
self.set_name(ModelName.GhostNetCls)
self.data_channel = data_channel
self.class_number = class_num
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
def create_block_list(self):
self.clear_list()
backbone = self.factory.get_base_model(BackboneName.GhostNet,
self.data_channel)
base_out_channels = backbone.get_outchannel_list()
self.add_block_list(BlockType.BaseNet, backbone, base_out_channels[-1])
avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.add_block_list(LayerType.GlobalAvgPool, avgpool,
base_out_channels[-1])
output_channel = 1280
layer1 = FcLayer(base_out_channels[-1], output_channel)
self.add_block_list(layer1.get_name(), layer1, output_channel)
layer2 = NormalizeLayer(bn_name=NormalizationType.BatchNormalize1d,
out_channel=output_channel)
self.add_block_list(layer2.get_name(), layer2, output_channel)
layer3 = ActivationLayer(self.activation_name, inplace=False)
self.add_block_list(layer3.get_name(), layer3, output_channel)
layer4 = nn.Dropout(0.2)
self.add_block_list(LayerType.Dropout, layer4, output_channel)
layer5 = nn.Linear(output_channel, self.class_number)
self.add_block_list(LayerType.FcLinear, layer5, self.class_number)
self.create_loss()
def create_loss(self, input_dict=None):
self.lossList = []
loss = CrossEntropy2d(ignore_index=250)
self.add_block_list(LossType.CrossEntropy2d, loss,
self.block_out_channels[-1])
self.lossList.append(loss)
def forward(self, x):
base_outputs = []
layer_outputs = []
output = []
for key, block in self._modules.items():
if BlockType.BaseNet in key:
base_outputs = block(x)
x = base_outputs[-1]
elif LayerType.RouteLayer in key:
x = block(layer_outputs, base_outputs)
elif LayerType.ShortcutLayer in key:
x = block(layer_outputs)
elif LossType.CrossEntropy2d in key:
output.append(x)
else:
x = block(x)
layer_outputs.append(x)
return output
class PSPNetSeg(BaseModel):
def __init__(self, data_channel=3, class_num=2):
super().__init__()
self.set_name(ModelName.PSPNetSeg)
self.data_channel = data_channel
self.class_number = class_num
self.is_jpu = True
self.bn_name = NormalizationType.BatchNormalize2d
self.activation_name = ActivationType.ReLU
self.factory = BackboneFactory()
self.create_block_list()
def create_block_list(self):
self.clear_list()
backbone = self.factory.get_base_model(BackboneName.ResNet101)
base_out_channels = backbone.get_outchannel_list()
self.add_block_list(BlockType.BaseNet, backbone, base_out_channels[-1])
if self.is_jpu:
jup = JPUBlock(layers='4,8,31,34',
in_planes=(512, 1024, 2048),
width=512,
bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(jup.get_name(), jup, 512 + 512 + 512 + 512)
scale_factor = 8
else:
scale_factor = 32
psp = PyramidPooling(2048,
bn_name=self.bn_name,
activation_name=self.activation_name)
self.add_block_list(psp.get_name(), psp, 2048 * 2)
conv1 = ConvBNActivationBlock(in_channels=2048 * 2,
out_channels=512,
kernel_size=3,
padding=1,
bias=False,
bnName=self.bn_name,
activationName=self.activation_name)
self.add_block_list(conv1.get_name(), conv1, 512)
dropout = nn.Dropout(0.1)
self.add_block_list(LayerType.Dropout, dropout,
self.block_out_channels[-1])
conv2 = nn.Conv2d(512, self.class_number, 1)
self.add_block_list(LayerType.Convolutional, conv2, self.class_number)
layer = Upsample(scale_factor=scale_factor, mode='bilinear')
self.add_block_list(layer.get_name(), layer,
self.block_out_channels[-1])
self.create_loss()
def create_loss(self, input_dict=None):
self.lossList = []
loss = CrossEntropy2d(ignore_index=250)
self.add_block_list(LossType.CrossEntropy2d, loss,
self.block_out_channels[-1])
self.lossList.append(loss)
def forward(self, x):
base_outputs = []
layer_outputs = []
output = []
for key, block in self._modules.items():
if BlockType.BaseNet in key:
base_outputs = block(x)
x = base_outputs[-1]
elif LayerType.RouteLayer in key:
x = block(layer_outputs, base_outputs)
elif EncNetBlockName.JPUBlock in key:
x = block(layer_outputs, base_outputs)
elif LossType.CrossEntropy2d in key:
output.append(x)
else:
x = block(x)
layer_outputs.append(x)
print(key, x.shape)
return output