def __init__(self, N, M, **kwargs):
super().__init__()
self.g_a_conv1 = conv(3, N)
self.g_a_gdn1 = GDN(N)
self.g_a_conv2 = conv(N, N)
self.g_a_gdn2 = GDN(N)
self.g_a_conv3 = conv(N, N)
self.g_a_gdn3 = GDN(N)
self.g_a_conv4 = conv(N, M)
def __init__(self, N, M):
super().__init__()
self.encode_hyper = nn.Sequential(
#先abs 再输入进来
conv(in_channels=M, out_channels=N, kernel_size=5, stride=1),
nn.ReLU(),
conv(in_channels=N, out_channels=N, kernel_size=5), #stride = 2
nn.ReLU(),
conv(in_channels=N, out_channels=N, kernel_size=5), #stride = 2
#出去后接bottleneck
)
def __init__(self, N, M, **kwargs):
super().__init__()
self.pre_conv = conv(6,3,stride=1) #不缩放！
self.pre_gdn = GDN(3)
self.g_a_conv1 = conv(3, N)
self.g_a_gdn1 = GDN(N)
self.g_a_conv2 = conv(N, N)
self.g_a_gdn2 = GDN(N)
self.g_a_conv3 = conv(N, N)
self.g_a_gdn3 = GDN(N)
self.g_a_conv4 = conv(N, M)
def __init__(self, N, scale_factor, F,
C): #N=input_channels,scale_factor:Upsample_factor,
super(cost_volume, self).__init__()
self.N = N
self.scale_factor = scale_factor
self.F = F
self.F0 = F // 3
self.C = C
self.model1 = nn.Sequential(
conv(2 * self.N, self.N, kernel_size=5, stride=1),
nn.GroupNorm(num_channels=self.N, num_groups=4),
nn.ReLU(),
conv(self.N, self.N, kernel_size=5, stride=1),
nn.GroupNorm(num_channels=self.N, num_groups=4),
nn.ReLU(),
)
self.upsample_layer = nn.UpsamplingBilinear2d(
scale_factor=self.scale_factor)
self.model2 = nn.Sequential(
#先取 tensor[0]需要4维的 之后再恢复5维，输入到model2里来 #or:torch.squeeze(x) torch.unsqueeze(x,dim=0)
# nn.UpsamplingBilinear2d(scale_factor=scale_factor),
nn.Conv3d(in_channels=self.F0,
out_channels=self.F0,
kernel_size=5,
stride=1,
padding=5 // 2),
nn.GroupNorm(num_channels=self.F0, num_groups=1),
nn.ReLU(),
nn.Conv3d(in_channels=self.F0,
out_channels=self.F0,
kernel_size=5,
stride=1,
padding=5 // 2),
nn.GroupNorm(num_channels=self.F0, num_groups=1),
nn.ReLU(),
#输出后记得先转2d再concat
)
self.model3 = nn.Sequential(
conv((self.F0 * self.C + self.N), self.N, kernel_size=5, stride=1),
nn.GroupNorm(num_channels=self.N, num_groups=4),
nn.ReLU(),
conv(self.N, self.N, kernel_size=5, stride=1),
nn.GroupNorm(num_channels=self.N, num_groups=4),
nn.ReLU(),
#最后变成C通道，即为disparity
conv(self.N, self.C, kernel_size=5, stride=1),
#softmax over dispartiy dimenstion (C is disparity maxnum)
# nn.functional.softmax(dim=-3),
)
def __init__(self, N, M, K): #K表示GMM对应正态分布个数
super().__init__()
self.N = N
self.M = M
self.K = K
# 每个都是上采样4倍才行
self.gmm_sigma = nn.Sequential(
deconv(in_channels=N, out_channels=N, kernel_size=5
), #stride=2,padding=kernel_size//2,output_padding=stride-1
nn.ReLU(),
deconv(
in_channels=N, out_channels=N, kernel_size=5
), # stride=2,padding=kernel_size//2,output_padding=stride-1
nn.ReLU(),
conv(in_channels=N, out_channels=(M * K), kernel_size=5,
stride=1), #padding=kernel_size//2
nn.ReLU(),
)
self.gmm_means = nn.Sequential(
deconv(in_channels=N, out_channels=N, kernel_size=5),
# stride=2,padding=kernel_size//2,output_padding=stride-1
nn.LeakyReLU(),
deconv(in_channels=N, out_channels=N, kernel_size=5),
# stride=2,padding=kernel_size//2,output_padding=stride-1
nn.LeakyReLU(),
conv(in_channels=N, out_channels=(M * K), kernel_size=5,
stride=1), # padding=kernel_size//2
)
self.gmm_weights = nn.Sequential(
deconv(in_channels=N, out_channels=N, kernel_size=5),
# stride=2,padding=kernel_size//2,output_padding=stride-1
nn.LeakyReLU(),
deconv(in_channels=N, out_channels=(M * K), kernel_size=5),
# stride=2,padding=kernel_size//2,output_padding=stride-1
# nn.MaxPool2d(kernel_size=(H//16,W//16)), # ?? 换图像分辨率就要换模型了
spatial_pool2d(),
nn.LeakyReLU(),
conv(in_channels=(M * K),
out_channels=(M * K),
kernel_size=1,
stride=1), # padding=kernel_size//2
#出去后要接一个softmax层，表示概率！！
)
def __init__(self, N=128):
super().__init__(entropy_bottleneck_channels=N)
self.encode = nn.Sequential(
conv(3, N, kernel_size=9, stride=4),
GDN(N),
conv(N, N),
GDN(N),
conv(N, N),
)
self.decode = nn.Sequential(
deconv(N, N),
GDN(N, inverse=True),
deconv(N, N),
GDN(N, inverse=True),
deconv(N, 3, kernel_size=9, stride=4),
)
def __init__(self,N,M,K): #K表示GMM对应正态分布个数
super().__init__()
self.N = N
self.M = M
self.K = K
self.upsample_layer = nn.UpsamplingBilinear2d(scale_factor=4) #固定为4倍 因为z和y分辨率本身就差4倍
# 输入与y1分辨率相同，但通道数是2倍
self.gmm_sigma = nn.Sequential(
conv(in_channels=(N+M),out_channels=N,kernel_size=5,stride=1),
nn.ReLU(),
conv(in_channels=N,out_channels=N,kernel_size=5,stride=1),
nn.ReLU(),
conv(in_channels=N,out_channels=(M*K),kernel_size=5,stride=1), #padding=kernel_size//2
nn.ReLU(),
)
self.gmm_means = nn.Sequential(
conv(in_channels=(N+M),out_channels=N,kernel_size=5,stride=1),
nn.LeakyReLU(),
conv(in_channels=N,out_channels=N,kernel_size=5,stride=1),
# stride=2,padding=kernel_size//2,output_padding=stride-1
nn.LeakyReLU(),
conv(in_channels=N, out_channels=(M * K), kernel_size=5, stride=1), # padding=kernel_size//2
)
self.gmm_weights = nn.Sequential(
conv(in_channels=(N+M),out_channels=N,kernel_size=5,stride=1),
nn.LeakyReLU(),
conv(in_channels=N,out_channels= (M*K),kernel_size=5,stride=1),
# nn.MaxPool2d(kernel_size=(H//16,W//16)), # ?? 换图像分辨率就要换模型了
spatial_pool2d(),
nn.LeakyReLU(),
conv(in_channels=(M*K), out_channels=(M * K), kernel_size=1, stride=1), # padding=kernel_size//2
#出去后要接一个softmax层，表示概率！！
)
def __init__(self, M, F, C):
super().__init__()
self.F = F
self.F0 = F // 3
self.M = M
self.C = C
# 定义新网络 global_context (输入y1)
self.global_net = nn.Sequential(
# nn.Conv2d(M,(F*C),kernel_size=5,stride=1,padding=kernel_size // 2) 等价于下行
conv(M, F * C, kernel_size=5,
stride=1), # padding = kernel_size // 2
nn.GroupNorm(num_channels=F * C, num_groups=F),
nn.ReLU(),
conv(F * C, F * C, kernel_size=5,
stride=1), # padding = kernel_size // 2
nn.GroupNorm(num_channels=F * C, num_groups=F),
nn.ReLU(),
conv(F * C, F * C, kernel_size=5,
stride=1), # padding = kernel_size // 2
nn.GroupNorm(num_channels=F * C, num_groups=F),
nn.ReLU(),
conv(F * C, F * C, kernel_size=5,
stride=1), # padding = kernel_size // 2
)
def __init__(self, N, M, **kwargs):
super().__init__(entropy_bottleneck_channels=N, **kwargs)
self.g_a_conv1 = conv(3, N)
self.g_a_gdn1 = GDN(N)
self.g_a_conv2 = conv(N, N)
self.g_a_gdn2 = GDN(N)
self.g_a_conv3 = conv(N, N)
self.g_a_gdn3 = GDN(N)
self.g_a_conv4 = conv(N, M)
self.g_s_conv1 = deconv(M, N)
self.g_s_gdn1 = GDN(N, inverse=True)
self.g_s_conv2 = deconv(N, N)
self.g_s_gdn2 = GDN(N, inverse=True)
self.g_s_conv3 = deconv(N, N)
self.g_s_gdn3 = GDN(N, inverse=True)
self.g_s_conv4 = deconv(N, 3)
self.h_a = nn.Sequential(
conv(M, N, stride=1, kernel_size=3),
nn.ReLU(inplace=True),
conv(N, N),
nn.ReLU(inplace=True),
conv(N, N),
)
self.h_s = nn.Sequential(
deconv(N, N),
nn.ReLU(inplace=True),
deconv(N, N),
nn.ReLU(inplace=True),
conv(N, M, stride=1, kernel_size=3),
nn.ReLU(inplace=True),
)
self.gaussian_conditional = GaussianConditional(None)
self.N = int(N)
self.M = int(M)
def __init__(self,
N=128,
M=192,
F=21,
C=32,
K=5,
**kwargs): #'cuda:0' or 'cpu'
super().__init__(entropy_bottleneck_channels=N, **kwargs)
# super(DSIC, self).__init__()
# self.entropy_bottleneck1 = CompressionModel(entropy_bottleneck_channels=N)
# self.entropy_bottleneck2 = CompressionModel(entropy_bottleneck_channels=N)
self.gaussian1 = GaussianMixtureConditional(K=K)
self.gaussian2 = GaussianMixtureConditional(K=K)
self.N = int(N)
self.M = int(M)
self.F = F
self.C = C
self.K = K
#定义组件
self.encoder1 = Encoder1(N, M)
# self.encoder2 = Encoder2(N,M)
self.decoder1 = Decoder1(N, M)
# self.decoder2 = Decoder2(N,M)
# pic2 需要的组件
self.pic2_g_a_conv1 = conv(3, N)
self.pic2_g_a_gdn1 = GDN(N)
self.pic2_g_a_conv2 = conv(2 * N, N)
self.pic2_g_a_gdn2 = GDN(N)
self.pic2_g_a_conv3 = conv(2 * N, N)
self.pic2_g_a_gdn3 = GDN(N)
self.pic2_g_a_conv4 = conv(2 * N, M)
#
self.pic2_g_s_conv1 = deconv(M, N)
self.pic2_g_s_gdn1 = GDN(N, inverse=True)
self.pic2_g_s_conv2 = deconv(2 * N, N)
self.pic2_g_s_gdn2 = GDN(N, inverse=True)
self.pic2_g_s_conv3 = deconv(2 * N, N)
self.pic2_g_s_gdn3 = GDN(N, inverse=True)
self.pic2_g_s_conv4 = deconv(2 * N, 3)
#end of pic2
#######
self._global_context = global_context(M, F, C)
#scale_factor 超分辨几倍 (from H,W/16)
self._cost_volume1 = cost_volume(N, 8, F, C) #最外层
self._cost_volume2 = cost_volume(N, 4, F, C)
self._cost_volume3 = cost_volume(N, 2, F, C) #最里层
self._cost_volume4 = cost_volume(N, 2, F, C) # 最里层
self._cost_volume5 = cost_volume(N, 4, F, C)
self._cost_volume6 = cost_volume(N, 8, F, C) # 最外层
self._warp1 = dense_warp()
self._warp2 = dense_warp()
self._warp3 = dense_warp()
self._warp4 = dense_warp()
self._warp5 = dense_warp()
self._warp6 = dense_warp()
#hyper
self._h_a1 = encode_hyper(N=N, M=M)
self._h_a2 = encode_hyper(N=N, M=M)
self._h_s1 = gmm_hyper_y1(N=N, M=M, K=K)
self._h_s2 = gmm_hyper_y2(N=N, M=M, K=K)
def __init__(self,N=128,M=192,K=5,**kwargs): #'cuda:0' or 'cpu'
super().__init__(entropy_bottleneck_channels=N, **kwargs)
# super(DSIC, self).__init__()
# self.entropy_bottleneck1 = CompressionModel(entropy_bottleneck_channels=N)
# self.entropy_bottleneck2 = CompressionModel(entropy_bottleneck_channels=N)
self.gaussian1 = GaussianMixtureConditional(K = K)
self.gaussian2 = GaussianMixtureConditional(K = K)
self.N = int(N)
self.M = int(M)
self.K = int(K)
#定义组件
self.encoder1 = Encoder1(N,M)
self.encoder2 = Encoder2(N,M)
self.decoder1 = Decoder1(N,M)
self.decoder2 = Decoder2(N,M)
# pic2 需要的组件
# #hyper
# self._h_a1 = encode_hyper(N=N,M=M)
# self._h_a2 = encode_hyper(N=N,M=M)
# self._h_s1 = gmm_hyper_y1(N=N,M=M,K=K)
# self._h_s2 = gmm_hyper_y2(N=N,M=M,K=K)
######################################################################
self.h_a1 = nn.Sequential(
conv(M, N, stride=1, kernel_size=3),
nn.LeakyReLU(inplace=True),
conv(N, N, stride=2, kernel_size=5),
nn.LeakyReLU(inplace=True),
conv(N, N, stride=2, kernel_size=5),
)
self.h_s1 = nn.Sequential(
deconv(N, M, stride=2, kernel_size=5),
nn.LeakyReLU(inplace=True),
deconv(M, M * 3 // 2, stride=2, kernel_size=5),
nn.LeakyReLU(inplace=True),
conv(M * 3 // 2, M * 2, stride=1, kernel_size=3),
)
self.entropy_parameters1 = nn.Sequential(
nn.Conv2d(M * 12 // 3, M * 10 // 3, 1),
nn.LeakyReLU(inplace=True),
nn.Conv2d(M * 10 // 3, M * 8 // 3, 1),
nn.LeakyReLU(inplace=True),
nn.Conv2d(M * 8 // 3, M * 6 // 3, 1),
)
self.context_prediction1 = MaskedConv2d(M,
2 * M,
kernel_size=5,
padding=2,
stride=1)
self.gaussian_conditional1 = GaussianConditional(None)
self.h_a2 = nn.Sequential(
conv(M, N, stride=1, kernel_size=3),
nn.LeakyReLU(inplace=True),
conv(N, N, stride=2, kernel_size=5),
nn.LeakyReLU(inplace=True),
conv(N, N, stride=2, kernel_size=5),
)
self.h_s2 = nn.Sequential(
deconv(N, M, stride=2, kernel_size=5),
nn.LeakyReLU(inplace=True),
deconv(M, M * 3 // 2, stride=2, kernel_size=5),
nn.LeakyReLU(inplace=True),
conv(M * 3 // 2, M * 2, stride=1, kernel_size=3),
)
self.entropy_parameters2 = nn.Sequential(
nn.Conv2d(M * 18 // 3, M * 10 // 3, 1), # (M * 12 // 3, M * 10 // 3, 1),
nn.LeakyReLU(inplace=True),
nn.Conv2d(M * 10 // 3, M * 8 // 3, 1),
nn.LeakyReLU(inplace=True),
nn.Conv2d(M * 8 // 3, M * 6 // 3, 1),
)
self.context_prediction2 = MaskedConv2d(M,
2 * M,
kernel_size=5,
padding=2,
stride=1)
self.gaussian_conditional2 = GaussianConditional(None)
