Update models/DehazeFormer.py

models/DehazeFormer.py

@@ -471,4 +471,117 @@ class MCT(nn.Module):
 		x_d = F.interpolate(x, (self.ts, self.ts), mode='area')
 		param = self.basenet(x_d)
 		out = self.mapping(x, param)
-		return out
+		return out
+class DehazeFormer(nn.Module):
+	def __init__(self, in_chans=3, out_chans=4, window_size=8,
+				 embed_dims=[24, 48, 96, 48, 24],
+				 mlp_ratios=[2., 4., 4., 2., 2.],
+				 depths=[16, 16, 16, 8, 8],
+				 num_heads=[2, 4, 6, 1, 1],
+				 attn_ratio=[1/4, 1/2, 3/4, 0, 0],
+				 conv_type=['DWConv', 'DWConv', 'DWConv', 'DWConv', 'DWConv'],
+				 norm_layer=[RLN, RLN, RLN, RLN, RLN]):
+		super(DehazeFormer, self).__init__()
+
+		# setting
+		self.patch_size = 4
+		self.window_size = window_size
+		self.mlp_ratios = mlp_ratios
+
+		# split image into non-overlapping patches
+		self.patch_embed = PatchEmbed(
+			patch_size=1, in_chans=in_chans, embed_dim=embed_dims[0], kernel_size=3)
+
+		# backbone
+		self.layer1 = BasicLayer(network_depth=sum(depths), dim=embed_dims[0], depth=depths[0],
+					   			 num_heads=num_heads[0], mlp_ratio=mlp_ratios[0],
+					   			 norm_layer=norm_layer[0], window_size=window_size,
+					   			 attn_ratio=attn_ratio[0], attn_loc='last', conv_type=conv_type[0])
+
+		self.patch_merge1 = PatchEmbed(
+			patch_size=2, in_chans=embed_dims[0], embed_dim=embed_dims[1])
+
+		self.skip1 = nn.Conv2d(embed_dims[0], embed_dims[0], 1)
+
+		self.layer2 = BasicLayer(network_depth=sum(depths), dim=embed_dims[1], depth=depths[1],
+								 num_heads=num_heads[1], mlp_ratio=mlp_ratios[1],
+								 norm_layer=norm_layer[1], window_size=window_size,
+								 attn_ratio=attn_ratio[1], attn_loc='last', conv_type=conv_type[1])
+
+		self.patch_merge2 = PatchEmbed(
+			patch_size=2, in_chans=embed_dims[1], embed_dim=embed_dims[2])
+
+		self.skip2 = nn.Conv2d(embed_dims[1], embed_dims[1], 1)
+
+		self.layer3 = BasicLayer(network_depth=sum(depths), dim=embed_dims[2], depth=depths[2],
+								 num_heads=num_heads[2], mlp_ratio=mlp_ratios[2],
+								 norm_layer=norm_layer[2], window_size=window_size,
+								 attn_ratio=attn_ratio[2], attn_loc='last', conv_type=conv_type[2])
+
+		self.patch_split1 = PatchUnEmbed(
+			patch_size=2, out_chans=embed_dims[3], embed_dim=embed_dims[2])
+
+		assert embed_dims[1] == embed_dims[3]
+		self.fusion1 = SKFusion(embed_dims[3])
+
+		self.layer4 = BasicLayer(network_depth=sum(depths), dim=embed_dims[3], depth=depths[3],
+								 num_heads=num_heads[3], mlp_ratio=mlp_ratios[3],
+								 norm_layer=norm_layer[3], window_size=window_size,
+								 attn_ratio=attn_ratio[3], attn_loc='last', conv_type=conv_type[3])
+
+		self.patch_split2 = PatchUnEmbed(
+			patch_size=2, out_chans=embed_dims[4], embed_dim=embed_dims[3])
+
+		assert embed_dims[0] == embed_dims[4]
+		self.fusion2 = SKFusion(embed_dims[4])
+
+		self.layer5 = BasicLayer(network_depth=sum(depths), dim=embed_dims[4], depth=depths[4],
+					   			 num_heads=num_heads[4], mlp_ratio=mlp_ratios[4],
+					   			 norm_layer=norm_layer[4], window_size=window_size,
+					   			 attn_ratio=attn_ratio[4], attn_loc='last', conv_type=conv_type[4])
+
+		# merge non-overlapping patches into image
+		self.patch_unembed = PatchUnEmbed(
+			patch_size=1, out_chans=out_chans, embed_dim=embed_dims[4], kernel_size=3)
+
+
+	def check_image_size(self, x):
+		# NOTE: for I2I test
+		_, _, h, w = x.size()
+		mod_pad_h = (self.patch_size - h % self.patch_size) % self.patch_size
+		mod_pad_w = (self.patch_size - w % self.patch_size) % self.patch_size
+		x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), 'reflect')
+		return x
+
+	def forward_features(self, x):
+		x = self.patch_embed(x)
+		x = self.layer1(x)
+		skip1 = x
+
+		x = self.patch_merge1(x)
+		x = self.layer2(x)
+		skip2 = x
+
+		x = self.patch_merge2(x)
+		x = self.layer3(x)
+		x = self.patch_split1(x)
+
+		x = self.fusion1([x, self.skip2(skip2)]) + x
+		x = self.layer4(x)
+		x = self.patch_split2(x)
+
+		x = self.fusion2([x, self.skip1(skip1)]) + x
+		x = self.layer5(x)
+		x = self.patch_unembed(x)
+		return x
+
+	def forward(self, x):
+		H, W = x.shape[2:]
+		x = self.check_image_size(x)
+
+		feat = self.forward_features(x)
+		K, B = torch.split(feat, (1, 3), dim=1)
+
+		x = K * x - B + x
+		x = x[:, :, :H, :W]
+		return x