第一次提交

app.py

@@ -0,0 +1,141 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/8/4 下午2:38
+# @Author  : xiaoshun
+# @Email   : 3038523973@qq.com
+# @File    : app.py
+# @Software: PyCharm
+
+from glob import glob
+import gradio as gr
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+import albumentations as albu
+from albumentations.pytorch.transforms import ToTensorV2
+from src.data.components.hrcwhu import HRCWHU
+from src.data.hrcwhu_datamodule import HRCWHUDataModule
+from src.models.components.cdnetv1 import CDnetV1
+from src.models.components.cdnetv2 import CDnetV2
+from src.models.components.dual_branch import Dual_Branch
+from src.models.components.hrcloud import HRcloudNet
+from src.models.components.mcdnet import MCDNet
+from src.models.components.scnn import SCNNNet
+
+
+class Application:
+    def __init__(self):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.models = {
+            "cdnetv1": CDnetV1(num_classes=2).to(self.device),
+            "cdnetv2": CDnetV2(num_classes=2).to(self.device),
+            "hrcloud": HRcloudNet(num_classes=2).to(self.device),
+            "mcdnet": MCDNet(in_channels=3, num_classes=2).to(self.device),
+            "scnn": SCNNNet(num_classes=2).to(self.device),
+            "dual_branch": Dual_Branch(img_size=256, in_channels=3, num_classes=2).to(
+                self.device
+            ),
+        }
+        self.__load_weight()
+        self.transform = albu.Compose(
+            [
+                albu.Resize(256, 256, always_apply=True),
+                ToTensorV2(),
+            ]
+        )
+
+    def __load_weight(self):
+        """
+        将模型权重加载进来
+        """
+        for model_name, model in self.models.items():
+            weight_path = glob(
+                f"logs/train/runs/*{model_name}*/*/checkpoints/*epoch*.ckpt"
+            )[0]
+            weight = torch.load(weight_path, map_location=self.device)
+            state_dict = {}
+            for key, value in weight["state_dict"].items():
+                new_key = key[4:]
+                state_dict[new_key] = value
+            model.load_state_dict(state_dict)
+            model.eval()
+            print(f"{model_name} weight loaded!")
+
+    @torch.no_grad
+    def inference(self, image: torch.Tensor, model_name: str):
+        x = image.float()
+        x = x.unsqueeze(0)
+        x = x.to(self.device)
+        logits = self.models[model_name](x)
+        if isinstance(logits, tuple):
+            logits = logits[0]
+        fake_mask = torch.argmax(logits, 1).detach().cpu().squeeze(0).numpy()
+        return fake_mask
+
+    def give_colors_to_mask(self, mask: np.ndarray):
+        """
+        赋予mask颜色
+        """
+        assert len(mask.shape) == 2, "Value Error,mask的形状为(height,width)"
+        colors_mask = np.zeros((mask.shape[0], mask.shape[1], 3)).astype(np.float32)
+        colors = ((255, 255, 255), (128, 192, 128))
+        for color in range(2):
+            segc = mask == color
+            colors_mask[:, :, 0] += segc * (colors[color][0])
+            colors_mask[:, :, 1] += segc * (colors[color][1])
+            colors_mask[:, :, 2] += segc * (colors[color][2])
+        return colors_mask
+
+    def to_pil(self, image: np.ndarray, width=None, height=None):
+        colors_np = self.give_colors_to_mask(image)
+        pil_np = Image.fromarray(np.uint8(colors_np))
+        if width and height:
+            pil_np = pil_np.resize((width, height))
+        return pil_np
+
+    def flip(self, image_pil: Image.Image, model_name: str):
+        if image_pil is None:
+            return Image.fromarray(np.uint8(np.random.random((32,32,3)) * 255)), "请上传一张图片"
+        if model_name is None:
+            return Image.fromarray(np.uint8(np.random.random((32,32,3)) * 255)), "请选择模型名称"
+        image = np.array(image_pil)
+        raw_height, raw_width = image.shape[0], image.shape[1]
+        transform = self.transform(image=image)
+        image = transform["image"]
+        image = image / 255.0
+        fake_image = self.inference(image, model_name)
+        fake_image = self.to_pil(fake_image, raw_width, raw_height)
+        return fake_image,"success"
+
+    def tiff_to_png(image: Image.Image):
+        if image.format == "TIFF":
+            image = image.convert("RGB")
+        return np.array(image)
+
+    def run(self):
+        app = gr.Interface(
+            self.flip,
+            [
+                gr.Image(sources=["clipboard", "upload"], type="pil"),
+                gr.Radio(
+                    ["cdnetv1", "cdnetv2", "hrcloud", "mcdnet", "scnn", "dual_branch"],
+                    label="model_name",
+                    info="选择使用的模型",
+                ),
+            ],
+            [gr.Image(), gr.Textbox(label="提示信息")],
+            examples=[
+                ["examples_png/barren_11.png", "dual_branch"],
+                ["examples_png/snow_10.png", "scnn"],
+                ["examples_png/vegetation_21.png", "cdnetv2"],
+                ["examples_png/water_22.png", "hrcloud"],
+            ],
+            title="云检测模型在线演示",
+            submit_btn=gr.Button("Submit", variant="primary")
+        )
+        app.launch(share=True)
+
+
+if __name__ == "__main__":
+    app = Application()
+    app.run()

logs/train/runs/hrcwhu_cdnetv1/2024-08-03_18-18-09/checkpoints/epoch_063.ckpt

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

logs/train/runs/hrcwhu_cdnetv1/2024-08-03_18-18-09/checkpoints/last.ckpt

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

logs/train/runs/hrcwhu_cdnetv2/2024-08-03_18-26-19/checkpoints/epoch_150.ckpt

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

logs/train/runs/hrcwhu_cdnetv2/2024-08-03_18-26-19/checkpoints/last.ckpt

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

logs/train/runs/hrcwhu_dual_branch/2024-08-03_18-40-06/checkpoints/epoch_009.ckpt

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

logs/train/runs/hrcwhu_dual_branch/2024-08-03_18-40-06/checkpoints/last.ckpt

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

logs/train/runs/hrcwhu_hrcloud/2024-08-03_18-55-52/checkpoints/epoch_024.ckpt

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

logs/train/runs/hrcwhu_hrcloud/2024-08-03_18-55-52/checkpoints/last.ckpt

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

logs/train/runs/hrcwhu_mcdnet/2024-08-04_14-30-56/checkpoints/epoch_032.ckpt

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

logs/train/runs/hrcwhu_mcdnet/2024-08-04_14-30-56/checkpoints/last.ckpt

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

logs/train/runs/hrcwhu_scnn/2024-08-03_19-29-26/checkpoints/epoch_196.ckpt

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

logs/train/runs/hrcwhu_scnn/2024-08-03_19-29-26/checkpoints/last.ckpt

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

requirements.txt

@@ -0,0 +1,289 @@
+absl-py==2.1.0
+addict==2.4.0
+aiofiles==23.2.1
+aiohttp==3.9.5
+aiosignal==1.3.1
+albucore==0.0.12
+albumentations==1.4.12
+alembic==1.13.2
+aliyun-python-sdk-core==2.15.1
+aliyun-python-sdk-kms==2.16.3
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.4.0
+argon2-cffi==23.1.0
+argon2-cffi-bindings==21.2.0
+arrow==1.3.0
+asttokens==2.4.1
+astunparse==1.6.3
+async-lru==2.0.4
+async-timeout==4.0.3
+attrs==23.2.0
+autopage==0.5.2
+Babel==2.15.0
+beautifulsoup4==4.12.3
+bleach==6.1.0
+blessed==1.20.0
+Brotli @ file:///croot/brotli-split_1714483155106/work
+certifi @ file:///croot/certifi_1717618050233/work/certifi
+cffi==1.16.0
+cfgv==3.4.0
+chardet==5.2.0
+charset-normalizer @ file:///tmp/build/80754af9/charset-normalizer_1630003229654/work
+click==8.1.7
+cliff==4.7.0
+cmaes==0.10.0
+cmd2==2.4.3
+colorama==0.4.6
+colorlog==6.8.2
+comm==0.2.2
+contourpy==1.2.1
+crcmod==1.7
+cryptography==42.0.8
+cycler==0.12.1
+debugpy==1.8.1
+decorator==5.1.1
+defusedxml==0.7.1
+dill==0.3.8
+distlib==0.3.8
+docker-pycreds==0.4.0
+efficientnet_pytorch==0.7.1
+einops @ file:///home/conda/feedstock_root/build_artifacts/einops_1714285159399/work
+eval_type_backport==0.2.0
+exceptiongroup==1.2.1
+executing==2.0.1
+fastapi==0.112.0
+fastjsonschema==2.20.0
+ffmpy==0.4.0
+filelock==3.14.0
+flatbuffers==24.3.25
+fonttools==4.53.0
+fqdn==1.5.1
+frozenlist==1.4.1
+fsspec==2024.5.0
+ftfy==6.2.0
+future==1.0.0
+gast==0.6.0
+gitdb==4.0.11
+GitPython==3.1.43
+gmpy2 @ file:///tmp/build/80754af9/gmpy2_1645455533097/work
+google-pasta==0.2.0
+gpustat==1.1.1
+gradio==4.40.0
+gradio_client==1.2.0
+greenlet==3.0.3
+grpcio==1.64.1
+h11==0.14.0
+h5py==3.11.0
+hf_transfer==0.1.8
+httpcore==1.0.5
+httpx==0.27.0
+huggingface-hub==0.23.4
+hydra-colorlog==1.2.0
+hydra-core==1.3.2
+hydra-optuna-sweeper==1.2.0
+identify==2.6.0
+idna @ file:///croot/idna_1714398848350/work
+image-dehazer==0.0.9
+imageio==2.34.2
+imgaug==0.4.0
+importlib_metadata==8.0.0
+importlib_resources==6.4.0
+iniconfig==2.0.0
+ipykernel==6.29.4
+ipython==8.25.0
+isoduration==20.11.0
+jedi==0.19.1
+Jinja2 @ file:///croot/jinja2_1716993405101/work
+jmespath==0.10.0
+joblib==1.4.2
+json5==0.9.25
+jsonpointer==3.0.0
+jsonschema==4.22.0
+jsonschema-specifications==2023.12.1
+jupyter-events==0.10.0
+jupyter-lsp==2.2.5
+jupyter_client==8.6.2
+jupyter_core==5.7.2
+jupyter_server==2.14.1
+jupyter_server_terminals==0.5.3
+jupyterlab==4.2.2
+jupyterlab_pygments==0.3.0
+jupyterlab_server==2.27.2
+keras==3.4.1
+kiwisolver==1.4.5
+kornia==0.7.3
+kornia_rs==0.1.5
+lazy_loader==0.4
+libclang==18.1.1
+lightning==2.3.3
+lightning-utilities==0.11.3.post0
+Mako==1.3.5
+mamba-ssm==2.2.2
+Markdown==3.6
+markdown-it-py==3.0.0
+MarkupSafe @ file:///croot/markupsafe_1704205993651/work
+matplotlib==3.9.0
+matplotlib-inline==0.1.7
+mdurl==0.1.2
+mistune==3.0.2
+mkl-fft @ file:///croot/mkl_fft_1695058164594/work
+mkl-random @ file:///croot/mkl_random_1695059800811/work
+mkl-service==2.4.0
+ml-dtypes==0.4.0
+mmcv==2.1.0
+mmengine==0.10.4
+mmsegmentation==1.2.2
+model-index==0.1.11
+mpmath @ file:///croot/mpmath_1690848262763/work
+multidict==6.0.5
+multiprocess==0.70.16
+munch==4.0.0
+namex==0.0.8
+natsort==8.4.0
+nbclient==0.10.0
+nbconvert==7.16.4
+nbformat==5.10.4
+nest-asyncio==1.6.0
+networkx @ file:///croot/networkx_1717597493534/work
+ninja==1.11.1.1
+nodeenv==1.9.1
+notebook_shim==0.2.4
+numpy @ file:///croot/numpy_and_numpy_base_1708638617955/work/dist/numpy-1.26.4-cp310-cp310-linux_x86_64.whl#sha256=d8cd837ed43e87f77e6efaa08e8de927ca030a1c9c5d04624432d6fb9a74a5ee
+nvidia-ml-py==12.555.43
+omegaconf==2.3.0
+opencv-python==4.10.0.84
+opencv-python-headless==4.10.0.84
+opendatalab==0.0.10
+openmim==0.3.9
+openxlab==0.1.1
+opt-einsum==3.3.0
+optree==0.12.1
+optuna==2.10.1
+ordered-set==4.1.0
+orjson==3.10.6
+oss2==2.17.0
+overrides==7.7.0
+packaging==24.1
+pandas==2.2.2
+pandocfilters==1.5.1
+parso==0.8.4
+patsy==0.5.6
+pbr==6.0.0
+pexpect==4.9.0
+pillow @ file:///croot/pillow_1714398848491/work
+platformdirs==4.2.2
+plotly==5.22.0
+pluggy==1.5.0
+pre-commit==3.7.1
+pretrainedmodels==0.7.4
+prettytable==3.10.0
+prometheus_client==0.20.0
+prompt_toolkit==3.0.47
+protobuf==4.25.3
+psutil==6.0.0
+ptyprocess==0.7.0
+pure-eval==0.2.2
+pyarrow==16.1.0
+pyarrow-hotfix==0.6
+pycparser==2.22
+pycryptodome==3.20.0
+pydantic==2.8.0
+pydantic_core==2.20.0
+pydub==0.25.1
+Pygments==2.18.0
+pyparsing==3.1.2
+pyperclip==1.9.0
+pyproject==1.3.1
+PySocks @ file:///home/builder/ci_310/pysocks_1640793678128/work
+pytest==8.2.2
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-json-logger==2.0.7
+python-multipart==0.0.9
+pytorch-lightning==2.3.3
+pytz==2023.4
+pyvips==2.2.3
+PyYAML @ file:///croot/pyyaml_1698096049011/work
+pyzmq==26.0.3
+referencing==0.35.1
+regex==2024.5.15
+requests==2.32.3
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+rich==13.4.2
+rootutils==1.0.7
+rpds-py==0.18.1
+ruff==0.5.6
+safetensors==0.4.3
+scikit-image==0.24.0
+scikit-learn==1.5.0
+scipy==1.13.1
+seaborn==0.13.2
+segmentation-models-pytorch==0.3.3
+semantic-version==2.10.0
+Send2Trash==1.8.3
+sentry-sdk==2.9.0
+setproctitle==1.3.3
+shapely==2.0.4
+shellingham==1.5.4
+six==1.16.0
+smmap==5.0.1
+sniffio==1.3.1
+soupsieve==2.5
+SQLAlchemy==2.0.31
+stack-data==0.6.3
+starlette==0.37.2
+statsmodels==0.14.2
+stevedore==5.2.0
+sympy @ file:///croot/sympy_1701397643339/work
+tabulate==0.9.0
+tenacity==8.4.1
+tensorboard==2.17.0
+tensorboard-data-server==0.7.2
+tensorboardX==2.6.2.2
+tensorflow==2.17.0
+tensorflow-io-gcs-filesystem==0.37.1
+termcolor==2.4.0
+terminado==0.18.1
+thop==0.1.1.post2209072238
+threadpoolctl==3.5.0
+tifffile==2024.6.18
+timm==0.9.2
+tinycss2==1.3.0
+tokenizers==0.19.1
+tomli==2.0.1
+tomlkit==0.12.0
+torch==2.3.1
+torchaudio==2.3.1
+torchinfo==1.8.0
+torchmetrics==1.4.0.post0
+torchsummary==1.5.1
+torchvision==0.18.1
+tornado==6.4.1
+tqdm==4.65.2
+traitlets==5.14.3
+transformers==4.42.4
+triton==2.3.1
+typer==0.12.3
+types-python-dateutil==2.9.0.20240316
+typing_extensions @ file:///croot/typing_extensions_1715268824938/work
+tzdata==2024.1
+uri-template==1.3.0
+urllib3==2.2.2
+uvicorn==0.30.5
+virtualenv==20.26.3
+wandb==0.17.4
+wcwidth==0.2.13
+webcolors==24.6.0
+webencodings==0.5.1
+websocket-client==1.8.0
+websockets==12.0
+Werkzeug==3.0.3
+wrapt==1.16.0
+xlrd==2.0.1
+xxhash==3.4.1
+yacs==0.1.8
+yapf==0.40.2
+yarl==1.9.4
+zipp==3.19.2

