track and test model

.gitattributes

@@ -1,27 +1,5 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zstandard filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+DTM_exp_train10%_model_a/d-best.pth filter=lfs diff=lfs merge=lfs -text
+DTM_exp_train10%_model_a/g-best.pth filter=lfs diff=lfs merge=lfs -text
+DTM_exp_train10%_model_c/d-best.pth filter=lfs diff=lfs merge=lfs -text
+DTM_exp_train10%_model_c/g-best.pth filter=lfs diff=lfs merge=lfs -text
+DTM_exp_train10%_model_b/g-best.pth filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,3 +1,5 @@
 venv
 *.pyc
 __pycache__
+sr.png
+test.png

DTM_exp_train10%_model_a/d-best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ecac9857f6f35d88afec8cea345c81ab1c7a8761ddda8225de820c3d3b48f4c
+size 27401785

DTM_exp_train10%_model_a/g-best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c0dab4fd03de2f0189e87db212a8fce5f1c7771aea0a23b804a146d9c00097df
+size 61648584

DTM_exp_train10%_model_b/g-best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:939adfbd0619db38fc5720ad1c2cb558a605e8e8aefac4703cfa98bbb651ec93
+size 61648714

DTM_exp_train10%_model_c/d-best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:65f96e37fdd9b1bc8c40b7e257a92c9b706f585b791586a50e2ec2463991e059
+size 27401786

DTM_exp_train10%_model_c/g-best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f9652227743f248f3e8aa514c236b8738e11e062c2f270fceec385daedad7c3
+size 49787998

models/modelNetA.py

