|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from .Attention import Block, CrossBlock |
| from util.util import PositionalEncoding, PosCNN |
| from .blocks import Conv2dBlock, ResBlocks, ActFirstResBlock |
| from .Unifront import UnifontModule |
| from params import * |
|
|
| class Generator(nn.Module): |
|
|
| def __init__( |
| self, |
| arg = None, |
| embed_dim=256, |
| depth=3, |
| num_heads=4, |
| mlp_ratio=4, |
| drop=0.0, |
| norm_layer=nn.LayerNorm, |
| max_num_patch=100, |
| ): |
| super().__init__() |
|
|
| |
| |
| self.layer_norm = None |
| self.grid_size = None |
| self.embed_dim = [256, 256,128, 128, 64, 32, 16] |
| num_block = 4 |
| self.pos_enc = PositionalEncoding(embed_dim, drop, max_num_patch) |
| self.query_embed = UnifontModule( |
| embed_dim, |
| ALPHABET, |
| input_type="unifont", |
| linear=True, |
| ) |
|
|
| """Block 1""" |
| index = 1 |
| self.blocks_2 = nn.ModuleList( |
| [ |
| CrossBlock( |
| self.embed_dim[index], |
| num_heads=num_heads, |
| mlp_ratio=mlp_ratio, |
| qkv_bias=True, |
| norm_layer=norm_layer, |
| ) |
| for i in range(depth+2) |
| ] |
| ) |
| self.layer_norm2 = nn.LayerNorm(self.embed_dim[index]) |
| self.tRGB_1 = nn.Sequential( |
| nn.Conv2d(self.embed_dim[index], self.embed_dim[num_block], 3, 1, 1) |
| ) |
| self.conv_1 = self._make_upsample_block(self.embed_dim[index], self.embed_dim[index+1]) |
|
|
| """Block 2""" |
| index+=1 |
| self.blocks_3 = nn.ModuleList( |
| [ |
| Block( |
| dim=self.embed_dim[index], |
| num_heads=num_heads, |
| mlp_ratio=mlp_ratio, |
| qkv_bias=True, |
| norm_layer=norm_layer, |
| ) |
| for i in range(depth) |
| ] |
| ) |
| self.layer_norm3 = nn.LayerNorm(self.embed_dim[index]) |
| self.tRGB_2 = nn.Sequential( |
| nn.Conv2d(self.embed_dim[index], self.embed_dim[num_block], 3, 1, 1) |
| ) |
| self.conv_2 = self._make_upsample_block(self.embed_dim[index], self.embed_dim[index+1]) |
|
|
|
|
| """Block 3""" |
| index+=1 |
| self.blocks_4 = nn.ModuleList( |
| [ |
| Block( |
| dim=self.embed_dim[index], |
| num_heads=num_heads, |
| mlp_ratio=mlp_ratio, |
| qkv_bias=True, |
| norm_layer=norm_layer, |
| ) |
| for i in range(depth) |
| ] |
| ) |
| self.layer_norm4 = nn.LayerNorm(self.embed_dim[index]) |
| self.tRGB_3 = nn.Sequential( |
| nn.Conv2d(self.embed_dim[index], self.embed_dim[num_block], 3, 1, 1) |
| ) |
| self.conv_3 = self._make_upsample_block(self.embed_dim[index], self.embed_dim[index+1]) |
|
|
|
|
| """Block 4""" |
| index+=1 |
| self.blocks_5 = nn.ModuleList( |
| [ |
| Block( |
| dim=self.embed_dim[index], |
| num_heads=num_heads, |
| mlp_ratio=mlp_ratio, |
| qkv_bias=True, |
| norm_layer=norm_layer, |
| ) |
| for i in range(depth) |
| ] |
| ) |
| self.layer_norm5 = nn.LayerNorm(self.embed_dim[index]) |
|
|
| self.pos_block = nn.ModuleList([PosCNN(i, i) for i in self.embed_dim]) |
| self.norm = norm_layer(embed_dim, elementwise_affine=True) |
| self.noise = torch.distributions.Normal( |
| loc=torch.tensor([0.0]), scale=torch.tensor([1.0]) |
| ) |
|
|
| self.deconv = nn.Sequential( |
| ResBlocks( |
| 2, self.embed_dim[index], norm="in", activation="relu", pad_type="reflect" |
| ), |
| nn.Upsample(scale_factor=2), |
| Conv2dBlock( |
| self.embed_dim[index], |
| self.embed_dim[index + 1], |
| 3, |
| 1, |
| 1, |
| norm="in", |
| activation="none", |
| pad_type="reflect", |
| ), |
| Conv2dBlock( |
| self.embed_dim[5], |
| self.embed_dim[5], |
| 5, |
| 1, |
| 2, |
| norm="in", |
| activation="relu", |
| pad_type="reflect", |
| ), |
| Conv2dBlock( |
| self.embed_dim[5], |
| 1, |
| 7, |
| 1, |
| 3, |
| norm="none", |
| activation="tanh", |
| pad_type="reflect", |
| ), |
| ) |
| self.initialize_weights() |
|
|
| def _make_upsample_block(self, in_dim, out_dim): |
| return nn.Sequential( |
| nn.Upsample(scale_factor=2), |
| Conv2dBlock(in_dim, out_dim, 3, 1, 1, norm="in", activation="none", pad_type="reflect"), |
| Conv2dBlock(out_dim, out_dim, 3, 1, 1, norm="in", activation="relu", pad_type="reflect"), |
| ) |
|
|
| def initialize_weights(self): |
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, m): |
| if isinstance(m, nn.Linear): |
| nn.init.xavier_uniform_(m.weight) |
| if m.bias is not None: |
| nn.init.constant_(m.bias, 0) |
| elif isinstance(m, nn.LayerNorm): |
| nn.init.constant_(m.bias, 0) |
| nn.init.constant_(m.weight, 1.0) |
|
|
| def _generate_features(self, src, tgt): |
| b = src.size(0) |
| start_h = 2 |
| start_w = tgt.size(1) |
|
|
| src = src |
| tmp = self.query_embed(tgt.clone()) |
| tgt = self.pos_enc(self.query_embed(tgt)) |
|
|
| stack_output = [] |
| for blk in self.blocks_2: |
| tgt = blk(tgt, src) |
| stack_output.append(tgt) |
| h2 = stack_output[-1] |
|
|
| tgt = torch.cat([h2, tmp], dim=1) |
| tgt = self.layer_norm2(tgt) |
| tgt = tgt.permute(0, 2, 1).view(b, self.embed_dim[1], start_h, start_w) |
| x_1 = self.tRGB_1(tgt) |
|
|
| tgt = self.conv_1(tgt) |
| b, c, h, w = tgt.shape |
| tgt = tgt.view(b, c, -1).permute(0, 2, 1) |
|
|
| for j, blk in enumerate(self.blocks_3): |
| tgt = blk(tgt) |
| if j == 0: |
| tgt = self.pos_block[2](tgt, h, w) |
| tgt = self.layer_norm3(tgt).permute(0, 2, 1).view(b, self.embed_dim[2], h, w) |
| x_2 = self.tRGB_2(tgt) |
|
|
| tgt = self.conv_2(tgt) |
| b, c, h, w = tgt.shape |
| tgt = tgt.view(b, c, -1).permute(0, 2, 1) |
|
|
| for j, blk in enumerate(self.blocks_4): |
| tgt = blk(tgt) |
| if j == 0: |
| tgt = self.pos_block[3](tgt, h, w) |
| tgt = self.layer_norm4(tgt).permute(0, 2, 1).view(b, self.embed_dim[3], h, w) |
| x_3 = self.tRGB_3(tgt) |
|
|
| tgt = self.conv_3(tgt) |
| b, c, h, w = tgt.shape |
| tgt = tgt.view(b, c, -1).permute(0, 2, 1) |
|
|
| for j, blk in enumerate(self.blocks_5): |
| tgt = blk(tgt) |
| if j == 0: |
| tgt = self.pos_block[4](tgt, h, w) |
| tgt = self.layer_norm5(tgt).permute(0, 2, 1).view(b, self.embed_dim[4], h, w) |
|
|
| fused = ( |
| F.interpolate(x_1, scale_factor=8) |
| + F.interpolate(x_2, scale_factor=4) |
| + F.interpolate(x_3, scale_factor=2) |
| + tgt |
| ) |
| noise = self.noise.sample(fused.size()).squeeze(-1).to(fused.device) |
| return fused + noise |
|
|
| def forward(self, src_w, tgt): |
| features = self._generate_features(src_w, tgt) |
| return self.deconv(features) |
|
|
| def Eval(self, xw, QRS): |
| outputs = [] |
| for i in range(QRS.shape[1]): |
| tgt = QRS[:, i, :].squeeze(1) |
| features = self._generate_features(xw, tgt) |
| outputs.append(self.deconv(features).detach()) |
| return outputs |