src/models/components/cdnetv1.py

@@ -0,0 +1,396 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/7/24 上午11:36
+# @Author  : xiaoshun
+# @Email   : 3038523973@qq.com
+# @File    : cdnetv1.py
+# @Software: PyCharm
+
+"""Cloud detection Network"""
+
+"""Cloud detection Network"""
+
+"""
+This is the implementation of CDnetV1 without multi-scale inputs. This implementation uses ResNet by default.
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+from torch.utils import data, model_zoo
+from torch.autograd import Variable
+import math
+import numpy as np
+
+affine_par = True
+from torch.autograd import Function
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes, affine=affine_par)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes, affine=affine_par)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride, bias=False)  # change
+        self.bn1 = nn.BatchNorm2d(planes, affine=affine_par)
+        for i in self.bn1.parameters():
+            i.requires_grad = False
+
+        padding = dilation
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1,  # change
+                               padding=padding, bias=False, dilation=dilation)
+        self.bn2 = nn.BatchNorm2d(planes, affine=affine_par)
+        for i in self.bn2.parameters():
+            i.requires_grad = False
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4, affine=affine_par)
+        for i in self.bn3.parameters():
+            i.requires_grad = False
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Classifier_Module(nn.Module):
+
+    def __init__(self, dilation_series, padding_series, num_classes):
+        super(Classifier_Module, self).__init__()
+        self.conv2d_list = nn.ModuleList()
+        for dilation, padding in zip(dilation_series, padding_series):
+            self.conv2d_list.append(
+                nn.Conv2d(2048, num_classes, kernel_size=3, stride=1, padding=padding, dilation=dilation, bias=True))
+
+        for m in self.conv2d_list:
+            m.weight.data.normal_(0, 0.01)
+
+    def forward(self, x):
+        out = self.conv2d_list[0](x)
+        for i in range(len(self.conv2d_list) - 1):
+            out += self.conv2d_list[i + 1](x)
+            return out
+
+
+class _ConvBNReLU(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0,
+                 dilation=1, groups=1, norm_layer=nn.BatchNorm2d):
+        super(_ConvBNReLU, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=False)
+        self.bn = norm_layer(out_channels)
+        self.relu = nn.ReLU(True)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+class _ASPPConv(nn.Module):
+    def __init__(self, in_channels, out_channels, atrous_rate, norm_layer):
+        super(_ASPPConv, self).__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, padding=atrous_rate, dilation=atrous_rate, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU(True)
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class _AsppPooling(nn.Module):
+    def __init__(self, in_channels, out_channels, norm_layer):
+        super(_AsppPooling, self).__init__()
+        self.gap = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(in_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU(True)
+        )
+
+    def forward(self, x):
+        size = x.size()[2:]
+        pool = self.gap(x)
+        out = F.interpolate(pool, size, mode='bilinear', align_corners=True)
+        return out
+
+
+class _ASPP(nn.Module):
+    def __init__(self, in_channels, atrous_rates, norm_layer):
+        super(_ASPP, self).__init__()
+        out_channels = 512  # changed from 256
+        self.b0 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU(True)
+        )
+
+        rate1, rate2, rate3 = tuple(atrous_rates)
+        self.b1 = _ASPPConv(in_channels, out_channels, rate1, norm_layer)
+        self.b2 = _ASPPConv(in_channels, out_channels, rate2, norm_layer)
+        self.b3 = _ASPPConv(in_channels, out_channels, rate3, norm_layer)
+        self.b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer)
+
+        # self.project = nn.Sequential(
+        # nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
+        # norm_layer(out_channels),
+        # nn.ReLU(True),
+        # nn.Dropout(0.5))
+        self.dropout2d = nn.Dropout2d(0.3)
+
+    def forward(self, x):
+        feat1 = self.dropout2d(self.b0(x))
+        feat2 = self.dropout2d(self.b1(x))
+        feat3 = self.dropout2d(self.b2(x))
+        feat4 = self.dropout2d(self.b3(x))
+        feat5 = self.dropout2d(self.b4(x))
+        x = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
+        # x = self.project(x)
+        return x
+
+
+class _FPM(nn.Module):
+    def __init__(self, in_channels, num_classes, norm_layer=nn.BatchNorm2d):
+        super(_FPM, self).__init__()
+        self.aspp = _ASPP(in_channels, [6, 12, 18], norm_layer=norm_layer)
+        # self.dropout2d = nn.Dropout2d(0.5)
+
+    def forward(self, x):
+        x = torch.cat((x, self.aspp(x)), dim=1)
+        # x = self.dropout2d(x) # added
+        return x
+
+
+class BR(nn.Module):
+    def __init__(self, num_classes, stride=1, downsample=None):
+        super(BR, self).__init__()
+        self.conv1 = conv3x3(num_classes, num_classes * 16, stride)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(num_classes * 16, num_classes)
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out += residual
+
+        return out
+
+
+class CDnetV1(nn.Module):
+    def __init__(self, block=Bottleneck, layers=[3, 4, 6, 3], num_classes=21, aux=True):
+        self.inplanes = 64
+        self.aux = aux
+        super().__init__()
+        # self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        # self.bn1 = nn.BatchNorm2d(64, affine = affine_par)
+
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64, affine=affine_par)
+        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(64, affine=affine_par)
+        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(64, affine=affine_par)
+
+        for i in self.bn1.parameters():
+            i.requires_grad = False
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=True)  # change
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4)
+        # self.layer5 = self._make_pred_layer(Classifier_Module, [6,12,18,24],[6,12,18,24],num_classes)
+
+        self.res5_con1x1 = nn.Sequential(
+            nn.Conv2d(1024 + 2048, 512, kernel_size=1, stride=1, padding=0),
+            nn.BatchNorm2d(512),
+            nn.ReLU(True)
+        )
+
+        self.fpm1 = _FPM(512, num_classes)
+        self.fpm2 = _FPM(512, num_classes)
+        self.fpm3 = _FPM(256, num_classes)
+
+        self.br1 = BR(num_classes)
+        self.br2 = BR(num_classes)
+        self.br3 = BR(num_classes)
+        self.br4 = BR(num_classes)
+        self.br5 = BR(num_classes)
+        self.br6 = BR(num_classes)
+        self.br7 = BR(num_classes)
+
+        self.predict1 = self._predict_layer(512 * 6, num_classes)
+        self.predict2 = self._predict_layer(512 * 6, num_classes)
+        self.predict3 = self._predict_layer(512 * 5 + 256, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, 0.01)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+        #        for i in m.parameters():
+        #            i.requires_grad = False
+
+    def _predict_layer(self, in_channels, num_classes):
+        return nn.Sequential(nn.Conv2d(in_channels, 256, kernel_size=1, stride=1, padding=0),
+                             nn.BatchNorm2d(256),
+                             nn.ReLU(True),
+                             nn.Dropout2d(0.1),
+                             nn.Conv2d(256, num_classes, kernel_size=3, stride=1, padding=1, bias=True))
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion or dilation == 2 or dilation == 4:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion, affine=affine_par))
+        for i in downsample._modules['1'].parameters():
+            i.requires_grad = False
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, dilation=dilation, downsample=downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, dilation=dilation))
+
+        return nn.Sequential(*layers)
+
+    # def _make_pred_layer(self,block, dilation_series, padding_series,num_classes):
+    # return block(dilation_series,padding_series,num_classes)
+
+    def base_forward(self, x):
+        x = self.relu(self.bn1(self.conv1(x)))
+        size_conv1 = x.size()[2:]
+        x = self.relu(self.bn2(self.conv2(x)))
+        x = self.relu(self.bn3(self.conv3(x)))
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        res2 = x
+        x = self.layer2(x)
+        res3 = x
+        x = self.layer3(x)
+        res4 = x
+        x = self.layer4(x)
+        x = self.res5_con1x1(torch.cat([x, res4], dim=1))
+
+        return x, res3, res2, size_conv1
+
+    def forward(self, x):
+        size = x.size()[2:]
+        score1, score2, score3, size_conv1 = self.base_forward(x)
+        # outputs = list()
+        score1 = self.fpm1(score1)
+        score1 = self.predict1(score1)  # 1/8
+        predict1 = score1
+        score1 = self.br1(score1)
+
+        score2 = self.fpm2(score2)
+        score2 = self.predict2(score2)  # 1/8
+        predict2 = score2
+
+        # first fusion
+        score2 = self.br2(score2) + score1
+        score2 = self.br3(score2)
+
+        score3 = self.fpm3(score3)
+        score3 = self.predict3(score3)  # 1/4
+        predict3 = score3
+        score3 = self.br4(score3)
+
+        # second fusion
+        size_score3 = score3.size()[2:]
+        score3 = score3 + F.interpolate(score2, size_score3, mode='bilinear', align_corners=True)
+        score3 = self.br5(score3)
+
+        # upsampling + BR
+        score3 = F.interpolate(score3, size_conv1, mode='bilinear', align_corners=True)
+        score3 = self.br6(score3)
+        score3 = F.interpolate(score3, size, mode='bilinear', align_corners=True)
+        score3 = self.br7(score3)
+
+        # if self.aux:
+        # auxout = self.dsn(mid)
+        # auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True)
+        # #outputs.append(auxout)
+        return score3
+        # return score3, predict1, predict2, predict3
+
+
+if __name__ == '__main__':
+    model = CDnetV1(num_classes=21)
+    fake_image = torch.randn(2, 3, 224, 224)
+    outputs = model(fake_image)
+    for out in outputs:
+        print(out.shape)
+    # torch.Size([2, 21, 224, 224])
+    # torch.Size([2, 21, 29, 29])
+    # torch.Size([2, 21, 29, 29])
+    # torch.Size([2, 21, 57, 57])

src/models/components/cdnetv2.py

@@ -0,0 +1,699 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/7/24 下午3:41
+# @Author  : xiaoshun
+# @Email   : 3038523973@qq.com
+# @File    : cdnetv2.py
+# @Software: PyCharm
+
+"""Cloud detection Network"""
+
+"""
+This is the implementation of CDnetV2 without multi-scale inputs. This implementation uses ResNet by default.
+"""
+# nn.GroupNorm
+
+import torch
+from torch import nn
+# import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+from torch.utils import data, model_zoo
+from torch.autograd import Variable
+import math
+import numpy as np
+
+affine_par = True
+from torch.autograd import Function
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes, affine=affine_par)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes, affine=affine_par)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride, bias=False)  # change
+        self.bn1 = nn.BatchNorm2d(planes, affine=affine_par)
+        for i in self.bn1.parameters():
+            i.requires_grad = False
+
+        padding = dilation
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1,  # change
+                               padding=padding, bias=False, dilation=dilation)
+        self.bn2 = nn.BatchNorm2d(planes, affine=affine_par)
+        for i in self.bn2.parameters():
+            i.requires_grad = False
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4, affine=affine_par)
+        for i in self.bn3.parameters():
+            i.requires_grad = False
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+        # self.layerx_1 = Bottleneck_nosample(64, 64, stride=1, dilation=1)
+        # self.layerx_2 = Bottleneck(256, 64, stride=1, dilation=1, downsample=None)
+        # self.layerx_3 = Bottleneck_downsample(256, 64, stride=2, dilation=1)
+
+
+class Res_block_1(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes=64, planes=64, stride=1, dilation=1):
+        super(Res_block_1, self).__init__()
+
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(inplanes, planes, kernel_size=1, stride=1, bias=False),
+            nn.GroupNorm(8, planes),
+            nn.ReLU(inplace=True))
+
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(planes, planes, kernel_size=3, stride=1,
+                      padding=1, bias=False, dilation=1),
+            nn.GroupNorm(8, planes),
+            nn.ReLU(inplace=True))
+
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False),
+            nn.GroupNorm(8, planes * 4))
+
+        self.relu = nn.ReLU(inplace=True)
+
+        self.down_sample = nn.Sequential(
+            nn.Conv2d(inplanes, planes * 4,
+                      kernel_size=1, stride=1, bias=False),
+            nn.GroupNorm(8, planes * 4))
+
+    def forward(self, x):
+        # residual = x
+
+        out = self.conv1(x)
+        out = self.conv2(out)
+        out = self.conv3(out)
+        residual = self.down_sample(x)
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Res_block_2(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes=256, planes=64, stride=1, dilation=1):
+        super(Res_block_2, self).__init__()
+
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(inplanes, planes, kernel_size=1, stride=1, bias=False),
+            nn.GroupNorm(8, planes),
+            nn.ReLU(inplace=True))
+
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(planes, planes, kernel_size=3, stride=1,
+                      padding=1, bias=False, dilation=1),
+            nn.GroupNorm(8, planes),
+            nn.ReLU(inplace=True))
+
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False),
+            nn.GroupNorm(8, planes * 4))
+
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.conv2(out)
+        out = self.conv3(out)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Res_block_3(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes=256, planes=64, stride=1, dilation=1):
+        super(Res_block_3, self).__init__()
+
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride, bias=False),
+            nn.GroupNorm(8, planes),
+            nn.ReLU(inplace=True))
+
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(planes, planes, kernel_size=3, stride=1,
+                      padding=1, bias=False, dilation=1),
+            nn.GroupNorm(8, planes),
+            nn.ReLU(inplace=True))
+
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False),
+            nn.GroupNorm(8, planes * 4))
+
+        self.relu = nn.ReLU(inplace=True)
+
+        self.downsample = nn.Sequential(
+            nn.Conv2d(inplanes, planes * 4,
+                      kernel_size=1, stride=stride, bias=False),
+            nn.GroupNorm(8, planes * 4))
+
+    def forward(self, x):
+        # residual = x
+
+        out = self.conv1(x)
+        out = self.conv2(out)
+        out = self.conv3(out)
+        # residual = self.downsample(x)
+        out += self.downsample(x)
+        out = self.relu(out)
+
+        return out
+
+
+class Classifier_Module(nn.Module):
+
+    def __init__(self, dilation_series, padding_series, num_classes):
+        super(Classifier_Module, self).__init__()
+        self.conv2d_list = nn.ModuleList()
+        for dilation, padding in zip(dilation_series, padding_series):
+            self.conv2d_list.append(
+                nn.Conv2d(2048, num_classes, kernel_size=3, stride=1, padding=padding, dilation=dilation, bias=True))
+
+        for m in self.conv2d_list:
+            m.weight.data.normal_(0, 0.01)
+
+    def forward(self, x):
+        out = self.conv2d_list[0](x)
+        for i in range(len(self.conv2d_list) - 1):
+            out += self.conv2d_list[i + 1](x)
+            return out
+
+
+class _ConvBNReLU(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0,
+                 dilation=1, groups=1, relu6=False, norm_layer=nn.BatchNorm2d):
+        super(_ConvBNReLU, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=False)
+        self.bn = norm_layer(out_channels)
+        self.relu = nn.ReLU6(True) if relu6 else nn.ReLU(True)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+class _ASPPConv(nn.Module):
+    def __init__(self, in_channels, out_channels, atrous_rate, norm_layer):
+        super(_ASPPConv, self).__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, padding=atrous_rate, dilation=atrous_rate, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU(True)
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class _AsppPooling(nn.Module):
+    def __init__(self, in_channels, out_channels, norm_layer):
+        super(_AsppPooling, self).__init__()
+        self.gap = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(in_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU(True)
+        )
+
+    def forward(self, x):
+        size = x.size()[2:]
+        pool = self.gap(x)
+        out = F.interpolate(pool, size, mode='bilinear', align_corners=True)
+        return out
+
+
+class _ASPP(nn.Module):
+    def __init__(self, in_channels, atrous_rates, norm_layer):
+        super(_ASPP, self).__init__()
+        out_channels = 256
+        self.b0 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU(True)
+        )
+
+        rate1, rate2, rate3 = tuple(atrous_rates)
+        self.b1 = _ASPPConv(in_channels, out_channels, rate1, norm_layer)
+        self.b2 = _ASPPConv(in_channels, out_channels, rate2, norm_layer)
+        self.b3 = _ASPPConv(in_channels, out_channels, rate3, norm_layer)
+        self.b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer)
+
+        self.project = nn.Sequential(
+            nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU(True),
+            nn.Dropout(0.5)
+        )
+
+    def forward(self, x):
+        feat1 = self.b0(x)
+        feat2 = self.b1(x)
+        feat3 = self.b2(x)
+        feat4 = self.b3(x)
+        feat5 = self.b4(x)
+        x = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
+        x = self.project(x)
+        return x
+
+
+class _DeepLabHead(nn.Module):
+    def __init__(self, num_classes, c1_channels=256, norm_layer=nn.BatchNorm2d):
+        super(_DeepLabHead, self).__init__()
+        self.aspp = _ASPP(2048, [12, 24, 36], norm_layer=norm_layer)
+        self.c1_block = _ConvBNReLU(c1_channels, 48, 3, padding=1, norm_layer=norm_layer)
+        self.block = nn.Sequential(
+            _ConvBNReLU(304, 256, 3, padding=1, norm_layer=norm_layer),
+            nn.Dropout(0.5),
+            _ConvBNReLU(256, 256, 3, padding=1, norm_layer=norm_layer),
+            nn.Dropout(0.1),
+            nn.Conv2d(256, num_classes, 1))
+
+    def forward(self, x, c1):
+        size = c1.size()[2:]
+        c1 = self.c1_block(c1)
+        x = self.aspp(x)
+        x = F.interpolate(x, size, mode='bilinear', align_corners=True)
+        return self.block(torch.cat([x, c1], dim=1))
+
+
+class _CARM(nn.Module):
+    def __init__(self, in_planes, ratio=8):
+        super(_CARM, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+
+        self.fc1_1 = nn.Linear(in_planes, in_planes // ratio)
+        self.fc1_2 = nn.Linear(in_planes // ratio, in_planes)
+
+        self.fc2_1 = nn.Linear(in_planes, in_planes // ratio)
+        self.fc2_2 = nn.Linear(in_planes // ratio, in_planes)
+        self.relu = nn.ReLU(True)
+
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = self.avg_pool(x)
+        avg_out = avg_out.view(avg_out.size(0), -1)
+        avg_out = self.fc1_2(self.relu(self.fc1_1(avg_out)))
+
+        max_out = self.max_pool(x)
+        max_out = max_out.view(max_out.size(0), -1)
+        max_out = self.fc2_2(self.relu(self.fc2_1(max_out)))
+
+        max_out_size = max_out.size()[1]
+        avg_out = torch.reshape(avg_out, (-1, max_out_size, 1, 1))
+        max_out = torch.reshape(max_out, (-1, max_out_size, 1, 1))
+
+        out = self.sigmoid(avg_out + max_out)
+
+        x = out * x
+        return x
+
+
+class FSFB_CH(nn.Module):
+    def __init__(self, in_planes, num, ratio=8):
+        super(FSFB_CH, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+
+        self.fc1_1 = nn.Linear(in_planes, in_planes // ratio)
+        self.fc1_2 = nn.Linear(in_planes // ratio, num * in_planes)
+
+        self.fc2_1 = nn.Linear(in_planes, in_planes // ratio)
+        self.fc2_2 = nn.Linear(in_planes // ratio, num * in_planes)
+        self.relu = nn.ReLU(True)
+
+        self.fc3 = nn.Linear(num * in_planes, 2 * num * in_planes)
+        self.fc4 = nn.Linear(2 * num * in_planes, 2 * num * in_planes)
+        self.fc5 = nn.Linear(2 * num * in_planes, num * in_planes)
+
+        self.softmax = nn.Softmax(dim=3)
+
+    def forward(self, x, num):
+        avg_out = self.avg_pool(x)
+        avg_out = avg_out.view(avg_out.size(0), -1)
+        avg_out = self.fc1_2(self.relu(self.fc1_1(avg_out)))
+
+        max_out = self.max_pool(x)
+        max_out = max_out.view(max_out.size(0), -1)
+        max_out = self.fc2_2(self.relu(self.fc2_1(max_out)))
+
+        out = avg_out + max_out
+        out = self.relu(self.fc3(out))
+        out = self.relu(self.fc4(out))
+        out = self.relu(self.fc5(out))  # (N, num*in_planes)
+
+        out_size = out.size()[1]
+        out = torch.reshape(out, (-1, out_size // num, 1, num))  # (N, in_planes, 1, num )
+        out = self.softmax(out)
+
+        channel_scale = torch.chunk(out, num, dim=3)  # (N, in_planes, 1, 1 )
+
+        return channel_scale
+
+
+class FSFB_SP(nn.Module):
+    def __init__(self, num, norm_layer=nn.BatchNorm2d):
+        super(FSFB_SP, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(2, 2 * num, kernel_size=3, padding=1, bias=False),
+            norm_layer(2 * num),
+            nn.ReLU(True),
+            nn.Conv2d(2 * num, 4 * num, kernel_size=3, padding=1, bias=False),
+            norm_layer(4 * num),
+            nn.ReLU(True),
+            nn.Conv2d(4 * num, 4 * num, kernel_size=3, padding=1, bias=False),
+            norm_layer(4 * num),
+            nn.ReLU(True),
+            nn.Conv2d(4 * num, 2 * num, kernel_size=3, padding=1, bias=False),
+            norm_layer(2 * num),
+            nn.ReLU(True),
+            nn.Conv2d(2 * num, num, kernel_size=3, padding=1, bias=False)
+        )
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, x, num):
+        avg_out = torch.mean(x, dim=1, keepdim=True)
+        max_out, _ = torch.max(x, dim=1, keepdim=True)
+        x = torch.cat([avg_out, max_out], dim=1)
+        x = self.conv(x)
+        x = self.softmax(x)
+        spatial_scale = torch.chunk(x, num, dim=1)
+        return spatial_scale
+
+
+##################################################################################################################
+
+
+class _HFFM(nn.Module):
+    def __init__(self, in_channels, atrous_rates, norm_layer=nn.BatchNorm2d):
+        super(_HFFM, self).__init__()
+        out_channels = 256
+        self.b0 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 1, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU(True)
+        )
+
+        rate1, rate2, rate3 = tuple(atrous_rates)
+        self.b1 = _ASPPConv(in_channels, out_channels, rate1, norm_layer)
+        self.b2 = _ASPPConv(in_channels, out_channels, rate2, norm_layer)
+        self.b3 = _ASPPConv(in_channels, out_channels, rate3, norm_layer)
+        self.b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer)
+        self.carm = _CARM(in_channels)
+        self.sa = FSFB_SP(4, norm_layer)
+        self.ca = FSFB_CH(out_channels, 4, 8)
+
+    def forward(self, x, num):
+        x = self.carm(x)
+        # feat1 = self.b0(x)
+        feat1 = self.b1(x)
+        feat2 = self.b2(x)
+        feat3 = self.b3(x)
+        feat4 = self.b4(x)
+        feat = feat1 + feat2 + feat3 + feat4
+        spatial_atten = self.sa(feat, num)
+        channel_atten = self.ca(feat, num)
+
+        feat_ca = channel_atten[0] * feat1 + channel_atten[1] * feat2 + channel_atten[2] * feat3 + channel_atten[
+            3] * feat4
+        feat_sa = spatial_atten[0] * feat1 + spatial_atten[1] * feat2 + spatial_atten[2] * feat3 + spatial_atten[
+            3] * feat4
+        feat_sa = feat_sa + feat_ca
+
+        return feat_sa
+
+
+class _AFFM(nn.Module):
+    def __init__(self, in_channels=256, norm_layer=nn.BatchNorm2d):
+        super(_AFFM, self).__init__()
+
+        self.sa = FSFB_SP(2, norm_layer)
+        self.ca = FSFB_CH(in_channels, 2, 8)
+        self.carm = _CARM(in_channels)
+
+    def forward(self, feat1, feat2, hffm, num):
+        feat = feat1 + feat2
+        spatial_atten = self.sa(feat, num)
+        channel_atten = self.ca(feat, num)
+
+        feat_ca = channel_atten[0] * feat1 + channel_atten[1] * feat2
+        feat_sa = spatial_atten[0] * feat1 + spatial_atten[1] * feat2
+        output = self.carm(feat_sa + feat_ca + hffm)
+        # output = self.carm (feat_sa  + hffm)
+
+        return output, channel_atten, spatial_atten
+
+
+class block_Conv3x3(nn.Module):
+    def __init__(self, in_channels):
+        super(block_Conv3x3, self).__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(in_channels, 256, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.BatchNorm2d(256),
+            nn.ReLU(True)
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class CDnetV2(nn.Module):
+    def __init__(self, block=Bottleneck, layers=[3, 4, 6, 3], num_classes=21, aux=True):
+        self.inplanes = 256  # change
+        self.aux = aux
+        super().__init__()
+        # self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        # self.bn1 = nn.BatchNorm2d(64, affine = affine_par)
+
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64, affine=affine_par)
+
+        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(64, affine=affine_par)
+
+        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(64, affine=affine_par)
+
+        self.relu = nn.ReLU(inplace=True)
+
+        self.dropout = nn.Dropout(0.3)
+        for i in self.bn1.parameters():
+            i.requires_grad = False
+
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=True)  # change
+
+        # self.layer1 = self._make_layer(block, 64, layers[0])
+
+        self.layerx_1 = Res_block_1(64, 64, stride=1, dilation=1)
+        self.layerx_2 = Res_block_2(256, 64, stride=1, dilation=1)
+        self.layerx_3 = Res_block_3(256, 64, stride=2, dilation=1)
+
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4)
+        # self.layer5 = self._make_pred_layer(Classifier_Module, [6,12,18,24],[6,12,18,24],num_classes)
+
+        self.hffm = _HFFM(2048, [6, 12, 18])
+        self.affm_1 = _AFFM()
+        self.affm_2 = _AFFM()
+        self.affm_3 = _AFFM()
+        self.affm_4 = _AFFM()
+        self.carm = _CARM(256)
+
+        self.con_layer1_1 = block_Conv3x3(256)
+        self.con_res2 = block_Conv3x3(256)
+        self.con_res3 = block_Conv3x3(512)
+        self.con_res4 = block_Conv3x3(1024)
+        self.con_res5 = block_Conv3x3(2048)
+
+        self.dsn1 = nn.Sequential(
+            nn.Conv2d(256, num_classes, kernel_size=1, stride=1, padding=0)
+        )
+
+        self.dsn2 = nn.Sequential(
+            nn.Conv2d(256, num_classes, kernel_size=1, stride=1, padding=0)
+        )
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, 0.01)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+        #        for i in m.parameters():
+        #            i.requires_grad = False
+
+        # self.inplanes = 256 # change
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion or dilation == 2 or dilation == 4:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion, affine=affine_par))
+        for i in downsample._modules['1'].parameters():
+            i.requires_grad = False
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, dilation=dilation, downsample=downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, dilation=dilation))
+
+        return nn.Sequential(*layers)
+
+    # def _make_pred_layer(self,block, dilation_series, padding_series,num_classes):
+    # return block(dilation_series,padding_series,num_classes)
+
+    def base_forward(self, x):
+        x = self.relu(self.bn1(self.conv1(x)))  # 1/2
+        x = self.relu(self.bn2(self.conv2(x)))
+        x = self.relu(self.bn3(self.conv3(x)))
+        x = self.maxpool(x)  # 1/4
+
+        # x = self.layer1(x)	# 1/8
+
+        # layer1
+        x = self.layerx_1(x)  # 1/4
+        layer1_0 = x
+
+        x = self.layerx_2(x)  # 1/4
+        layer1_0 = self.con_layer1_1(x + layer1_0)  # 256
+        size_layer1_0 = layer1_0.size()[2:]
+
+        x = self.layerx_3(x)  # 1/8
+        res2 = self.con_res2(x)  # 256
+        size_res2 = res2.size()[2:]
+
+        # layer2-4
+        x = self.layer2(x)  # 1/16
+        res3 = self.con_res3(x)  # 256
+        x = self.layer3(x)  # 1/16
+
+        res4 = self.con_res4(x)  # 256
+        x = self.layer4(x)  # 1/16
+        res5 = self.con_res5(x)  # 256
+
+        # x = self.res5_con1x1(torch.cat([x, res4], dim=1))
+        return layer1_0, res2, res3, res4, res5, x, size_layer1_0, size_res2
+
+        # return res2, res3, res4, res5, x, layer_1024,  size_res2
+
+    def forward(self, x):
+        # size = x.size()[2:]
+        layer1_0, res2, res3, res4, res5, layer4, size_layer1_0, size_res2 = self.base_forward(x)
+
+        hffm = self.hffm(layer4, 4)  # 256 HFFM
+        res5 = res5 + hffm
+        aux_feature = res5  # loss_aux
+        # res5 = self.carm(res5)
+        res5, _, _ = self.affm_1(res4, res5, hffm, 2)  # 1/16
+        # aux_feature = res5
+        res5, _, _ = self.affm_2(res3, res5, hffm, 2)  # 1/16
+
+        res5 = F.interpolate(res5, size_res2, mode='bilinear', align_corners=True)
+        res5, _, _ = self.affm_3(res2, res5, F.interpolate(hffm, size_res2, mode='bilinear', align_corners=True), 2)
+
+        res5 = F.interpolate(res5, size_layer1_0, mode='bilinear', align_corners=True)
+        res5, _, _ = self.affm_4(layer1_0, res5,
+                                 F.interpolate(hffm, size_layer1_0, mode='bilinear', align_corners=True), 2)
+
+        output = self.dsn1(res5)
+
+        if self.aux:
+            auxout = self.dsn2(aux_feature)
+            # auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True)
+            # outputs.append(auxout)
+        size = x.size()[2:]
+        pred, pred_aux = output, auxout
+        pred = F.interpolate(pred, size, mode='bilinear', align_corners=True)
+        pred_aux = F.interpolate(pred_aux, size, mode='bilinear', align_corners=True)
+        return pred, pred_aux
+
+
+if __name__ == '__main__':
+    model = CDnetV2(num_classes=3)
+    fake_image = torch.rand(2, 3, 256, 256)
+    output = model(fake_image)
+    for out in output:
+        print(out.shape)
+        # torch.Size([2, 3, 256, 256])
+        # torch.Size([2, 3, 256, 256])

src/models/components/cnn.py

@@ -0,0 +1,26 @@
+import torch
+from torch import nn
+
+
+class CNN(nn.Module):
+    def __init__(self, dim=32):
+        super(CNN, self).__init__()
+        self.conv1 = nn.Conv2d(1, dim, 5)
+        self.conv2 = nn.Conv2d(dim, dim * 2, 5)
+        self.fc1 = nn.Linear(dim * 2 * 4 * 4, 10)
+    
+    def forward(self, x):
+        x = torch.relu(self.conv1(x))
+        x = torch.max_pool2d(x, 2)
+        x = torch.relu(self.conv2(x))
+        x = torch.max_pool2d(x, 2)
+        x = x.view(-1, x.shape[1] * x.shape[2] * x.shape[3])
+        x = self.fc1(x)
+        return x
+
+
+if __name__ == "__main__":
+    input = torch.randn(2, 1, 28, 28)
+    model = CNN()
+    output = model(input)
+    assert output.shape == (2, 10)

src/models/components/dual_branch.py

@@ -0,0 +1,680 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/7/26 上午11:19
+# @Author  : xiaoshun
+# @Email   : 3038523973@qq.com
+# @File    : dual_branch.py
+# @Software: PyCharm
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# from models.Transformer.ViT import truncated_normal_
+
+# Decoder细化卷积模块
+class SBR(nn.Module):
+    def __init__(self, in_ch):
+        super(SBR, self).__init__()
+        self.conv1x3 = nn.Sequential(
+            nn.Conv2d(in_ch, in_ch, kernel_size=(1, 3), stride=1, padding=(0, 1)),
+            nn.BatchNorm2d(in_ch),
+            nn.ReLU(True)
+        )
+        self.conv3x1 = nn.Sequential(
+            nn.Conv2d(in_ch, in_ch, kernel_size=(3, 1), stride=1, padding=(1, 0)),
+            nn.BatchNorm2d(in_ch),
+            nn.ReLU(True)
+        )
+
+    def forward(self, x):
+        out = self.conv3x1(self.conv1x3(x))  # 先进行1x3的卷积，得到结果并将结果再进行3x1的卷积
+        return out + x
+
+
+# 下采样卷积模块 stage 1,2,3
+class c_stage123(nn.Module):
+    def __init__(self, in_chans, out_chans):
+        super().__init__()
+        self.stage123 = nn.Sequential(
+            nn.Conv2d(in_channels=in_chans, out_channels=out_chans, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(out_chans),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=out_chans, out_channels=out_chans, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(out_chans),
+            nn.ReLU(),
+        )
+        self.conv1x1_123 = nn.Conv2d(in_channels=in_chans, out_channels=out_chans, kernel_size=1)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+    def forward(self, x):
+        stage123 = self.stage123(x)  # 3*3卷积，两倍下采样    3*224*224-->64*112*112
+        max = self.maxpool(x)  # 最大值池化，两倍下采样   3*224*224-->3*112*112
+        max = self.conv1x1_123(max)  # 1*1卷积     3*112*112-->64*112*112
+        stage123 = stage123 + max  # 残差结构，广播机制
+        return stage123
+
+
+# 下采样卷积模块 stage4,5
+class c_stage45(nn.Module):
+    def __init__(self, in_chans, out_chans):
+        super().__init__()
+        self.stage45 = nn.Sequential(
+            nn.Conv2d(in_channels=in_chans, out_channels=out_chans, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(out_chans),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=out_chans, out_channels=out_chans, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(out_chans),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=out_chans, out_channels=out_chans, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(out_chans),
+            nn.ReLU(),
+        )
+        self.conv1x1_45 = nn.Conv2d(in_channels=in_chans, out_channels=out_chans, kernel_size=1)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+    def forward(self, x):
+        stage45 = self.stage45(x)  # 3*3卷积模块 2倍下采样
+        max = self.maxpool(x)  # 最大值池化，两倍下采样
+        max = self.conv1x1_45(max)  # 1*1卷积模块 调整通道数
+        stage45 = stage45 + max  # 残差结构
+        return stage45
+
+
+class Identity(nn.Module):  # 恒等映射
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x
+
+
+# 轻量卷积模块
+class DepthwiseConv2d(nn.Module):  # 用于自注意力机制
+    def __init__(self, in_chans, out_chans, kernel_size=1, stride=1, padding=0, dilation=1):
+        super().__init__()
+        # depthwise conv
+        self.depthwise = nn.Conv2d(
+            in_channels=in_chans,
+            out_channels=in_chans,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,  # 深层卷积的膨胀率
+            groups=in_chans  # 指定分组卷积的组数
+        )
+        # batch norm
+        self.bn = nn.BatchNorm2d(num_features=in_chans)
+
+        # pointwise conv   逐点卷积
+        self.pointwise = nn.Conv2d(
+            in_channels=in_chans,
+            out_channels=out_chans,
+            kernel_size=1
+        )
+
+    def forward(self, x):
+        x = self.depthwise(x)
+        x = self.bn(x)
+        x = self.pointwise(x)
+        return x
+
+
+# residual skip connection 残差跳跃连接
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, input, **kwargs):
+        x = self.fn(input, **kwargs)
+        return (x + input)
+
+
+# layer norm plus 层归一化
+class PreNorm(nn.Module):  # 代表神经网络层
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+
+    def forward(self, input, **kwargs):
+        return self.fn(self.norm(input), **kwargs)
+
+
+# FeedForward层使得representation的表达能力更强
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout=0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(in_features=dim, out_features=hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(in_features=hidden_dim, out_features=dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, input):
+        return self.net(input)
+
+
+class ConvAttnetion(nn.Module):
+    '''
+    using the Depth_Separable_Wise Conv2d to produce the q, k, v instead of using Linear Project in ViT
+    '''
+
+    def __init__(self, dim, img_size, heads=8, dim_head=64, kernel_size=3, q_stride=1, k_stride=1, v_stride=1,
+                 dropout=0., last_stage=False):
+        super().__init__()
+        self.last_stage = last_stage
+        self.img_size = img_size
+        inner_dim = dim_head * heads  # 512
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** (-0.5)
+
+        pad = (kernel_size - q_stride) // 2
+
+        self.to_q = DepthwiseConv2d(in_chans=dim, out_chans=inner_dim, kernel_size=kernel_size, stride=q_stride,
+                                    padding=pad)  # 自注意力机制
+        self.to_k = DepthwiseConv2d(in_chans=dim, out_chans=inner_dim, kernel_size=kernel_size, stride=k_stride,
+                                    padding=pad)
+        self.to_v = DepthwiseConv2d(in_chans=dim, out_chans=inner_dim, kernel_size=kernel_size, stride=v_stride,
+                                    padding=pad)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(
+                in_features=inner_dim,
+                out_features=dim
+            ),
+            nn.Dropout(dropout)
+        ) if project_out else Identity()
+
+    def forward(self, x):
+        b, n, c, h = *x.shape, self.heads  # * 星号的作用大概是去掉 tuple 属性吧
+
+        # print(x.shape)
+        # print('+++++++++++++++++++++++++++++++++')
+
+        # if语句内容没有使用
+        if self.last_stage:
+            cls_token = x[:, 0]
+            # print(cls_token.shape)
+            # print('+++++++++++++++++++++++++++++++++')
+            x = x[:, 1:]  # 去掉每个数组的第一个元素
+
+            cls_token = rearrange(torch.unsqueeze(cls_token, dim=1), 'b n (h d) -> b h n d', h=h)
+
+        # rearrange:用于对张量的维度进行重新变换排序，可用于替换pytorch中的reshape，view，transpose和permute等操作
+        x = rearrange(x, 'b (l w) n -> b n l w', l=self.img_size, w=self.img_size)  # [1, 3136, 64]-->1*64*56*56
+        # batch_size,N(通道数),h,w
+
+        q = self.to_q(x)  # 1*64*56*56-->1*64*56*56
+        # print(q.shape)
+        # print('++++++++++++++')
+        q = rearrange(q, 'b (h d) l w -> b h (l w) d', h=h)  # 1*64*56*56-->1*1*3136*64
+        # print(q.shape)
+        # print('=====================')
+        # batch_size,head,h*w,dim_head
+
+        k = self.to_k(x)  # 操作和q一样
+        k = rearrange(k, 'b (h d) l w -> b h (l w) d', h=h)
+        # batch_size,head,h*w,dim_head
+
+        v = self.to_v(x)  ##操作和q一样
+        # print(v.shape)
+        # print('[[[[[[[[[[[[[[[[[[[[[[[[[[[[')
+        v = rearrange(v, 'b (h d) l w -> b h (l w) d', h=h)
+        # print(v.shape)
+        # print(']]]]]]]]]]]]]]]]]]]]]]]]]]]')
+        # batch_size,head,h*w,dim_head
+
+        if self.last_stage:
+            # print(q.shape)
+            # print('================')
+            q = torch.cat([cls_token, q], dim=2)
+            # print(q.shape)
+            # print('++++++++++++++++++')
+            v = torch.cat([cls_token, v], dim=2)
+            k = torch.cat([cls_token, k], dim=2)
+
+        # calculate attention by matmul + scale
+        # permute:(batch_size,head,dim_head,h*w
+        # print(k.shape)
+        # print('++++++++++++++++++++')
+        k = k.permute(0, 1, 3, 2)  # 1*1*3136*64-->1*1*64*3136
+        # print(k.shape)
+        # print('====================')
+        attention = (q.matmul(k))  # 1*1*3136*3136
+        # print(attention.shape)
+        # print('--------------------')
+        attention = attention * self.scale  # 可以得到一个logit的向量，避免出现梯度下降和梯度爆炸
+        # print(attention.shape)
+        # print('####################')
+        # pass a softmax
+        attention = F.softmax(attention, dim=-1)
+        # print(attention.shape)
+        # print('********************')
+
+        # matmul v
+        # attention.matmul(v):(batch_size,head,h*w,dim_head)
+        # permute:(batch_size,h*w,head,dim_head)
+        out = (attention.matmul(v)).permute(0, 2, 1, 3).reshape(b, n,
+                                                                c)  # 1*3136*64  这些操作的目的是将注意力权重和值向量相乘后得到的结果进行重塑，得到一个形状为 (batch size, 序列长度, 值向量或矩阵的维度) 的张量
+
+        # linear project
+        out = self.to_out(out)
+        return out
+
+
+# Reshape Layers
+class Rearrange(nn.Module):
+    def __init__(self, string, h, w):
+        super().__init__()
+        self.string = string
+        self.h = h
+        self.w = w
+
+    def forward(self, input):
+
+        if self.string == 'b c h w -> b (h w) c':
+            N, C, H, W = input.shape
+            # print(input.shape)
+            x = torch.reshape(input, shape=(N, -1, self.h * self.w)).permute(0, 2, 1)
+            # print(x.shape)
+            # print('+++++++++++++++++++')
+        if self.string == 'b (h w) c -> b c h w':
+            N, _, C = input.shape
+            # print(input.shape)
+            x = torch.reshape(input, shape=(N, self.h, self.w, -1)).permute(0, 3, 1, 2)
+            # print(x.shape)
+            # print('=====================')
+        return x
+
+
+# Transformer layers
+class Transformer(nn.Module):
+    def __init__(self, dim, img_size, depth, heads, dim_head, mlp_dim, dropout=0., last_stage=False):
+        super().__init__()
+        self.layers = nn.ModuleList([  # 管理子模块，参数注册
+            nn.ModuleList([
+                PreNorm(dim=dim, fn=ConvAttnetion(dim, img_size, heads=heads, dim_head=dim_head, dropout=dropout,
+                                                  last_stage=last_stage)),  # 归一化，重参数化
+                PreNorm(dim=dim, fn=FeedForward(dim=dim, hidden_dim=mlp_dim, dropout=dropout))
+            ]) for _ in range(depth)
+        ])
+
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = x + attn(x)
+            x = x + ff(x)
+        return x
+
+
+class Dual_Branch(nn.Module):  # 最主要的大函数
+    def __init__(self, img_size, in_channels, num_classes, dim=64, kernels=[7, 3, 3, 3], strides=[4, 2, 2, 2],
+                 heads=[1, 3, 6, 6],
+                 depth=[1, 2, 10, 10], pool='cls', dropout=0., emb_dropout=0., scale_dim=4, ):
+        super().__init__()
+
+        assert pool in ['cls', 'mean'], f'pool type must be either cls or mean pooling'
+        self.pool = pool
+        self.dim = dim
+
+        # stage1
+        # k:7 s:4    in: 1, 64, 56, 56  out: 1, 3136, 64
+        self.stage1_conv_embed = nn.Sequential(
+            nn.Conv2d(  # 1*3*224*224-->[1, 64, 56, 56]
+                in_channels=in_channels,
+                out_channels=dim,
+                kernel_size=kernels[0],
+                stride=strides[0],
+                padding=2
+            ),
+            Rearrange('b c h w -> b (h w) c', h=img_size // 4, w=img_size // 4),  # [1, 64, 56, 56]-->[1, 3136, 64]
+            nn.LayerNorm(dim)  # 对每个batch归一化
+        )
+
+        self.stage1_transformer = nn.Sequential(
+            Transformer(  #
+                dim=dim,
+                img_size=img_size // 4,
+                depth=depth[0],  # Transformer层中的编码器和解码器层数。
+                heads=heads[0],
+                dim_head=self.dim,  # 它是每个注意力头的维度大小，通常是嵌入维度除以头数。
+                mlp_dim=dim * scale_dim,  # mlp_dim:它是Transformer中前馈神经网络的隐藏层维度大小，通常是嵌入维度乘以一个缩放因子。
+                dropout=dropout,
+                # last_stage=last_stage     #它是一个标志位，用于表示该Transformer层是否是最后一层。
+            ),
+            Rearrange('b (h w) c -> b c h w', h=img_size // 4, w=img_size // 4)
+        )
+
+        # stage2
+        # k:3 s:2  in: 1, 192, 28, 28  out: 1, 784, 192
+        in_channels = dim
+        scale = heads[1] // heads[0]
+        dim = scale * dim
+
+        self.stage2_conv_embed = nn.Sequential(
+            nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=dim,
+                kernel_size=kernels[1],
+                stride=strides[1],
+                padding=1
+            ),
+            Rearrange('b c h w -> b (h w) c', h=img_size // 8, w=img_size // 8),
+            nn.LayerNorm(dim)
+        )
+
+        self.stage2_transformer = nn.Sequential(
+            Transformer(
+                dim=dim,
+                img_size=img_size // 8,
+                depth=depth[1],
+                heads=heads[1],
+                dim_head=self.dim,
+                mlp_dim=dim * scale_dim,
+                dropout=dropout
+            ),
+            Rearrange('b (h w) c -> b c h w', h=img_size // 8, w=img_size // 8)
+        )
+
+        # stage3
+        in_channels = dim
+        scale = heads[2] // heads[1]
+        dim = scale * dim
+
+        self.stage3_conv_embed = nn.Sequential(
+            nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=dim,
+                kernel_size=kernels[2],
+                stride=strides[2],
+                padding=1
+            ),
+            Rearrange('b c h w -> b (h w) c', h=img_size // 16, w=img_size // 16),
+            nn.LayerNorm(dim)
+        )
+
+        self.stage3_transformer = nn.Sequential(
+            Transformer(
+                dim=dim,
+                img_size=img_size // 16,
+                depth=depth[2],
+                heads=heads[2],
+                dim_head=self.dim,
+                mlp_dim=dim * scale_dim,
+                dropout=dropout
+            ),
+            Rearrange('b (h w) c -> b c h w', h=img_size // 16, w=img_size // 16)
+        )
+
+        # stage4
+        in_channels = dim
+        scale = heads[3] // heads[2]
+        dim = scale * dim
+
+        self.stage4_conv_embed = nn.Sequential(
+            nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=dim,
+                kernel_size=kernels[3],
+                stride=strides[3],
+                padding=1
+            ),
+            Rearrange('b c h w -> b (h w) c', h=img_size // 32, w=img_size // 32),
+            nn.LayerNorm(dim)
+        )
+
+        self.stage4_transformer = nn.Sequential(
+            Transformer(
+                dim=dim, img_size=img_size // 32,
+                depth=depth[3],
+                heads=heads[3],
+                dim_head=self.dim,
+                mlp_dim=dim * scale_dim,
+                dropout=dropout,
+            ),
+            Rearrange('b (h w) c -> b c h w', h=img_size // 32, w=img_size // 32)
+        )
+
+        ### CNN Branch ###
+        self.c_stage1 = c_stage123(in_chans=3, out_chans=64)
+        self.c_stage2 = c_stage123(in_chans=64, out_chans=128)
+        self.c_stage3 = c_stage123(in_chans=128, out_chans=384)
+        self.c_stage4 = c_stage45(in_chans=384, out_chans=512)
+        self.c_stage5 = c_stage45(in_chans=512, out_chans=1024)
+        self.c_max = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.up_conv1 = nn.Conv2d(in_channels=192, out_channels=128, kernel_size=1)
+        self.up_conv2 = nn.Conv2d(in_channels=384, out_channels=512, kernel_size=1)
+
+        ### CTmerge ###
+        self.CTmerge1 = nn.Sequential(
+            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+        )
+        self.CTmerge2 = nn.Sequential(
+            nn.Conv2d(in_channels=320, out_channels=128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+        )
+        self.CTmerge3 = nn.Sequential(
+            nn.Conv2d(in_channels=768, out_channels=512, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(512),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=512, out_channels=384, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(384),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(384),
+            nn.ReLU(),
+        )
+
+        self.CTmerge4 = nn.Sequential(
+            nn.Conv2d(in_channels=896, out_channels=640, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(640),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=640, out_channels=512, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(512),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(512),
+            nn.ReLU(),
+        )
+
+        # decoder
+        self.decoder4 = nn.Sequential(
+            DepthwiseConv2d(
+                in_chans=1408,
+                out_chans=1024,
+                kernel_size=3,
+                stride=1,
+                padding=1
+            ),
+            DepthwiseConv2d(
+                in_chans=1024,
+                out_chans=512,
+                kernel_size=3,
+                stride=1,
+                padding=1
+            ),
+            nn.GELU()
+        )
+        self.decoder3 = nn.Sequential(
+            DepthwiseConv2d(
+                in_chans=896,
+                out_chans=512,
+                kernel_size=3,
+                stride=1,
+                padding=1
+            ),
+            DepthwiseConv2d(
+                in_chans=512,
+                out_chans=384,
+                kernel_size=3,
+                stride=1,
+                padding=1
+            ),
+            nn.GELU()
+        )
+
+        self.decoder2 = nn.Sequential(
+            DepthwiseConv2d(
+                in_chans=576,
+                out_chans=256,
+                kernel_size=3,
+                stride=1,
+                padding=1
+            ),
+            DepthwiseConv2d(
+                in_chans=256,
+                out_chans=192,
+                kernel_size=3,
+                stride=1,
+                padding=1
+            ),
+            nn.GELU()
+        )
+
+        self.decoder1 = nn.Sequential(
+            DepthwiseConv2d(
+                in_chans=256,
+                out_chans=64,
+                kernel_size=3,
+                stride=1,
+                padding=1
+            ),
+            DepthwiseConv2d(
+                in_chans=64,
+                out_chans=16,
+                kernel_size=3,
+                stride=1,
+                padding=1
+            ),
+            nn.GELU()
+        )
+        self.sbr4 = SBR(512)
+        self.sbr3 = SBR(384)
+        self.sbr2 = SBR(192)
+        self.sbr1 = SBR(16)
+
+        self.head = nn.Conv2d(in_channels=16, out_channels=num_classes, kernel_size=1)
+
+    def forward(self, input):
+        ### encoder ###
+        # stage1 = ts1  cat  cs1
+        # t_s1 = self.t_stage1(input)
+        # print(input.shape)
+        # print('++++++++++++++++++++++')
+
+        t_s1 = self.stage1_conv_embed(input)  # 1*3*224*224-->1*3136*64
+
+        # print(t_s1.shape)
+        # print('======================')
+
+        t_s1 = self.stage1_transformer(t_s1)  # 1*3136*64-->1*64*56*56
+
+        # print(t_s1.shape)
+        # print('----------------------')
+
+        c_s1 = self.c_stage1(input)  # 1*3*224*224-->1*64*112*112
+
+        # print(c_s1.shape)
+        # print('!!!!!!!!!!!!!!!!!!!!!!!')
+
+        stage1 = self.CTmerge1(torch.cat([t_s1, self.c_max(c_s1)], dim=1))  # 1*64*56*56   # 拼接两条分支
+
+        # print(stage1.shape)
+        # print('[[[[[[[[[[[[[[[[[[[[[[[')
+
+        # stage2 = ts2 up cs2
+        # t_s2 = self.t_stage2(stage1)
+        t_s2 = self.stage2_conv_embed(stage1)  # 1*64*56*56-->1*784*192   # stage2_conv_embed是转化为序列操作
+
+        # print(t_s2.shape)
+        # print('[[[[[[[[[[[[[[[[[[[[[[[')
+        t_s2 = self.stage2_transformer(t_s2)  # 1*784*192-->1*192*28*28
+        # print(t_s2.shape)
+        # print('+++++++++++++++++++++++++')
+
+        c_s2 = self.c_stage2(c_s1)  # 1*64*112*112-->1*128*56*56
+        stage2 = self.CTmerge2(
+            torch.cat([c_s2, F.interpolate(t_s2, size=c_s2.size()[2:], mode='bilinear', align_corners=True)],
+                      dim=1))  # mode='bilinear'表示使用双线性插值  1*128*56*56
+
+        # stage3 = ts3 cat cs3
+        # t_s3 = self.t_stage3(t_s2)
+        t_s3 = self.stage3_conv_embed(t_s2)  # 1*192*28*28-->1*196*384
+        # print(t_s3.shape)
+        # print('///////////////////////')
+        t_s3 = self.stage3_transformer(t_s3)  # 1*196*384-->1*384*14*14
+        # print(t_s3.shape)
+        # print('....................')
+        c_s3 = self.c_stage3(stage2)  # 1*128*56*56-->1*384*28*28
+        stage3 = self.CTmerge3(torch.cat([t_s3, self.c_max(c_s3)], dim=1))  # 1*384*14*14
+
+        # stage4 = ts4 up cs4
+        # t_s4 = self.t_stage4(stage3)
+        t_s4 = self.stage4_conv_embed(stage3)  # 1*384*14*14-->1*49*384
+        # print(t_s4.shape)
+        # print(';;;;;;;;;;;;;;;;;;;;;;;')
+        t_s4 = self.stage4_transformer(t_s4)  # 1*49*384-->1*384*7*7
+        # print(t_s4.shape)
+        # print('::::::::::::::::::::')
+
+        c_s4 = self.c_stage4(c_s3)  # 1*384*28*28-->1*512*14*14
+        stage4 = self.CTmerge4(
+            torch.cat([c_s4, F.interpolate(t_s4, size=c_s4.size()[2:], mode='bilinear', align_corners=True)],
+                      dim=1))  # 1*512*14*14
+
+        # cs5
+        c_s5 = self.c_stage5(stage4)  # 1*512*14*14-->1*1024*7*7
+
+        ### decoder ###
+        decoder4 = torch.cat([c_s5, t_s4], dim=1)  # 1*1408*7*7
+        decoder4 = self.decoder4(decoder4)  # 1*1408*7*7-->1*512*7*7
+        decoder4 = F.interpolate(decoder4, size=c_s3.size()[2:], mode='bilinear',
+                                 align_corners=True)  # 1*512*7*7-->1*512*28*28
+        decoder4 = self.sbr4(decoder4)  # 1*512*28*28
+        # print(decoder4.shape)
+
+        decoder3 = torch.cat([decoder4, c_s3], dim=1)  # 1*896*28*28
+        decoder3 = self.decoder3(decoder3)  # 1*384*28*28
+        decoder3 = F.interpolate(decoder3, size=t_s2.size()[2:], mode='bilinear', align_corners=True)  # 1*384*28*28
+        decoder3 = self.sbr3(decoder3)  # 1*384*28*28
+        # print(decoder3.shape)
+
+        decoder2 = torch.cat([decoder3, t_s2], dim=1)  # 1*576*28*28
+        decoder2 = self.decoder2(decoder2)  # 1*192*28*28
+        decoder2 = F.interpolate(decoder2, size=c_s1.size()[2:], mode='bilinear', align_corners=True)  # 1*192*112*112
+        decoder2 = self.sbr2(decoder2)  # 1*192*112*112
+        # print(decoder2.shape)
+
+        decoder1 = torch.cat([decoder2, c_s1], dim=1)  # 1*256*112*112
+        decoder1 = self.decoder1(decoder1)  # 1*16*112*112
+        # print(decoder1.shape)
+        final = F.interpolate(decoder1, size=input.size()[2:], mode='bilinear', align_corners=True)  # 1*16*224*224
+        # print(final.shape)
+        # final = self.sbr1(decoder1)
+        # print(final.shape)
+        final = self.head(final)  # 1*3*224*224
+
+        return final
+
+
+if __name__ == '__main__':
+    x = torch.rand(1, 3, 224, 224).cuda()
+    model = Dual_Branch(img_size=224, in_channels=3, num_classes=7).cuda()
+    y = model(x)
+    print(y.shape)
+    # torch.Size([1, 7, 224, 224])

src/models/components/hrcloud.py

@@ -0,0 +1,751 @@
+# 论文地址：https://arxiv.org/abs/2407.07365
+#
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+import os
+
+import numpy as np
+import torch
+import torch._utils
+import torch.nn as nn
+import torch.nn.functional as F
+
+BatchNorm2d = nn.BatchNorm2d
+# BN_MOMENTUM = 0.01
+relu_inplace = True
+BN_MOMENTUM = 0.1
+ALIGN_CORNERS = True
+
+logger = logging.getLogger(__name__)
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+from yacs.config import CfgNode as CN
+import math
+from einops import rearrange
+
+# configs for HRNet48
+HRNET_48 = CN()
+HRNET_48.FINAL_CONV_KERNEL = 1
+
+HRNET_48.STAGE1 = CN()
+HRNET_48.STAGE1.NUM_MODULES = 1
+HRNET_48.STAGE1.NUM_BRANCHES = 1
+HRNET_48.STAGE1.NUM_BLOCKS = [4]
+HRNET_48.STAGE1.NUM_CHANNELS = [64]
+HRNET_48.STAGE1.BLOCK = 'BOTTLENECK'
+HRNET_48.STAGE1.FUSE_METHOD = 'SUM'
+
+HRNET_48.STAGE2 = CN()
+HRNET_48.STAGE2.NUM_MODULES = 1
+HRNET_48.STAGE2.NUM_BRANCHES = 2
+HRNET_48.STAGE2.NUM_BLOCKS = [4, 4]
+HRNET_48.STAGE2.NUM_CHANNELS = [48, 96]
+HRNET_48.STAGE2.BLOCK = 'BASIC'
+HRNET_48.STAGE2.FUSE_METHOD = 'SUM'
+
+HRNET_48.STAGE3 = CN()
+HRNET_48.STAGE3.NUM_MODULES = 4
+HRNET_48.STAGE3.NUM_BRANCHES = 3
+HRNET_48.STAGE3.NUM_BLOCKS = [4, 4, 4]
+HRNET_48.STAGE3.NUM_CHANNELS = [48, 96, 192]
+HRNET_48.STAGE3.BLOCK = 'BASIC'
+HRNET_48.STAGE3.FUSE_METHOD = 'SUM'
+
+HRNET_48.STAGE4 = CN()
+HRNET_48.STAGE4.NUM_MODULES = 3
+HRNET_48.STAGE4.NUM_BRANCHES = 4
+HRNET_48.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]
+HRNET_48.STAGE4.NUM_CHANNELS = [48, 96, 192, 384]
+HRNET_48.STAGE4.BLOCK = 'BASIC'
+HRNET_48.STAGE4.FUSE_METHOD = 'SUM'
+
+HRNET_32 = CN()
+HRNET_32.FINAL_CONV_KERNEL = 1
+
+HRNET_32.STAGE1 = CN()
+HRNET_32.STAGE1.NUM_MODULES = 1
+HRNET_32.STAGE1.NUM_BRANCHES = 1
+HRNET_32.STAGE1.NUM_BLOCKS = [4]
+HRNET_32.STAGE1.NUM_CHANNELS = [64]
+HRNET_32.STAGE1.BLOCK = 'BOTTLENECK'
+HRNET_32.STAGE1.FUSE_METHOD = 'SUM'
+
+HRNET_32.STAGE2 = CN()
+HRNET_32.STAGE2.NUM_MODULES = 1
+HRNET_32.STAGE2.NUM_BRANCHES = 2
+HRNET_32.STAGE2.NUM_BLOCKS = [4, 4]
+HRNET_32.STAGE2.NUM_CHANNELS = [32, 64]
+HRNET_32.STAGE2.BLOCK = 'BASIC'
+HRNET_32.STAGE2.FUSE_METHOD = 'SUM'
+
+HRNET_32.STAGE3 = CN()
+HRNET_32.STAGE3.NUM_MODULES = 4
+HRNET_32.STAGE3.NUM_BRANCHES = 3
+HRNET_32.STAGE3.NUM_BLOCKS = [4, 4, 4]
+HRNET_32.STAGE3.NUM_CHANNELS = [32, 64, 128]
+HRNET_32.STAGE3.BLOCK = 'BASIC'
+HRNET_32.STAGE3.FUSE_METHOD = 'SUM'
+
+HRNET_32.STAGE4 = CN()
+HRNET_32.STAGE4.NUM_MODULES = 3
+HRNET_32.STAGE4.NUM_BRANCHES = 4
+HRNET_32.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]
+HRNET_32.STAGE4.NUM_CHANNELS = [32, 64, 128, 256]
+HRNET_32.STAGE4.BLOCK = 'BASIC'
+HRNET_32.STAGE4.FUSE_METHOD = 'SUM'
+
+HRNET_18 = CN()
+HRNET_18.FINAL_CONV_KERNEL = 1
+
+HRNET_18.STAGE1 = CN()
+HRNET_18.STAGE1.NUM_MODULES = 1
+HRNET_18.STAGE1.NUM_BRANCHES = 1
+HRNET_18.STAGE1.NUM_BLOCKS = [4]
+HRNET_18.STAGE1.NUM_CHANNELS = [64]
+HRNET_18.STAGE1.BLOCK = 'BOTTLENECK'
+HRNET_18.STAGE1.FUSE_METHOD = 'SUM'
+
+HRNET_18.STAGE2 = CN()
+HRNET_18.STAGE2.NUM_MODULES = 1
+HRNET_18.STAGE2.NUM_BRANCHES = 2
+HRNET_18.STAGE2.NUM_BLOCKS = [4, 4]
+HRNET_18.STAGE2.NUM_CHANNELS = [18, 36]
+HRNET_18.STAGE2.BLOCK = 'BASIC'
+HRNET_18.STAGE2.FUSE_METHOD = 'SUM'
+
+HRNET_18.STAGE3 = CN()
+HRNET_18.STAGE3.NUM_MODULES = 4
+HRNET_18.STAGE3.NUM_BRANCHES = 3
+HRNET_18.STAGE3.NUM_BLOCKS = [4, 4, 4]
+HRNET_18.STAGE3.NUM_CHANNELS = [18, 36, 72]
+HRNET_18.STAGE3.BLOCK = 'BASIC'
+HRNET_18.STAGE3.FUSE_METHOD = 'SUM'
+
+HRNET_18.STAGE4 = CN()
+HRNET_18.STAGE4.NUM_MODULES = 3
+HRNET_18.STAGE4.NUM_BRANCHES = 4
+HRNET_18.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]
+HRNET_18.STAGE4.NUM_CHANNELS = [18, 36, 72, 144]
+HRNET_18.STAGE4.BLOCK = 'BASIC'
+HRNET_18.STAGE4.FUSE_METHOD = 'SUM'
+
+
+class PPM(nn.Module):
+    def __init__(self, in_dim, reduction_dim, bins):
+        super(PPM, self).__init__()
+        self.features = []
+        for bin in bins:
+            self.features.append(nn.Sequential(
+                nn.AdaptiveAvgPool2d(bin),
+                nn.Conv2d(in_dim, reduction_dim, kernel_size=1, bias=False),
+                nn.BatchNorm2d(reduction_dim),
+                nn.ReLU(inplace=True)
+            ))
+        self.features = nn.ModuleList(self.features)
+
+    def forward(self, x):
+        x_size = x.size()
+        out = [x]
+        for f in self.features:
+            out.append(F.interpolate(f(x), x_size[2:], mode='bilinear', align_corners=True))
+        return torch.cat(out, 1)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=relu_inplace)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        out = out + residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=1, bias=False)
+        self.bn2 = BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
+                               bias=False)
+        self.bn3 = BatchNorm2d(planes * self.expansion,
+                               momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=relu_inplace)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+            # att = self.downsample(att)
+        out = out + residual
+        out = self.relu(out)
+
+        return out
+
+
+class HighResolutionModule(nn.Module):
+    def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
+                 num_channels, fuse_method, multi_scale_output=True):
+        super(HighResolutionModule, self).__init__()
+        self._check_branches(
+            num_branches, blocks, num_blocks, num_inchannels, num_channels)
+
+        self.num_inchannels = num_inchannels
+        self.fuse_method = fuse_method
+        self.num_branches = num_branches
+
+        self.multi_scale_output = multi_scale_output
+
+        self.branches = self._make_branches(
+            num_branches, blocks, num_blocks, num_channels)
+        self.fuse_layers = self._make_fuse_layers()
+        self.relu = nn.ReLU(inplace=relu_inplace)
+
+    def _check_branches(self, num_branches, blocks, num_blocks,
+                        num_inchannels, num_channels):
+        if num_branches != len(num_blocks):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(
+                num_branches, len(num_blocks))
+            logger.error(error_msg)
+            raise ValueError(error_msg)
+
+        if num_branches != len(num_channels):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(
+                num_branches, len(num_channels))
+            logger.error(error_msg)
+            raise ValueError(error_msg)
+
+        if num_branches != len(num_inchannels):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(
+                num_branches, len(num_inchannels))
+            logger.error(error_msg)
+            raise ValueError(error_msg)
+
+    def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
+                         stride=1):
+        downsample = None
+        if stride != 1 or \
+                self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.num_inchannels[branch_index],
+                          num_channels[branch_index] * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                BatchNorm2d(num_channels[branch_index] * block.expansion,
+                            momentum=BN_MOMENTUM),
+            )
+
+        layers = []
+        layers.append(block(self.num_inchannels[branch_index],
+                            num_channels[branch_index], stride, downsample))
+        self.num_inchannels[branch_index] = \
+            num_channels[branch_index] * block.expansion
+        for i in range(1, num_blocks[branch_index]):
+            layers.append(block(self.num_inchannels[branch_index],
+                                num_channels[branch_index]))
+
+        return nn.Sequential(*layers)
+
+    # 创建平行层
+    def _make_branches(self, num_branches, block, num_blocks, num_channels):
+        branches = []
+
+        for i in range(num_branches):
+            branches.append(
+                self._make_one_branch(i, block, num_blocks, num_channels))
+
+        return nn.ModuleList(branches)
+
+    def _make_fuse_layers(self):
+        if self.num_branches == 1:
+            return None
+        num_branches = self.num_branches  # 3
+        num_inchannels = self.num_inchannels  # [48, 96, 192]
+        fuse_layers = []
+        for i in range(num_branches if self.multi_scale_output else 1):
+            fuse_layer = []
+            for j in range(num_branches):
+                if j > i:
+                    fuse_layer.append(nn.Sequential(
+                        nn.Conv2d(num_inchannels[j],
+                                  num_inchannels[i],
+                                  1,
+                                  1,
+                                  0,
+                                  bias=False),
+                        BatchNorm2d(num_inchannels[i], momentum=BN_MOMENTUM)))
+                elif j == i:
+                    fuse_layer.append(None)
+                else:
+                    conv3x3s = []
+                    for k in range(i - j):
+                        if k == i - j - 1:
+                            num_outchannels_conv3x3 = num_inchannels[i]
+                            conv3x3s.append(nn.Sequential(
+                                nn.Conv2d(num_inchannels[j],
+                                          num_outchannels_conv3x3,
+                                          3, 2, 1, bias=False),
+                                BatchNorm2d(num_outchannels_conv3x3,
+                                            momentum=BN_MOMENTUM)))
+                        else:
+                            num_outchannels_conv3x3 = num_inchannels[j]
+                            conv3x3s.append(nn.Sequential(
+                                nn.Conv2d(num_inchannels[j],
+                                          num_outchannels_conv3x3,
+                                          3, 2, 1, bias=False),
+                                BatchNorm2d(num_outchannels_conv3x3,
+                                            momentum=BN_MOMENTUM),
+                                nn.ReLU(inplace=relu_inplace)))
+                    fuse_layer.append(nn.Sequential(*conv3x3s))
+            fuse_layers.append(nn.ModuleList(fuse_layer))
+
+        return nn.ModuleList(fuse_layers)
+
+    def get_num_inchannels(self):
+        return self.num_inchannels
+
+    def forward(self, x):
+        if self.num_branches == 1:
+            return [self.branches[0](x[0])]
+
+        for i in range(self.num_branches):
+            x[i] = self.branches[i](x[i])
+
+        x_fuse = []
+        for i in range(len(self.fuse_layers)):
+            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
+            for j in range(1, self.num_branches):
+                if i == j:
+                    y = y + x[j]
+                elif j > i:
+                    width_output = x[i].shape[-1]
+                    height_output = x[i].shape[-2]
+                    y = y + F.interpolate(
+                        self.fuse_layers[i][j](x[j]),
+                        size=[height_output, width_output],
+                        mode='bilinear', align_corners=ALIGN_CORNERS)
+                else:
+                    y = y + self.fuse_layers[i][j](x[j])
+            x_fuse.append(self.relu(y))
+
+        return x_fuse
+
+
+blocks_dict = {
+    'BASIC': BasicBlock,
+    'BOTTLENECK': Bottleneck
+}
+
+
+class HRcloudNet(nn.Module):
+
+    def __init__(self, num_classes=2, base_c=48, **kwargs):
+        global ALIGN_CORNERS
+        extra = HRNET_48
+        super(HRcloudNet, self).__init__()
+        ALIGN_CORNERS = True
+        # ALIGN_CORNERS = config.MODEL.ALIGN_CORNERS
+        self.num_classes = num_classes
+        # stem net
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
+                               bias=False)
+        self.bn1 = BatchNorm2d(64, momentum=BN_MOMENTUM)
+        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
+                               bias=False)
+        self.bn2 = BatchNorm2d(64, momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=relu_inplace)
+
+        self.stage1_cfg = extra['STAGE1']
+        num_channels = self.stage1_cfg['NUM_CHANNELS'][0]
+        block = blocks_dict[self.stage1_cfg['BLOCK']]
+        num_blocks = self.stage1_cfg['NUM_BLOCKS'][0]
+        self.layer1 = self._make_layer(block, 64, num_channels, num_blocks)
+        stage1_out_channel = block.expansion * num_channels
+
+        self.stage2_cfg = extra['STAGE2']
+        num_channels = self.stage2_cfg['NUM_CHANNELS']
+        block = blocks_dict[self.stage2_cfg['BLOCK']]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))]
+        self.transition1 = self._make_transition_layer(
+            [stage1_out_channel], num_channels)
+        self.stage2, pre_stage_channels = self._make_stage(
+            self.stage2_cfg, num_channels)
+
+        self.stage3_cfg = extra['STAGE3']
+        num_channels = self.stage3_cfg['NUM_CHANNELS']
+        block = blocks_dict[self.stage3_cfg['BLOCK']]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))]
+        self.transition2 = self._make_transition_layer(
+            pre_stage_channels, num_channels)  # 只在pre[-1]与cur[-1]之间下采样？
+        self.stage3, pre_stage_channels = self._make_stage(
+            self.stage3_cfg, num_channels)
+
+        self.stage4_cfg = extra['STAGE4']
+        num_channels = self.stage4_cfg['NUM_CHANNELS']
+        block = blocks_dict[self.stage4_cfg['BLOCK']]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))]
+        self.transition3 = self._make_transition_layer(
+            pre_stage_channels, num_channels)
+        self.stage4, pre_stage_channels = self._make_stage(
+            self.stage4_cfg, num_channels, multi_scale_output=True)
+        self.out_conv = OutConv(base_c, num_classes)
+        last_inp_channels = int(np.sum(pre_stage_channels))
+
+        self.corr = Corr(nclass=2)
+        self.proj = nn.Sequential(
+            # 512 32
+            nn.Conv2d(720, 48, kernel_size=3, stride=1, padding=1, bias=True),
+            nn.BatchNorm2d(48),
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(0.1),
+        )
+        # self.up1 = Up(base_c * 16, base_c * 8 // factor, bilinear)
+        self.up2 = Up(base_c * 8, base_c * 4, True)
+        self.up3 = Up(base_c * 4, base_c * 2, True)
+        self.up4 = Up(base_c * 2, base_c, True)
+        fea_dim = 720
+        bins = (1, 2, 3, 6)
+        self.ppm = PPM(fea_dim, int(fea_dim / len(bins)), bins)
+        fea_dim *= 2
+        self.cls = nn.Sequential(
+            nn.Conv2d(fea_dim, 512, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(512),
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(p=0.1),
+            nn.Conv2d(512, 2, kernel_size=1)
+        )
+
+    '''
+    转换层的作用有两种情况：
+
+    当前分支数小于之前分支数时，仅对前几个分支进行通道数调整。
+    当前分支数大于之前分支数时，新建一些转换层，对多余的分支进行下采样，改变通道数以适应后续的连接。
+    最终，这些转换层会被组合成一个 nn.ModuleList 对象，并在网络的构建过程中使用。
+    这有助于确保每个分支的通道数在不同阶段之间能够正确匹配，以便进行特征的融合和连接
+    '''
+
+    def _make_transition_layer(
+            self, num_channels_pre_layer, num_channels_cur_layer):
+        # 现在的分支数
+        num_branches_cur = len(num_channels_cur_layer)  # 3
+        # 处理前的分支数
+        num_branches_pre = len(num_channels_pre_layer)  # 2
+
+        transition_layers = []
+        for i in range(num_branches_cur):
+            # 如果当前分支数小于之前分支数，仅针对第一到第二阶段
+            if i < num_branches_pre:
+                # 如果对应层的通道数不一致，则进行转化（
+                if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
+                    transition_layers.append(nn.Sequential(
+
+                        nn.Conv2d(num_channels_pre_layer[i],
+                                  num_channels_cur_layer[i],
+                                  3,
+                                  1,
+                                  1,
+                                  bias=False),
+                        BatchNorm2d(
+                            num_channels_cur_layer[i], momentum=BN_MOMENTUM),
+                        nn.ReLU(inplace=relu_inplace)))
+                else:
+                    transition_layers.append(None)
+            else:  # 在新建层下采样改变通道数
+                conv3x3s = []
+                for j in range(i + 1 - num_branches_pre):  # 3
+                    inchannels = num_channels_pre_layer[-1]
+                    outchannels = num_channels_cur_layer[i] \
+                        if j == i - num_branches_pre else inchannels
+                    conv3x3s.append(nn.Sequential(
+                        nn.Conv2d(
+                            inchannels, outchannels, 3, 2, 1, bias=False),
+                        BatchNorm2d(outchannels, momentum=BN_MOMENTUM),
+                        nn.ReLU(inplace=relu_inplace)))
+                transition_layers.append(nn.Sequential(*conv3x3s))
+
+        return nn.ModuleList(transition_layers)
+
+    '''
+    _make_layer 函数的主要作用是创建一个由多个相同类型的残差块（Residual Block）组成的层。
+    '''
+
+    def _make_layer(self, block, inplanes, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
+            )
+
+        layers = []
+        layers.append(block(inplanes, planes, stride, downsample))
+        inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    # 多尺度融合
+    def _make_stage(self, layer_config, num_inchannels,
+                    multi_scale_output=True):
+        num_modules = layer_config['NUM_MODULES']
+        num_branches = layer_config['NUM_BRANCHES']
+        num_blocks = layer_config['NUM_BLOCKS']
+        num_channels = layer_config['NUM_CHANNELS']
+        block = blocks_dict[layer_config['BLOCK']]
+        fuse_method = layer_config['FUSE_METHOD']
+
+        modules = []
+        for i in range(num_modules):  # 重复4次
+            # multi_scale_output is only used last module
+            if not multi_scale_output and i == num_modules - 1:
+                reset_multi_scale_output = False
+            else:
+                reset_multi_scale_output = True
+            modules.append(
+                HighResolutionModule(num_branches,
+                                     block,
+                                     num_blocks,
+                                     num_inchannels,
+                                     num_channels,
+                                     fuse_method,
+                                     reset_multi_scale_output)
+            )
+            num_inchannels = modules[-1].get_num_inchannels()
+
+        return nn.Sequential(*modules), num_inchannels
+
+    def forward(self, input, need_fp=True, use_corr=True):
+        # from ipdb import set_trace
+        # set_trace()
+        x = self.conv1(input)
+        x = self.bn1(x)
+        x = self.relu(x)
+        # x_176 = x
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+        x = self.layer1(x)
+
+        x_list = []
+        for i in range(self.stage2_cfg['NUM_BRANCHES']):  # 2
+            if self.transition1[i] is not None:
+                x_list.append(self.transition1[i](x))
+            else:
+                x_list.append(x)
+        y_list = self.stage2(x_list)
+        # Y1
+        x_list = []
+        for i in range(self.stage3_cfg['NUM_BRANCHES']):
+            if self.transition2[i] is not None:
+                if i < self.stage2_cfg['NUM_BRANCHES']:
+                    x_list.append(self.transition2[i](y_list[i]))
+                else:
+                    x_list.append(self.transition2[i](y_list[-1]))
+            else:
+                x_list.append(y_list[i])
+        y_list = self.stage3(x_list)
+
+        x_list = []
+        for i in range(self.stage4_cfg['NUM_BRANCHES']):
+            if self.transition3[i] is not None:
+                if i < self.stage3_cfg['NUM_BRANCHES']:
+                    x_list.append(self.transition3[i](y_list[i]))
+                else:
+                    x_list.append(self.transition3[i](y_list[-1]))
+            else:
+                x_list.append(y_list[i])
+        x = self.stage4(x_list)
+        dict_return = {}
+        # Upsampling
+        x0_h, x0_w = x[0].size(2), x[0].size(3)
+
+        x3 = F.interpolate(x[3], size=(x0_h, x0_w), mode='bilinear', align_corners=ALIGN_CORNERS)
+        # x = self.stage3_(x)
+        x[2] = self.up2(x[3], x[2])
+        x2 = F.interpolate(x[2], size=(x0_h, x0_w), mode='bilinear', align_corners=ALIGN_CORNERS)
+        # x = self.stage2_(x)
+        x[1] = self.up3(x[2], x[1])
+        x1 = F.interpolate(x[1], size=(x0_h, x0_w), mode='bilinear', align_corners=ALIGN_CORNERS)
+        x[0] = self.up4(x[1], x[0])
+        xk = torch.cat([x[0], x1, x2, x3], 1)
+        # PPM
+        feat = self.ppm(xk)
+        x = self.cls(feat)
+        # fp分支
+        if need_fp:
+            logits = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)
+            # logits = self.out_conv(torch.cat((x, nn.Dropout2d(0.5)(x))))
+            out = logits
+            out_fp = logits
+            if use_corr:
+                proj_feats = self.proj(xk)
+                corr_out = self.corr(proj_feats, out)
+                corr_out = F.interpolate(corr_out, size=(352, 352), mode="bilinear", align_corners=True)
+                dict_return['corr_out'] = corr_out
+            dict_return['out'] = out
+            dict_return['out_fp'] = out_fp
+
+            return dict_return['out']
+
+        out = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)
+        if use_corr:  # True
+            proj_feats = self.proj(xk)
+            # 计算
+            corr_out = self.corr(proj_feats, out)
+            corr_out = F.interpolate(corr_out, size=(352, 352), mode="bilinear", align_corners=True)
+            dict_return['corr_out'] = corr_out
+        dict_return['out'] = out
+        return dict_return['out']
+        # return x
+
+    def init_weights(self, pretrained='', ):
+        logger.info('=> init weights from normal distribution')
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.normal_(m.weight, std=0.001)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+        if os.path.isfile(pretrained):
+            pretrained_dict = torch.load(pretrained)
+            logger.info('=> loading pretrained model {}'.format(pretrained))
+            model_dict = self.state_dict()
+            pretrained_dict = {k: v for k, v in pretrained_dict.items()
+                               if k in model_dict.keys()}
+            for k, _ in pretrained_dict.items():
+                logger.info(
+                    '=> loading {} pretrained model {}'.format(k, pretrained))
+            model_dict.update(pretrained_dict)
+            self.load_state_dict(model_dict)
+
+
+class OutConv(nn.Sequential):
+    def __init__(self, in_channels, num_classes):
+        super(OutConv, self).__init__(
+            nn.Conv2d(720, num_classes, kernel_size=1)
+        )
+
+
+class DoubleConv(nn.Sequential):
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        if mid_channels is None:
+            mid_channels = out_channels
+        super(DoubleConv, self).__init__(
+            nn.Conv2d(in_channels + out_channels, mid_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+
+class Up(nn.Module):
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super(Up, self).__init__()
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels)
+
+    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+        x1 = self.up(x1)
+        # [N, C, H, W]
+        diff_y = x2.size()[2] - x1.size()[2]
+        diff_x = x2.size()[3] - x1.size()[3]
+
+        # padding_left, padding_right, padding_top, padding_bottom
+        x1 = F.pad(x1, [diff_x // 2, diff_x - diff_x // 2,
+                        diff_y // 2, diff_y - diff_y // 2])
+
+        x = torch.cat([x2, x1], dim=1)
+        x = self.conv(x)
+        return x
+
+
+class Corr(nn.Module):
+    def __init__(self, nclass=2):
+        super(Corr, self).__init__()
+        self.nclass = nclass
+        self.conv1 = nn.Conv2d(48, self.nclass, kernel_size=1, stride=1, padding=0, bias=True)
+        self.conv2 = nn.Conv2d(48, self.nclass, kernel_size=1, stride=1, padding=0, bias=True)
+
+    def forward(self, feature_in, out):
+        # in torch.Size([4, 32, 22, 22])
+        # out = [4 2 352 352]
+        h_in, w_in = math.ceil(feature_in.shape[2] / (1)), math.ceil(feature_in.shape[3] / (1))
+        out = F.interpolate(out.detach(), (h_in, w_in), mode='bilinear', align_corners=True)
+        feature = F.interpolate(feature_in, (h_in, w_in), mode='bilinear', align_corners=True)
+        f1 = rearrange(self.conv1(feature), 'n c h w -> n c (h w)')
+        f2 = rearrange(self.conv2(feature), 'n c h w -> n c (h w)')
+        out_temp = rearrange(out, 'n c h w -> n c (h w)')
+        corr_map = torch.matmul(f1.transpose(1, 2), f2) / torch.sqrt(torch.tensor(f1.shape[1]).float())
+        corr_map = F.softmax(corr_map, dim=-1)
+        # out_temp 2 2 484
+        # corr_map 4 484 484
+        out = rearrange(torch.matmul(out_temp, corr_map), 'n c (h w) -> n c h w', h=h_in, w=w_in)
+        # out torch.Size([4, 2, 22, 22])
+        return out
+
+
+if __name__ == '__main__':
+    input = torch.randn(4, 3, 352, 352)
+    cloud = HRcloudNet(num_classes=2)
+    output = cloud(input)
+    print(output.shape)
+    # torch.Size([4, 2, 352, 352]) torch.Size([4, 2, 352, 352]) torch.Size([4, 2, 352, 352])

src/models/components/lnn.py

@@ -0,0 +1,23 @@
+import torch
+from torch import nn
+
+
+class LNN(nn.Module):
+    # 创建一个全连接网络用于手写数字识别，并通过一个参数dim控制中间层的维度
+    def __init__(self, dim=32):
+        super(LNN, self).__init__()
+        self.fc1 = nn.Linear(28 * 28, dim)
+        self.fc2 = nn.Linear(dim, 10)
+    
+    def forward(self, x):
+        x = x.view(-1, x.shape[1] * x.shape[2] * x.shape[3])
+        x = torch.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+
+if __name__ == "__main__":
+    input = torch.randn(2, 1, 28, 28)
+    model = LNN()
+    output = model(input)
+    assert output.shape == (2, 10)

src/models/components/mcdnet.py

@@ -0,0 +1,448 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/7/21 下午3:51
+# @Author  : xiaoshun
+# @Email   : 3038523973@qq.com
+# @File    : mcdnet.py
+# @Software: PyCharm
+import cv2
+import image_dehazer
+import numpy as np
+# 论文地址：https://www.sciencedirect.com/science/article/pii/S1569843224001742?via%3Dihub
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class _DPFF(nn.Module):
+    def __init__(self, in_channels) -> None:
+        super(_DPFF, self).__init__()
+        self.cbr1 = nn.Conv2d(in_channels * 2, in_channels, 1, 1, bias=False)
+        self.cbr2 = nn.Conv2d(in_channels * 2, in_channels, 1, 1, bias=False)
+        # self.sigmoid = nn.Sigmoid()
+        self.cbr3 = nn.Conv2d(in_channels, in_channels, 1, 1, bias=False)
+        self.cbr4 = nn.Conv2d(in_channels * 2, in_channels, 1, 1, bias=False)
+
+    def forward(self, feature1, feature2):
+        d1 = torch.abs(feature1 - feature2)
+        d2 = self.cbr1(torch.cat([feature1, feature2], dim=1))
+        d = torch.cat([d1, d2], dim=1)
+        d = self.cbr2(d)
+        # d = self.sigmoid(d)
+
+        v1, v2 = self.cbr3(feature1), self.cbr3(feature2)
+        v1, v2 = v1 * d, v2 * d
+        features = torch.cat([v1, v2], dim=1)
+        features = self.cbr4(features)
+
+        return features
+
+
+class DPFF(nn.Module):
+    def __init__(self, layer_channels) -> None:
+        super(DPFF, self).__init__()
+        self.cfes = nn.ModuleList()
+        for layer_channel in layer_channels:
+            self.cfes.append(_DPFF(layer_channel))
+
+    def forward(self, features1, features2):
+        outputs = []
+        for feature1, feature2, cfe in zip(features1, features2, self.cfes):
+            outputs.append(cfe(feature1, feature2))
+        return outputs
+
+
+class DirectDPFF(nn.Module):
+    def __init__(self, layer_channels) -> None:
+        super(DirectDPFF, self).__init__()
+        self.fusions = nn.ModuleList(
+            [nn.Conv2d(layer_channel * 2, layer_channel, 1, 1) for layer_channel in layer_channels]
+        )
+
+    def forward(self, features1, features2):
+        outputs = []
+        for feature1, feature2, fusion in zip(features1, features2, self.fusions):
+            feature = torch.cat([feature1, feature2], dim=1)
+            outputs.append(fusion(feature))
+        return outputs
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, input_size, output_size, kernel_size=4, stride=2, padding=1, bias=True,
+                 bn=False, activation=True, maxpool=True):
+        super(ConvBlock, self).__init__()
+        self.module = []
+        if maxpool:
+            down = nn.Sequential(
+                *[
+                    nn.MaxPool2d(2),
+                    nn.Conv2d(input_size, output_size, 1, 1, 0, bias=bias)
+                ]
+            )
+        else:
+            down = nn.Conv2d(input_size, output_size, kernel_size, stride, padding, bias=bias)
+        self.module.append(down)
+        if bn:
+            self.module.append(nn.BatchNorm2d(output_size))
+        if activation:
+            self.module.append(nn.PReLU())
+        self.module = nn.Sequential(*self.module)
+
+    def forward(self, x):
+        out = self.module(x)
+
+        return out
+
+
+class DeconvBlock(nn.Module):
+    def __init__(self, input_size, output_size, kernel_size=4, stride=2, padding=1, bias=True,
+                 bn=False, activation=True, bilinear=True):
+        super(DeconvBlock, self).__init__()
+        self.module = []
+        if bilinear:
+            deconv = nn.Sequential(
+                *[
+                    nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
+                    nn.Conv2d(input_size, output_size, 1, 1, 0, bias=bias)
+                ]
+            )
+        else:
+            deconv = nn.ConvTranspose2d(input_size, output_size, kernel_size, stride, padding, bias=bias)
+        self.module.append(deconv)
+        if bn:
+            self.module.append(nn.BatchNorm2d(output_size))
+        if activation:
+            self.module.append(nn.PReLU())
+        self.module = nn.Sequential(*self.module)
+
+    def forward(self, x):
+        out = self.module(x)
+
+        return out
+
+
+class FusionBlock(torch.nn.Module):
+    def __init__(self, num_filter, num_ft, kernel_size=4, stride=2, padding=1, bias=True, maxpool=False,
+                 bilinear=False):
+        super(FusionBlock, self).__init__()
+        self.num_ft = num_ft
+        self.up_convs = nn.ModuleList()
+        self.down_convs = nn.ModuleList()
+        for i in range(self.num_ft):
+            self.up_convs.append(
+                DeconvBlock(num_filter // (2 ** i), num_filter // (2 ** (i + 1)), kernel_size, stride, padding,
+                            bias=bias, bilinear=bilinear)
+            )
+            self.down_convs.append(
+                ConvBlock(num_filter // (2 ** (i + 1)), num_filter // (2 ** i), kernel_size, stride, padding, bias=bias,
+                          maxpool=maxpool)
+            )
+
+    def forward(self, ft_l, ft_h_list):
+        ft_fusion = ft_l
+        for i in range(len(ft_h_list)):
+            ft = ft_fusion
+            for j in range(self.num_ft - i):
+                ft = self.up_convs[j](ft)
+            ft = ft - ft_h_list[i]
+            for j in range(self.num_ft - i):
+                ft = self.down_convs[self.num_ft - i - j - 1](ft)
+            ft_fusion = ft_fusion + ft
+
+        return ft_fusion
+
+
+class ConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, bias=True):
+        super(ConvLayer, self).__init__()
+        reflection_padding = kernel_size // 2
+        self.reflection_pad = nn.ReflectionPad2d(reflection_padding)
+        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, bias=bias)
+
+    def forward(self, x):
+        out = self.reflection_pad(x)
+        out = self.conv2d(out)
+        return out
+
+
+class UpsampleConvLayer(torch.nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride):
+        super(UpsampleConvLayer, self).__init__()
+        self.conv2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride)
+
+    def forward(self, x):
+        out = self.conv2d(x)
+        return out
+
+
+class AddRelu(nn.Module):
+    """It is for adding two feed forwards to the output of the two following conv layers in expanding path
+    """
+
+    def __init__(self) -> None:
+        super(AddRelu, self).__init__()
+        self.relu = nn.PReLU()
+
+    def forward(self, input_tensor1, input_tensor2, input_tensor3):
+        x = input_tensor1 + input_tensor2 + input_tensor3
+        return self.relu(x)
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super(BasicBlock, self).__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.conv1 = ConvLayer(in_channels, mid_channels, kernel_size=3, stride=1)
+        self.bn1 = nn.BatchNorm2d(mid_channels, momentum=0.1)
+        self.relu = nn.PReLU()
+
+        self.conv2 = ConvLayer(mid_channels, out_channels, kernel_size=3, stride=1)
+        self.bn2 = nn.BatchNorm2d(out_channels, momentum=0.1)
+
+        self.conv3 = ConvLayer(in_channels, out_channels, kernel_size=1, stride=1)
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        residual = self.conv3(x)
+
+        out = out + residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(Bottleneck, self).__init__()
+        self.conv1 = ConvLayer(in_channels, out_channels, kernel_size=3, stride=1)
+        self.bn1 = nn.BatchNorm2d(out_channels, momentum=0.1)
+
+        self.conv2 = ConvLayer(out_channels, out_channels, kernel_size=3, stride=1)
+        self.bn2 = nn.BatchNorm2d(out_channels, momentum=0.1)
+
+        self.conv3 = ConvLayer(out_channels, out_channels, kernel_size=3, stride=1)
+        self.bn3 = nn.BatchNorm2d(out_channels, momentum=0.1)
+
+        self.conv4 = ConvLayer(in_channels, out_channels, kernel_size=1, stride=1)
+
+        self.relu = nn.PReLU()
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        residual = self.conv4(x)
+
+        out = out + residual
+        out = self.relu(out)
+
+        return out
+
+
+class PPM(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(PPM, self).__init__()
+
+        self.pool_sizes = [1, 2, 3, 6]  # subregion size in each level
+        self.num_levels = len(self.pool_sizes)  # number of pyramid levels
+
+        self.conv_layers = nn.ModuleList()
+        for i in range(self.num_levels):
+            self.conv_layers.append(nn.Sequential(
+                nn.AdaptiveAvgPool2d(output_size=self.pool_sizes[i]),
+                nn.Conv2d(in_channels, in_channels // self.num_levels, kernel_size=1),
+                nn.BatchNorm2d(in_channels // self.num_levels),
+                nn.ReLU(inplace=True)
+            ))
+        self.out_conv = nn.Conv2d(in_channels * 2, out_channels, kernel_size=1, stride=1)
+
+    def forward(self, x):
+        input_size = x.size()[2:]  # get input size
+        output = [x]
+
+        # pyramid pooling
+        for i in range(self.num_levels):
+            out = self.conv_layers[i](x)
+            out = F.interpolate(out, size=input_size, mode='bilinear', align_corners=True)
+            output.append(out)
+
+        # concatenate features from different levels
+        output = torch.cat(output, dim=1)
+        output = self.out_conv(output)
+
+        return output
+
+
+class MCDNet(nn.Module):
+    def __init__(self, in_channels=4, num_classes=4, maxpool=False, bilinear=False) -> None:
+        super(MCDNet, self).__init__()
+        level = 1
+        # encoder
+        self.conv_input = ConvLayer(in_channels, 32 * level, kernel_size=3, stride=2)
+
+        self.dense0 = BasicBlock(32 * level, 32 * level)
+        self.conv2x = ConvLayer(32 * level, 64 * level, kernel_size=3, stride=2)
+
+        self.dense1 = BasicBlock(64 * level, 64 * level)
+        self.conv4x = ConvLayer(64 * level, 128 * level, kernel_size=3, stride=2)
+
+        self.dense2 = BasicBlock(128 * level, 128 * level)
+        self.conv8x = ConvLayer(128 * level, 256 * level, kernel_size=3, stride=2)
+
+        self.dense3 = BasicBlock(256 * level, 256 * level)
+        self.conv16x = ConvLayer(256 * level, 512 * level, kernel_size=3, stride=2)
+
+        self.dense4 = PPM(512 * level, 512 * level)
+
+        # dpff
+        self.dpffm = DPFF([32, 64, 128, 256, 512])
+
+        # decoder
+        self.convd16x = UpsampleConvLayer(512 * level, 256 * level, kernel_size=3, stride=2)
+        self.fusion4 = FusionBlock(256 * level, 3, maxpool=maxpool, bilinear=bilinear)
+        self.dense_4 = Bottleneck(512 * level, 256 * level)
+        self.add_block4 = AddRelu()
+
+        self.convd8x = UpsampleConvLayer(256 * level, 128 * level, kernel_size=3, stride=2)
+        self.fusion3 = FusionBlock(128 * level, 2, maxpool=maxpool, bilinear=bilinear)
+        self.dense_3 = Bottleneck(256 * level, 128 * level)
+        self.add_block3 = AddRelu()
+
+        self.convd4x = UpsampleConvLayer(128 * level, 64 * level, kernel_size=3, stride=2)
+        self.fusion2 = FusionBlock(64 * level, 1, maxpool=maxpool, bilinear=bilinear)
+        self.dense_2 = Bottleneck(128 * level, 64 * level)
+        self.add_block2 = AddRelu()
+
+        self.convd2x = UpsampleConvLayer(64 * level, 32 * level, kernel_size=3, stride=2)
+        self.dense_1 = Bottleneck(64 * level, 32 * level)
+        self.add_block1 = AddRelu()
+
+        self.head = UpsampleConvLayer(32 * level, num_classes, kernel_size=3, stride=2)
+        self.apply(self._weights_init)
+
+    def _weights_init(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.xavier_normal_(m.weight)
+            nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+        elif isinstance(m, nn.BatchNorm2d):
+            nn.init.constant_(m.weight, 1)
+            nn.init.constant_(m.bias, 0)
+
+    def get_lr_data(self, x: torch.Tensor) -> torch.Tensor:
+        images = x.cpu().permute(0, 2, 3, 1).numpy()
+        batch_size = images.shape[0]
+        lr = []
+        for i in range(batch_size):
+            lr_image = cv2.cvtColor(images[i], cv2.COLOR_RGB2BGR)
+            lr_image = image_dehazer.remove_haze(lr_image, showHazeTransmissionMap=False)[0]
+            lr_image = cv2.cvtColor(lr_image, cv2.COLOR_BGR2RGB)
+            max_pix = np.max(lr_image)
+            min_pix = np.min(lr_image)
+            lr_image = (lr_image - min_pix) / (max_pix - min_pix)
+            lr_image = np.clip(lr_image, 0, 1)
+            lr_tensor = torch.from_numpy(lr_image).permute(2, 0, 1).float()
+            lr.append(lr_tensor)
+        return torch.stack(lr, dim=0).to(x.device)
+
+    def forward(self, x1):
+        x2 = self.get_lr_data(x1)
+        # encoder1
+        res1x_1 = self.conv_input(x1)
+        res1x_1 = self.dense0(res1x_1)
+
+        res2x_1 = self.conv2x(res1x_1)
+        res2x_1 = self.dense1(res2x_1)
+
+        res4x_1 = self.conv4x(res2x_1)
+        res4x_1 = self.dense2(res4x_1)
+
+        res8x_1 = self.conv8x(res4x_1)
+        res8x_1 = self.dense3(res8x_1)
+
+        res16x_1 = self.conv16x(res8x_1)
+        res16x_1 = self.dense4(res16x_1)
+
+        # encoder2
+        res1x_2 = self.conv_input(x2)
+        res1x_2 = self.dense0(res1x_2)
+
+        res2x_2 = self.conv2x(res1x_2)
+        res2x_2 = self.dense1(res2x_2)
+
+        res4x_2 = self.conv4x(res2x_2)
+        res4x_2 = self.dense2(res4x_2)
+
+        res8x_2 = self.conv8x(res4x_2)
+        res8x_2 = self.dense3(res8x_2)
+
+        res16x_2 = self.conv16x(res8x_2)
+        res16x_2 = self.dense4(res16x_2)
+
+        # dual-perspective feature fusion
+        res1x, res2x, res4x, res8x, res16x = self.dpffm(
+            [res1x_1, res2x_1, res4x_1, res8x_1, res16x_1],
+            [res1x_2, res2x_2, res4x_2, res8x_2, res16x_2]
+        )
+
+        # decoder
+        res8x1 = self.convd16x(res16x)
+        res8x1 = F.interpolate(res8x1, res8x.size()[2:], mode='bilinear')
+        res8x2 = self.fusion4(res8x, [res1x, res2x, res4x])
+        res8x2 = torch.cat([res8x1, res8x2], dim=1)
+        res8x2 = self.dense_4(res8x2)
+        res8x2 = self.add_block4(res8x1, res8x, res8x2)
+
+        res4x1 = self.convd8x(res8x2)
+        res4x1 = F.interpolate(res4x1, res4x.size()[2:], mode='bilinear')
+        res4x2 = self.fusion3(res4x, [res1x, res2x])
+        res4x2 = torch.cat([res4x1, res4x2], dim=1)
+        res4x2 = self.dense_3(res4x2)
+        res4x2 = self.add_block3(res4x1, res4x, res4x2)
+
+        res2x1 = self.convd4x(res4x2)
+        res2x1 = F.interpolate(res2x1, res2x.size()[2:], mode='bilinear')
+        res2x2 = self.fusion2(res2x, [res1x])
+        res2x2 = torch.cat([res2x1, res2x2], dim=1)
+        res2x2 = self.dense_2(res2x2)
+        res2x2 = self.add_block2(res2x1, res2x, res2x2)
+
+        res1x1 = self.convd2x(res2x2)
+        res1x1 = F.interpolate(res1x1, res1x.size()[2:], mode='bilinear')
+        res1x2 = torch.cat([res1x1, res1x], dim=1)
+        res1x2 = self.dense_1(res1x2)
+        res1x2 = self.add_block1(res1x1, res1x, res1x2)
+
+        out = self.head(res1x2)
+        out = F.interpolate(out, x1.size()[2:], mode='bilinear')
+
+        return out
+
+
+def lr_lambda(epoch):
+    return (1 - epoch / 50) ** 0.9
+
+
+if __name__ == "__main__":
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    # device = 'cpu'
+    model = MCDNet(in_channels=3, num_classes=7).to(device)
+    fake_img = torch.randn(size=(2, 3, 256, 256)).to(device)
+    out = model(fake_img).detach().cpu()
+    print(out.shape)
+#     torch.Size([2, 7, 256, 256])

src/models/components/scnn.py

@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/7/21 下午5:11
+# @Author  : xiaoshun
+# @Email   : 3038523973@qq.com
+# @File    : scnn.py
+# @Software: PyCharm
+
+# 论文地址:https://www.sciencedirect.com/science/article/abs/pii/S0924271624000352?via%3Dihub#fn1
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SCNNNet(nn.Module):
+    def __init__(self, in_channels=3, num_classes=2, dropout_p=0.5):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=1)
+        self.conv2 = nn.Conv2d(64, num_classes, kernel_size=1)
+        self.conv3 = nn.Conv2d(num_classes, num_classes, kernel_size=3, padding=1)
+        self.dropout = nn.Dropout2d(p=dropout_p)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = self.dropout(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        return x
+
+
+if __name__ == '__main__':
+    model = SCNNNet(num_classes=7)
+    fake_img = torch.randn((2, 3, 224, 224))
+    out = model(fake_img)
+    print(out.shape)
+    # torch.Size([2, 7, 224, 224])

src/models/components/unet.py

@@ -0,0 +1,63 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class UNet(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(UNet, self).__init__()
+
+        def conv_block(in_channels, out_channels):
+            return nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+                nn.ReLU(inplace=True)
+            )
+
+        self.encoder1 = conv_block(in_channels, 64)
+        self.encoder2 = conv_block(64, 128)
+        self.encoder3 = conv_block(128, 256)
+        self.encoder4 = conv_block(256, 512)
+        self.bottleneck = conv_block(512, 1024)
+
+        self.upconv4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
+        self.decoder4 = conv_block(1024, 512)
+        self.upconv3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
+        self.decoder3 = conv_block(512, 256)
+        self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
+        self.decoder2 = conv_block(256, 128)
+        self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
+        self.decoder1 = conv_block(128, 64)
+
+        self.final = nn.Conv2d(64, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        enc1 = self.encoder1(x)
+        enc2 = self.encoder2(F.max_pool2d(enc1, kernel_size=2, stride=2))
+        enc3 = self.encoder3(F.max_pool2d(enc2, kernel_size=2, stride=2))
+        enc4 = self.encoder4(F.max_pool2d(enc3, kernel_size=2, stride=2))
+        bottleneck = self.bottleneck(F.max_pool2d(enc4, kernel_size=2, stride=2))
+
+        dec4 = self.upconv4(bottleneck)
+        dec4 = torch.cat((dec4, enc4), dim=1)
+        dec4 = self.decoder4(dec4)
+        dec3 = self.upconv3(dec4)
+        dec3 = torch.cat((dec3, enc3), dim=1)
+        dec3 = self.decoder3(dec3)
+        dec2 = self.upconv2(dec3)
+        dec2 = torch.cat((dec2, enc2), dim=1)
+        dec2 = self.decoder2(dec2)
+        dec1 = self.upconv1(dec2)
+        dec1 = torch.cat((dec1, enc1), dim=1)
+        dec1 = self.decoder1(dec1)
+
+        return self.final(dec1)
+
+if __name__ == "__main__":
+     model = UNet(in_channels=3,out_channels=7)
+     fake_img = torch.rand(size=(2,3,224,224))
+     print(fake_img.shape)
+    #  torch.Size([2, 3, 224, 224])
+     out = model(fake_img)
+     print(out.shape)
+    #  torch.Size([2, 7, 224, 224])

src/models/components/vae.py

@@ -0,0 +1,152 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from contextlib import contextmanager
+from typing import List, Dict
+from src.plugin.taming_transformers.taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
+
+from src.plugin.ldm.modules.diffusionmodules.model import Encoder, Decoder
+from src.plugin.ldm.modules.distributions.distributions import DiagonalGaussianDistribution
+
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class AutoencoderKL(nn.Module):
+    def __init__(
+        self,
+        double_z: bool = True,
+        z_channels: int = 3,
+        resolution: int = 512,
+        in_channels: int = 3,
+        out_ch: int = 3,
+        ch: int = 128,
+        ch_mult: List = [1, 2, 4, 4],
+        num_res_blocks: int = 2,
+        attn_resolutions: List = [],
+        dropout: float = 0.0,
+        embed_dim: int = 3,
+        ckpt_path: str = None,
+        ignore_keys: List = [],
+    ):
+        super(AutoencoderKL, self).__init__()
+        ddconfig = {
+            "double_z": double_z,
+            "z_channels": z_channels,
+            "resolution": resolution,
+            "in_channels": in_channels,
+            "out_ch": out_ch,
+            "ch": ch,
+            "ch_mult": ch_mult,
+            "num_res_blocks": num_res_blocks,
+            "attn_resolutions": attn_resolutions,
+            "dropout": dropout
+        }
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = Decoder(**ddconfig)
+        assert ddconfig["double_z"]
+        self.quant_conv = nn.Conv2d(
+            2 * ddconfig["z_channels"], 2 * embed_dim, 1)
+        self.post_quant_conv = nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+        self.embed_dim = embed_dim
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+
+    def init_from_ckpt(self, path, ignore_keys=list()):
+        sd = torch.load(path, map_location="cpu")["state_dict"]
+        keys = list(sd.keys())
+        for k in keys:
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    print(f"Deleting key {k} from state_dict.")
+                    del sd[k]
+        self.load_state_dict(sd, strict=False)
+        print(f"Restored from {path}")
+
+    def encode(self, x):
+        h = self.encoder(x)  # B, C, h, w
+        moments = self.quant_conv(h)  # B, 6, h, w
+        posterior = DiagonalGaussianDistribution(moments)
+        return posterior  # 分布
+
+    def decode(self, z):
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+        return dec
+
+    def forward(self, input, sample_posterior=True):
+        posterior = self.encode(input)  # 高斯分布
+        if sample_posterior:
+            z = posterior.sample()  # 采样
+        else:
+            z = posterior.mode()
+        dec = self.decode(z)
+        last_layer_weight = self.decoder.conv_out.weight
+        return dec, posterior, last_layer_weight
+
+
+if __name__ == '__main__':
+    # Test the input and output shapes of the model
+    model = AutoencoderKL()
+    x = torch.randn(1, 3, 512, 512)
+    dec, posterior, last_layer_weight = model(x)
+
+    assert dec.shape == (1, 3, 512, 512)
+    assert posterior.sample().shape == posterior.mode().shape == (1, 3, 64, 64)
+    assert last_layer_weight.shape == (3, 128, 3, 3)
+
+    # Plot the latent space and the reconstruction from the pretrained model
+    model = AutoencoderKL(ckpt_path="/mnt/chongqinggeminiceph1fs/geminicephfs/wx-mm-spr-xxxx/zouxuechao/Collaborative-Diffusion/outputs/512_vae/2024-06-27T06-02-04_512_vae/checkpoints/epoch=000036.ckpt")
+    model.eval()
+    image_path = "data/celeba/image/image_512_downsampled_from_hq_1024/0.jpg"
+
+    from PIL import Image
+    import numpy as np
+    from src.data.components.celeba import DalleTransformerPreprocessor
+    from src.data.components.celeba import CelebA
+    image = Image.open(image_path).convert('RGB')
+    image = np.array(image).astype(np.uint8)
+    import copy
+    original = copy.deepcopy(image)
+    transform = DalleTransformerPreprocessor(size=512, phase='test')
+    image = transform(image=image)['image']
+    image = image.astype(np.float32)/127.5 - 1.0
+    image = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0)
+
+    dec, posterior, last_layer_weight = model(image)
+
+    # original image
+    plt.subplot(1, 3, 1)
+    plt.imshow(original)
+    plt.title("Original")
+    plt.axis("off")
+
+    # sampled image from the latent space
+    plt.subplot(1, 3, 2)
+    x = model.decode(posterior.sample())
+    x = (x+1)/2
+    x = x.squeeze(0).permute(1, 2, 0).cpu()
+    x = x.detach().numpy()
+    x = x.clip(0, 1)
+    x = (x*255).astype(np.uint8)
+    plt.imshow(x)
+    plt.title("Sampled")
+    plt.axis("off")
+
+    # reconstructed image
+    plt.subplot(1, 3, 3)
+    x = dec
+    x = (x+1)/2
+    x = x.squeeze(0).permute(1, 2, 0).cpu()
+    x = x.detach().numpy()
+    x = x.clip(0, 1)
+    x = (x*255).astype(np.uint8)
+    plt.imshow(x)
+    plt.title("Reconstructed")
+    plt.axis("off")
+
+    plt.tight_layout()
+    plt.savefig("vae_reconstruction.png")