@@ -0,0 +1,381 @@
+# Copyright 2021 Dakewe Biotech Corporation. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# ==============================================================================
+# File description: Realize the model definition function.
+# ==============================================================================
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+from torch import Tensor
+
+__all__ = [
+    "ResidualDenseBlock", "ResidualResidualDenseBlock",
+    "Discriminator", "Generator",
+    "DownSamplingNetwork"
+]
+
+
+class ResidualDenseBlock(nn.Module):
+    """Achieves densely connected convolutional layers.
+    `Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993v5.pdf>` paper.
+
+    Args:
+            channels (int): The number of channels in the input image.
+            growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels + growths * 0, growths, (3, 3), (1, 1), (1, 1))
+        self.conv2 = nn.Conv2d(channels + growths * 1, growths, (3, 3), (1, 1), (1, 1))
+        self.conv3 = nn.Conv2d(channels + growths * 2, growths, (3, 3), (1, 1), (1, 1))
+        self.conv4 = nn.Conv2d(channels + growths * 3, growths, (3, 3), (1, 1), (1, 1))
+        self.conv5 = nn.Conv2d(channels + growths * 4, channels, (3, 3), (1, 1), (1, 1))
+
+        self.leaky_relu = nn.LeakyReLU(0.2, True)
+        self.identity = nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out1 = self.leaky_relu(self.conv1(x))
+        out2 = self.leaky_relu(self.conv2(torch.cat([x, out1], 1)))
+        out3 = self.leaky_relu(self.conv3(torch.cat([x, out1, out2], 1)))
+        out4 = self.leaky_relu(self.conv4(torch.cat([x, out1, out2, out3], 1)))
+        out5 = self.identity(self.conv5(torch.cat([x, out1, out2, out3, out4], 1)))
+        out = out5 * 0.2 + identity
+
+        return out
+
+
+
+class ResidualDenseBlock(nn.Module):
+    """Achieves densely connected convolutional layers.
+    `Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993v5.pdf>` paper.
+
+    Args:
+            channels (int): The number of channels in the input image.
+            growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels + growths * 0, growths, (3, 3), (1, 1), (1, 1))
+        self.conv2 = nn.Conv2d(channels + growths * 1, growths, (3, 3), (1, 1), (1, 1))
+        self.conv3 = nn.Conv2d(channels + growths * 2, growths, (3, 3), (1, 1), (1, 1))
+        self.conv4 = nn.Conv2d(channels + growths * 3, growths, (3, 3), (1, 1), (1, 1))
+        self.conv5 = nn.Conv2d(channels + growths * 4, channels, (3, 3), (1, 1), (1, 1))
+
+        self.leaky_relu = nn.LeakyReLU(0.2, True)
+        self.identity = nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out1 = self.leaky_relu(self.conv1(x))
+        out2 = self.leaky_relu(self.conv2(torch.cat([x, out1], 1)))
+        out3 = self.leaky_relu(self.conv3(torch.cat([x, out1, out2], 1)))
+        out4 = self.leaky_relu(self.conv4(torch.cat([x, out1, out2, out3], 1)))
+        out5 = self.identity(self.conv5(torch.cat([x, out1, out2, out3, out4], 1)))
+        out = out5 * 0.2 + identity
+
+        return out
+
+
+
+class MiniResidualDenseBlock(nn.Module):
+    """Achieves densely connected convolutional layers.
+    `Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993v5.pdf>` paper.
+
+    Args:
+            channels (int): The number of channels in the input image.
+            growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(MiniResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels + growths * 0, growths, (3, 3), (1, 1), (1, 1))
+        self.conv2 = nn.Conv2d(channels + growths * 1, growths, (3, 3), (1, 1), (1, 1))
+        self.conv3 = nn.Conv2d(channels + growths * 2, growths, (3, 3), (1, 1), (1, 1))
+        self.conv4 = nn.Conv2d(channels + growths * 3, growths, (3, 3), (1, 1), (1, 1))
+        self.conv5 = nn.Conv2d(channels + growths * 4, channels, (3, 3), (1, 1), (1, 1))
+
+        self.leaky_relu = nn.LeakyReLU(0.2, True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out1 = self.leaky_relu(self.conv1(x))
+        out2 = self.leaky_relu(self.conv2(torch.cat([x, out1], 1)))
+        out3 = self.leaky_relu(self.conv3(torch.cat([x, out1, out2], 1)))
+        out4 = self.leaky_relu(self.conv4(torch.cat([x, out1, out2, out3], 1)))
+        out5 = self.leaky_relu(self.conv5(torch.cat([x, out1, out2, out3, out4], 1)))
+        out = out5 * 0.2 + identity
+
+        return out
+
+
+
+class ResidualResidualDenseBlock(nn.Module):
+    """Multi-layer residual dense convolution block.
+
+    Args:
+        channels (int): The number of channels in the input image.
+        growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(ResidualResidualDenseBlock, self).__init__()
+        self.rdb1 = ResidualDenseBlock(channels, growths)
+        self.rdb2 = ResidualDenseBlock(channels, growths)
+        self.rdb3 = ResidualDenseBlock(channels, growths)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+
+        out = self.rdb1(x)
+        out = self.rdb2(out)
+        out = self.rdb3(out)
+        out = out * 0.2 + identity
+
+        return out
+
+
+class MiniResidualResidualDenseBlock(nn.Module):
+    """Multi-layer residual dense convolution block.
+
+    Args:
+        channels (int): The number of channels in the input image.
+        growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(MiniResidualResidualDenseBlock, self).__init__()
+        self.M_rdb1 = MiniResidualDenseBlock(channels, growths)
+        self.M_rdb2 = MiniResidualDenseBlock(channels, growths)
+        self.M_rdb3 = MiniResidualDenseBlock(channels, growths)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+        out = self.M_rdb1(x)
+        out = self.M_rdb2(out)
+        out = self.M_rdb3(out)
+        out = out * 0.2 + identity
+        return out
+
+
+
+class Discriminator(nn.Module):
+    def __init__(self) -> None:
+        super(Discriminator, self).__init__()
+        self.features = nn.Sequential(
+            # input size. (3) x 512 x 512
+            nn.Conv2d(2, 32, (3, 3), (1, 1), (1, 1), bias=True),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(32, 64, (4, 4), (2, 2), (1, 1), bias=False),
+            nn.BatchNorm2d(64),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1), bias=False),
+            nn.BatchNorm2d(64),
+            nn.LeakyReLU(0.2, True),
+            # state size. (128) x 256 x 256
+            nn.Conv2d(64, 128, (4, 4), (2, 2), (1, 1), bias=False),
+            nn.BatchNorm2d(128),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 128, (3, 3), (1, 1), (1, 1), bias=False),
+            nn.BatchNorm2d(128),
+            nn.LeakyReLU(0.2, True),
+            # state size. (256) x 64 x 64
+            nn.Conv2d(128, 256, (4, 4), (2, 2), (1, 1), bias=False),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(256, 256, (3, 3), (1, 1), (1, 1), bias=False),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(256, 256, (4, 4), (2, 2), (1, 1), bias=False),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(256, 256, (3, 3), (1, 1), (1, 1), bias=False),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+            # state size. (512) x 16 x 16
+            nn.Conv2d(256, 256, (4, 4), (2, 2), (1, 1), bias=False),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+
+            nn.Conv2d(256, 256, (4, 4), (2, 2), (1, 1), bias=False),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+            # state size. (512) x 8 x 8
+        )
+
+        self.classifier = nn.Sequential(
+            nn.Linear(256 * 8 * 8, 100),
+            nn.LeakyReLU(0.2, True),
+            nn.Linear(100, 1),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        out = self.features(x)
+        out = torch.flatten(out, 1)
+        out = self.classifier(out)
+        return out
+
+class Generator(nn.Module):
+    def __init__(self) -> None:
+        super(Generator, self).__init__()
+        #RLNet
+        self.RLNetconv_block1 = nn.Conv2d(1, 64, (3, 3), (1, 1), (1, 1))
+        RLNettrunk = []
+        for _ in range(4):
+            RLNettrunk += [ResidualResidualDenseBlock(64, 32)]
+        self.RLNettrunk = nn.Sequential(*RLNettrunk)
+        self.RLNetconv_block2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
+        self.RLNetconv_block3 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True)
+        )
+        self.RLNetconv_block4 = nn.Sequential(
+            nn.Conv2d(64, 1, (3, 3), (1, 1), (1, 1)),
+            nn.Tanh()
+        )
+
+        #############################################################################
+        #Generator
+        self.conv_block1 = nn.Conv2d(1, 64, (3, 3), (1, 1), (1, 1))
+
+        trunk = []
+        for _ in range(16):
+            trunk += [ResidualResidualDenseBlock(64, 32)]
+        self.trunk = nn.Sequential(*trunk)
+
+        # After the feature extraction network, reconnect a layer of convolutional blocks.
+        self.conv_block2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
+
+
+        # Upsampling convolutional layer.
+        self.upsampling = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True)
+        )
+
+        # Reconnect a layer of convolution block after upsampling.
+        self.conv_block3 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True)
+        )
+
+        self.conv_block4 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            #nn.Sigmoid()
+        )
+
+        self.conv_block0_branch0 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 128, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 128, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 64, (3, 3), (1, 1), (1, 1)),
+            nn.Tanh()
+        )
+
+        self.conv_block0_branch1 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 128, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 128, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 64, (3, 3), (1, 1), (1, 1)),
+            nn.Tanh()
+        )
+
+        self.conv_block1_branch0 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 1, (3, 3), (1, 1), (1, 1)),
+            #nn.LeakyReLU(0.2, True),
+            #nn.Conv2d(32, 1, (3, 3), (1, 1), (1, 1)),
+            nn.Sigmoid()
+        )
+
+
+
+        self.conv_block1_branch1 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 1, (3, 3), (1, 1), (1, 1)),
+            nn.Sigmoid())
+
+
+
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        #RLNet
+        out1 = self.RLNetconv_block1(x)
+        out = self.RLNettrunk(out1)
+        out2 = self.RLNetconv_block2(out)
+        out = out1 + out2
+        out = self.RLNetconv_block3(out)
+        out = self.RLNetconv_block4(out)
+        rlNet_out = out + x
+
+        #Generator
+        out1 = self.conv_block1(rlNet_out)
+        out = self.trunk(out1)
+        out2 = self.conv_block2(out)
+        out = out1 + out2
+        out = self.upsampling(F.interpolate(out, scale_factor=2, mode="bicubic"))
+        out = self.upsampling(F.interpolate(out, scale_factor=2, mode="bicubic"))
+        out = self.conv_block3(out)
+        #
+        out = self.conv_block4(out)
+
+        #demResidual = out[:, 1:2, :, :]
+        #grayResidual = out[:, 0:1, :, :]
+
+        # out = self.trunkRGB(out_4)
+        #
+        # out_dem = out[:, 3:4, :, :] * 0.2 + demResidual # DEM images extracted
+        # out_rgb = out[:, 0:3, :, :] * 0.2 + rgbResidual # RGB images extracted
+
+        #ra0
+        #out_rgb=  rgbResidual + self.conv_block0_branch0(rgbResidual)
+
+        out_dem = out + self.conv_block0_branch1(out) #out+ tanh()
+        out_gray = out + self.conv_block0_branch0(out) #out+ tanh()
+
+        out_gray = self.conv_block1_branch0(out_gray) #sigmoid()
+        out_dem = self.conv_block1_branch1(out_dem) #sigmoid()
+
+        return out_gray, out_dem, rlNet_out
+
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+    def _initialize_weights(self) -> None:
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+                m.weight.data *= 0.1
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                m.weight.data *= 0.1

models/modelNetB.py

@@ -0,0 +1,307 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+__all__ = [
+    "ResidualDenseBlock", "ResidualResidualDenseBlock", "Generator",
+    "DownSamplingNetwork"
+]
+
+
+class ResidualDenseBlock(nn.Module):
+    """Achieves densely connected convolutional layers.
+    `Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993v5.pdf>` paper.
+
+    Args:
+            channels (int): The number of channels in the input image.
+            growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels + growths * 0, growths, (3, 3), (1, 1), (1, 1))
+        self.conv2 = nn.Conv2d(channels + growths * 1, growths, (3, 3), (1, 1), (1, 1))
+        self.conv3 = nn.Conv2d(channels + growths * 2, growths, (3, 3), (1, 1), (1, 1))
+        self.conv4 = nn.Conv2d(channels + growths * 3, growths, (3, 3), (1, 1), (1, 1))
+        self.conv5 = nn.Conv2d(channels + growths * 4, channels, (3, 3), (1, 1), (1, 1))
+
+        self.leaky_relu = nn.LeakyReLU(0.2, True)
+        self.identity = nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out1 = self.leaky_relu(self.conv1(x))
+        out2 = self.leaky_relu(self.conv2(torch.cat([x, out1], 1)))
+        out3 = self.leaky_relu(self.conv3(torch.cat([x, out1, out2], 1)))
+        out4 = self.leaky_relu(self.conv4(torch.cat([x, out1, out2, out3], 1)))
+        out5 = self.identity(self.conv5(torch.cat([x, out1, out2, out3, out4], 1)))
+        out = out5 * 0.2 + identity
+
+        return out
+
+
+
+class ResidualDenseBlock(nn.Module):
+    """Achieves densely connected convolutional layers.
+    `Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993v5.pdf>` paper.
+
+    Args:
+            channels (int): The number of channels in the input image.
+            growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels + growths * 0, growths, (3, 3), (1, 1), (1, 1))
+        self.conv2 = nn.Conv2d(channels + growths * 1, growths, (3, 3), (1, 1), (1, 1))
+        self.conv3 = nn.Conv2d(channels + growths * 2, growths, (3, 3), (1, 1), (1, 1))
+        self.conv4 = nn.Conv2d(channels + growths * 3, growths, (3, 3), (1, 1), (1, 1))
+        self.conv5 = nn.Conv2d(channels + growths * 4, channels, (3, 3), (1, 1), (1, 1))
+
+        self.leaky_relu = nn.LeakyReLU(0.2, True)
+        self.identity = nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out1 = self.leaky_relu(self.conv1(x))
+        out2 = self.leaky_relu(self.conv2(torch.cat([x, out1], 1)))
+        out3 = self.leaky_relu(self.conv3(torch.cat([x, out1, out2], 1)))
+        out4 = self.leaky_relu(self.conv4(torch.cat([x, out1, out2, out3], 1)))
+        out5 = self.identity(self.conv5(torch.cat([x, out1, out2, out3, out4], 1)))
+        out = out5 * 0.2 + identity
+
+        return out
+
+
+
+class MiniResidualDenseBlock(nn.Module):
+    """Achieves densely connected convolutional layers.
+    `Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993v5.pdf>` paper.
+
+    Args:
+            channels (int): The number of channels in the input image.
+            growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(MiniResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels + growths * 0, growths, (3, 3), (1, 1), (1, 1))
+        self.conv2 = nn.Conv2d(channels + growths * 1, growths, (3, 3), (1, 1), (1, 1))
+        self.conv3 = nn.Conv2d(channels + growths * 2, growths, (3, 3), (1, 1), (1, 1))
+        self.conv4 = nn.Conv2d(channels + growths * 3, growths, (3, 3), (1, 1), (1, 1))
+        self.conv5 = nn.Conv2d(channels + growths * 4, channels, (3, 3), (1, 1), (1, 1))
+
+        self.leaky_relu = nn.LeakyReLU(0.2, True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out1 = self.leaky_relu(self.conv1(x))
+        out2 = self.leaky_relu(self.conv2(torch.cat([x, out1], 1)))
+        out3 = self.leaky_relu(self.conv3(torch.cat([x, out1, out2], 1)))
+        out4 = self.leaky_relu(self.conv4(torch.cat([x, out1, out2, out3], 1)))
+        out5 = self.leaky_relu(self.conv5(torch.cat([x, out1, out2, out3, out4], 1)))
+        out = out5 * 0.2 + identity
+
+        return out
+
+
+
+class ResidualResidualDenseBlock(nn.Module):
+    """Multi-layer residual dense convolution block.
+
+    Args:
+        channels (int): The number of channels in the input image.
+        growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(ResidualResidualDenseBlock, self).__init__()
+        self.rdb1 = ResidualDenseBlock(channels, growths)
+        self.rdb2 = ResidualDenseBlock(channels, growths)
+        self.rdb3 = ResidualDenseBlock(channels, growths)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+
+        out = self.rdb1(x)
+        out = self.rdb2(out)
+        out = self.rdb3(out)
+        out = out * 0.2 + identity
+
+        return out
+
+
+class MiniResidualResidualDenseBlock(nn.Module):
+    """Multi-layer residual dense convolution block.
+
+    Args:
+        channels (int): The number of channels in the input image.
+        growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(MiniResidualResidualDenseBlock, self).__init__()
+        self.M_rdb1 = MiniResidualDenseBlock(channels, growths)
+        self.M_rdb2 = MiniResidualDenseBlock(channels, growths)
+        self.M_rdb3 = MiniResidualDenseBlock(channels, growths)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+        out = self.M_rdb1(x)
+        out = self.M_rdb2(out)
+        out = self.M_rdb3(out)
+        out = out * 0.2 + identity
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(self) -> None:
+        super(Generator, self).__init__()
+
+        #RLNet
+        self.RLNetconv_block1 = nn.Conv2d(1, 64, (3, 3), (1, 1), (1, 1))
+        RLNettrunk = []
+        for _ in range(4):
+            RLNettrunk += [ResidualResidualDenseBlock(64, 32)]
+        self.RLNettrunk = nn.Sequential(*RLNettrunk)
+        self.RLNetconv_block2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
+        self.RLNetconv_block3 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True)
+        )
+        self.RLNetconv_block4 = nn.Sequential(
+            nn.Conv2d(64, 1, (3, 3), (1, 1), (1, 1)),
+            nn.Tanh()
+        )
+
+        #############################################################################
+        # Generator
+        self.conv_block1 = nn.Conv2d(1, 64, (3, 3), (1, 1), (1, 1))
+        trunk = []
+        for _ in range(16):
+            trunk += [ResidualResidualDenseBlock(64, 32)]
+        self.trunk = nn.Sequential(*trunk)
+
+        # After the feature extraction network, reconnect a layer of convolutional blocks.
+        self.conv_block2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
+
+
+        # Upsampling convolutional layer.
+        self.upsampling = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True)
+        )
+
+        # Reconnect a layer of convolution block after upsampling.
+        self.conv_block3 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True)
+        )
+
+        self.conv_block4 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            #nn.Sigmoid()
+        )
+
+        self.conv_block0_branch0 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 128, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 128, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 64, (3, 3), (1, 1), (1, 1)),
+            nn.Tanh()
+        )
+
+        self.conv_block0_branch1 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 128, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 128, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 64, (3, 3), (1, 1), (1, 1)),
+            nn.Tanh()
+        )
+
+        self.conv_block1_branch0 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 1, (3, 3), (1, 1), (1, 1)),
+            #nn.LeakyReLU(0.2, True),
+            #nn.Conv2d(32, 1, (3, 3), (1, 1), (1, 1)),
+            nn.Sigmoid()
+        )
+
+
+
+        self.conv_block1_branch1 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 1, (3, 3), (1, 1), (1, 1)),
+            nn.Sigmoid())
+
+
+
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        #RLNet
+        out1 = self.RLNetconv_block1(x)
+        out = self.RLNettrunk(out1)
+        out2 = self.RLNetconv_block2(out)
+        out = out1 + out2
+        out = self.RLNetconv_block3(out)
+        out = self.RLNetconv_block4(out)
+        rlNet_out = out + x
+
+        #Generator
+        out1 = self.conv_block1(rlNet_out)
+        out = self.trunk(out1)
+        out2 = self.conv_block2(out)
+        out = out1 + out2
+        out = self.upsampling(F.interpolate(out, scale_factor=2, mode="bicubic"))
+        out = self.upsampling(F.interpolate(out, scale_factor=2, mode="bicubic"))
+        out = self.conv_block3(out)
+        #
+        out = self.conv_block4(out)
+
+        #demResidual = out[:, 1:2, :, :]
+        #grayResidual = out[:, 0:1, :, :]
+
+        # out = self.trunkRGB(out_4)
+        #
+        # out_dem = out[:, 3:4, :, :] * 0.2 + demResidual # DEM images extracted
+        # out_rgb = out[:, 0:3, :, :] * 0.2 + rgbResidual # RGB images extracted
+
+        #ra0
+        #out_rgb=  rgbResidual + self.conv_block0_branch0(rgbResidual)
+
+        out_dem = out + self.conv_block0_branch1(out) #out+ tanh()
+        out_gray = out + self.conv_block0_branch0(out) #out+ tanh()
+
+        out_gray = self.conv_block1_branch0(out_gray) #sigmoid()
+        out_dem = self.conv_block1_branch1(out_dem) #sigmoid()
+
+        return out_gray, out_dem, rlNet_out
+
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+    def _initialize_weights(self) -> None:
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+                m.weight.data *= 0.1
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                m.weight.data *= 0.1
+
+

models/modelNetC.py

@@ -0,0 +1,335 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+__all__ = [
+    "ResidualDenseBlock", "ResidualResidualDenseBlock", "Generator",
+    "DownSamplingNetwork"
+]
+
+
+class ResidualDenseBlock(nn.Module):
+    """Achieves densely connected convolutional layers.
+    `Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993v5.pdf>` paper.
+
+    Args:
+            channels (int): The number of channels in the input image.
+            growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels + growths * 0, growths, (3, 3), (1, 1), (1, 1))
+        self.conv2 = nn.Conv2d(channels + growths * 1, growths, (3, 3), (1, 1), (1, 1))
+        self.conv3 = nn.Conv2d(channels + growths * 2, growths, (3, 3), (1, 1), (1, 1))
+        self.conv4 = nn.Conv2d(channels + growths * 3, growths, (3, 3), (1, 1), (1, 1))
+        self.conv5 = nn.Conv2d(channels + growths * 4, channels, (3, 3), (1, 1), (1, 1))
+
+        self.leaky_relu = nn.LeakyReLU(0.2, True)
+        self.identity = nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out1 = self.leaky_relu(self.conv1(x))
+        out2 = self.leaky_relu(self.conv2(torch.cat([x, out1], 1)))
+        out3 = self.leaky_relu(self.conv3(torch.cat([x, out1, out2], 1)))
+        out4 = self.leaky_relu(self.conv4(torch.cat([x, out1, out2, out3], 1)))
+        out5 = self.identity(self.conv5(torch.cat([x, out1, out2, out3, out4], 1)))
+        out = out5 * 0.2 + identity
+
+        return out
+
+
+
+class ResidualDenseBlock(nn.Module):
+    """Achieves densely connected convolutional layers.
+    `Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993v5.pdf>` paper.
+
+    Args:
+            channels (int): The number of channels in the input image.
+            growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels + growths * 0, growths, (3, 3), (1, 1), (1, 1))
+        self.conv2 = nn.Conv2d(channels + growths * 1, growths, (3, 3), (1, 1), (1, 1))
+        self.conv3 = nn.Conv2d(channels + growths * 2, growths, (3, 3), (1, 1), (1, 1))
+        self.conv4 = nn.Conv2d(channels + growths * 3, growths, (3, 3), (1, 1), (1, 1))
+        self.conv5 = nn.Conv2d(channels + growths * 4, channels, (3, 3), (1, 1), (1, 1))
+
+        self.leaky_relu = nn.LeakyReLU(0.2, True)
+        self.identity = nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out1 = self.leaky_relu(self.conv1(x))
+        out2 = self.leaky_relu(self.conv2(torch.cat([x, out1], 1)))
+        out3 = self.leaky_relu(self.conv3(torch.cat([x, out1, out2], 1)))
+        out4 = self.leaky_relu(self.conv4(torch.cat([x, out1, out2, out3], 1)))
+        out5 = self.identity(self.conv5(torch.cat([x, out1, out2, out3, out4], 1)))
+        out = out5 * 0.2 + identity
+
+        return out
+
+
+
+class MiniResidualDenseBlock(nn.Module):
+    """Achieves densely connected convolutional layers.
+    `Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993v5.pdf>` paper.
+
+    Args:
+            channels (int): The number of channels in the input image.
+            growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(MiniResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels + growths * 0, growths, (3, 3), (1, 1), (1, 1))
+        self.conv2 = nn.Conv2d(channels + growths * 1, growths, (3, 3), (1, 1), (1, 1))
+        self.conv3 = nn.Conv2d(channels + growths * 2, growths, (3, 3), (1, 1), (1, 1))
+        self.conv4 = nn.Conv2d(channels + growths * 3, growths, (3, 3), (1, 1), (1, 1))
+        self.conv5 = nn.Conv2d(channels + growths * 4, channels, (3, 3), (1, 1), (1, 1))
+
+        self.leaky_relu = nn.LeakyReLU(0.2, True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out1 = self.leaky_relu(self.conv1(x))
+        out2 = self.leaky_relu(self.conv2(torch.cat([x, out1], 1)))
+        out3 = self.leaky_relu(self.conv3(torch.cat([x, out1, out2], 1)))
+        out4 = self.leaky_relu(self.conv4(torch.cat([x, out1, out2, out3], 1)))
+        out5 = self.leaky_relu(self.conv5(torch.cat([x, out1, out2, out3, out4], 1)))
+        out = out5 * 0.2 + identity
+
+        return out
+
+
+
+class ResidualResidualDenseBlock(nn.Module):
+    """Multi-layer residual dense convolution block.
+
+    Args:
+        channels (int): The number of channels in the input image.
+        growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(ResidualResidualDenseBlock, self).__init__()
+        self.rdb1 = ResidualDenseBlock(channels, growths)
+        self.rdb2 = ResidualDenseBlock(channels, growths)
+        self.rdb3 = ResidualDenseBlock(channels, growths)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+
+        out = self.rdb1(x)
+        out = self.rdb2(out)
+        out = self.rdb3(out)
+        out = out * 0.2 + identity
+
+        return out
+
+
+class MiniResidualResidualDenseBlock(nn.Module):
+    """Multi-layer residual dense convolution block.
+
+    Args:
+        channels (int): The number of channels in the input image.
+        growths (int): The number of channels that increase in each layer of convolution.
+    """
+
+    def __init__(self, channels: int, growths: int) -> None:
+        super(MiniResidualResidualDenseBlock, self).__init__()
+        self.M_rdb1 = MiniResidualDenseBlock(channels, growths)
+        self.M_rdb2 = MiniResidualDenseBlock(channels, growths)
+        self.M_rdb3 = MiniResidualDenseBlock(channels, growths)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+        out = self.M_rdb1(x)
+        out = self.M_rdb2(out)
+        out = self.M_rdb3(out)
+        out = out * 0.2 + identity
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(self) -> None:
+        super(Generator, self).__init__()
+        # Generator
+        self.conv_block1 = nn.Conv2d(1, 64, (3, 3), (1, 1), (1, 1))
+        trunk = []
+        for _ in range(16):
+            trunk += [ResidualResidualDenseBlock(64, 32)]
+        self.trunk = nn.Sequential(*trunk)
+
+        # After the feature extraction network, reconnect a layer of convolutional blocks.
+        self.conv_block2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
+
+
+        # Upsampling convolutional layer.
+        self.upsampling = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True)
+        )
+
+        # Reconnect a layer of convolution block after upsampling.
+        self.conv_block3 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True)
+        )
+
+        self.conv_block4 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            #nn.Sigmoid()
+        )
+
+        self.conv_block0_branch0 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 128, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 128, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 64, (3, 3), (1, 1), (1, 1)),
+            nn.Tanh()
+        )
+
+        self.conv_block0_branch1 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 128, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 128, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 64, (3, 3), (1, 1), (1, 1)),
+            nn.Tanh()
+        )
+
+        self.conv_block1_branch0 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 1, (3, 3), (1, 1), (1, 1)),
+            #nn.LeakyReLU(0.2, True),
+            #nn.Conv2d(32, 1, (3, 3), (1, 1), (1, 1)),
+            nn.Sigmoid()
+        )
+
+
+
+        self.conv_block1_branch1 = nn.Sequential(
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 1, (3, 3), (1, 1), (1, 1)),
+            nn.Sigmoid())
+
+
+
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        #Generator
+        out1 = self.conv_block1(x)
+        out = self.trunk(out1)
+        out2 = self.conv_block2(out)
+        out = out1 + out2
+        out = self.upsampling(F.interpolate(out, scale_factor=2, mode="bicubic"))
+        out = self.upsampling(F.interpolate(out, scale_factor=2, mode="bicubic"))
+        out = self.conv_block3(out)
+        #
+        out = self.conv_block4(out)
+
+        #demResidual = out[:, 1:2, :, :]
+        #grayResidual = out[:, 0:1, :, :]
+
+        # out = self.trunkRGB(out_4)
+        #
+        # out_dem = out[:, 3:4, :, :] * 0.2 + demResidual # DEM images extracted
+        # out_rgb = out[:, 0:3, :, :] * 0.2 + rgbResidual # RGB images extracted
+
+        #ra0
+        #out_rgb=  rgbResidual + self.conv_block0_branch0(rgbResidual)
+
+        out_dem = out + self.conv_block0_branch1(out) #out+ tanh()
+        out_gray = out + self.conv_block0_branch0(out) #out+ tanh()
+
+        out_gray = self.conv_block1_branch0(out_gray) #sigmoid()
+        out_dem = self.conv_block1_branch1(out_dem) #sigmoid()
+
+        return out_gray, out_dem
+
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+    def _initialize_weights(self) -> None:
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+                m.weight.data *= 0.1
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                m.weight.data *= 0.1
+
+class Discriminator(nn.Module):
+    def __init__(self) -> None:
+        super(Discriminator, self).__init__()
+        self.features = nn.Sequential(
+            # input size. (3) x 512 x 512
+            nn.Conv2d(2, 32, (3, 3), (1, 1), (1, 1), bias=True),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(32, 64, (4, 4), (2, 2), (1, 1), bias=False),
+            nn.BatchNorm2d(64),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1), bias=False),
+            nn.BatchNorm2d(64),
+            nn.LeakyReLU(0.2, True),
+            # state size. (128) x 256 x 256
+            nn.Conv2d(64, 128, (4, 4), (2, 2), (1, 1), bias=False),
+            nn.BatchNorm2d(128),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128, 128, (3, 3), (1, 1), (1, 1), bias=False),
+            nn.BatchNorm2d(128),
+            nn.LeakyReLU(0.2, True),
+            # state size. (256) x 64 x 64
+            nn.Conv2d(128, 256, (4, 4), (2, 2), (1, 1), bias=False),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(256, 256, (3, 3), (1, 1), (1, 1), bias=False),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(256, 256, (4, 4), (2, 2), (1, 1), bias=False),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(256, 256, (3, 3), (1, 1), (1, 1), bias=False),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+            # state size. (512) x 16 x 16
+            nn.Conv2d(256, 256, (4, 4), (2, 2), (1, 1), bias=False),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+
+            nn.Conv2d(256, 256, (4, 4), (2, 2), (1, 1), bias=False),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+            # state size. (512) x 8 x 8
+        )
+
+        self.classifier = nn.Sequential(
+            nn.Linear(256 * 8 * 8, 100),
+            nn.LeakyReLU(0.2, True),
+            nn.Linear(100, 1),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        out = self.features(x)
+        out = torch.flatten(out, 1)
+        out = self.classifier(out)
+        return out
+

requirements.txt

@@ -1,3 +1,4 @@
+matplotlib
 gradio
 torch
-torchvision
+torchvision

test.py

@@ -0,0 +1,55 @@
+import torch
+import torchvision
+from torchvision import transforms
+from PIL import Image
+import matplotlib.pyplot as plt
+import numpy as np
+from models.modelNetA import Generator as GA
+from models.modelNetB import Generator as GB
+from models.modelNetC import Generator as GC
+
+
+
+DEVICE='cpu'
+model_type = 'model_b'
+
+modeltype2path = {
+    'model_a': 'DTM_exp_train10%_model_a/g-best.pth',
+    'model_b': 'DTM_exp_train10%_model_b/g-best.pth',
+    'model_c': 'DTM_exp_train10%_model_c/g-best.pth',
+}
+
+if model_type == 'model_a':
+    generator = GA()
+if model_type == 'model_b':
+    generator = GB()
+if model_type == 'model_c':
+    generator = GC()
+
+generator = torch.nn.DataParallel(generator)
+state_dict_Gen = torch.load(modeltype2path[model_type], map_location=torch.device('cpu'))
+generator.load_state_dict(state_dict_Gen)
+generator = generator.module.to(DEVICE)
+# generator.to(DEVICE)
+generator.eval()
+
+preprocess = transforms.Compose([
+    transforms.Grayscale(),
+    transforms.Resize((512, 512)),
+    transforms.ToTensor()
+])
+input_img = Image.open('demo_imgs/fake.jpg')
+torch_img = preprocess(input_img).to(DEVICE).unsqueeze(0).to(DEVICE)
+with torch.no_grad():
+    output = generator(torch_img)
+sr, sr_dem_selected = output[0], output[1]
+sr = sr.squeeze(0).cpu()
+
+print(sr.shape)
+torchvision.utils.save_image(sr, 'sr.png')
+
+sr_dem_selected = sr_dem_selected.squeeze().cpu().detach().numpy()
+print(sr_dem_selected.shape)
+plt.imshow(sr_dem_selected, cmap='jet', vmin=0, vmax=np.max(sr_dem_selected))
+plt.colorbar()
+plt.savefig('test.png')