diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..e5f95609987a8d14ce77a31d975d41dd0fc35f16 --- /dev/null +++ b/app.py @@ -0,0 +1,273 @@ +from typing import List +from glob import glob +import numpy as np +from PIL import Image +from mmseg.models.segmentors.encoder_decoder import EncoderDecoder +import gradio as gr +import torch +import os +from models.cdnetv1 import CDnetV1 +from models.cdnetv2 import CDnetV2 +from models.dbnet import DBNet +from models.hrcloudnet import HRCloudNet +from models.kappamask import KappaMask +from models.mcdnet import MCDNet +from models.scnn import SCNN +from models.unetmobv2 import UNetMobV2 + + +class CloudAdapterGradio: + def __init__(self, device="cpu", example_inputs=None, num_classes=2, palette=None, other_model_weight_path=None): + self.device = device + self.example_inputs = example_inputs + self.img_size = 256 if num_classes == 2 else 512 + self.palette = palette + self.legend = self.html_legend(num_classes=num_classes) + + self.other_models = { + "cdnetv1": CDnetV1(num_classes=num_classes).to(self.device), + "cdnetv2": CDnetV2(num_classes=num_classes).to(self.device), + "hrcloudnet": HRCloudNet(num_classes=num_classes).to(self.device), + "mcdnet": MCDNet(in_channels=3, num_classes=num_classes).to(self.device), + "scnn": SCNN(num_classes=num_classes).to(self.device), + "dbnet": DBNet(img_size=self.img_size, in_channels=3, num_classes=num_classes).to( + self.device + ), + "unetmobv2": UNetMobV2(num_classes=num_classes).to(self.device), + "kappamask": KappaMask(num_classes=num_classes, in_channels=3).to(self.device) + } + self.name_mapping = { + "KappaMask": "kappamask", + "CDNetv1": "cdnetv1", + "CDNetv2": "cdnetv2", + "HRCloudNet": "hrcloudnet", + "MCDNet": "mcdnet", + "SCNN": "scnn", + "DBNet": "dbnet", + "UNetMobv2": "unetmobv2", + "Cloud-Adapter": "cloud-adapter", + } + + self.load_weights(other_model_weight_path) + + self.create_ui() + + def load_weights(self, checkpoint_path: str): + for model_name, model in self.other_models.items(): + weight_path = os.path.join(checkpoint_path, model_name+".bin") + weight_path = glob(weight_path)[0] + weight = torch.load(weight_path, map_location=self.device) + model.load_state_dict(weight) + model.eval() + print(f"Loaded {model_name} weights from {weight_path}") + + def html_legend(self, num_classes=2): + if num_classes == 2: + return """ +
+
+
+ Clear +
+
+
+ Cloud +
+
+ """ + return """ +
+
+
+ Clear Sky +
+
+
+ Thick Cloud +
+
+
+ Thin Cloud +
+
+
+ Cloud Shadow +
+
+""" + + def create_ui(self): + with gr.Row(): + # 左侧：输入图片和按钮 + with gr.Column(scale=1): # 左侧列 + in_image = gr.Image( + label='Input Image', + sources='upload', + elem_classes='input_image', + interactive=True, + type="pil", + ) + with gr.Row(): + # 增加一个下拉菜单 + model_choice = gr.Dropdown( + choices=[ + "DBNet", + "HRCloudNet", + "CDNetv2", + "UNetMobv2", + "CDNetv1", + "MCDNet", + "KappaMask", + "SCNN", + ], + value="DBNet", + label="Model", + elem_classes='model_type', + ) + run_button = gr.Button( + 'Run', + variant="primary", + ) + # 示例输入列表 + gr.Examples( + examples=self.example_inputs, + inputs=in_image, + label="Example Inputs" + ) + + # 右侧：输出图片 + with gr.Column(scale=1): # 右侧列 + with gr.Column(): + # 输出图片 + out_image = gr.Image( + label='Output Image', + elem_classes='output_image', + interactive=False + ) + # 图例 + legend = gr.HTML( + value=self.legend, + elem_classes="output_legend", + ) + + # 按钮点击逻辑：触发图像转换 + run_button.click( + self.inference, + inputs=[in_image, model_choice], + outputs=out_image, + ) + + @torch.no_grad() + def inference(self, image: Image.Image, model_choice: str) -> Image.Image: + return self.simple_model_forward(image, self.name_mapping[model_choice]) + + @torch.no_grad() + def simple_model_forward(self, image: Image.Image, model_choice: str) -> Image.Image: + """ + Simple Model Inference + """ + ori_size = image.size + image = image.resize((self.img_size, self.img_size), + resample=Image.Resampling.BILINEAR) + image = np.array(image) + image = (image - np.min(image)) / (np.max(image)-np.min(image)) + + image = torch.from_numpy(image).unsqueeze(0).to(self.device) + image = image.permute(0, 3, 1, 2).float() + + logits: torch.Tensor = self.other_models[model_choice].forward(image) + pred_mask = torch.argmax(logits, dim=1).squeeze( + 0).cpu().numpy().astype(np.uint8) + + del image + del logits + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + im = Image.fromarray(pred_mask).convert("P") + im.putpalette(self.palette) + return im.resize(ori_size, resample=Image.Resampling.BILINEAR) + + +def get_palette(dataset_name: str) -> List[int]: + if dataset_name in ["cloudsen12_high_l1c", "cloudsen12_high_l2a"]: + return [79, 253, 199, 77, 2, 115, 251, 255, 41, 221, 53, 223] + if dataset_name == "l8_biome": + return [79, 253, 199, 221, 53, 223, 251, 255, 41, 77, 2, 115] + if dataset_name in ["gf12ms_whu_gf1", "gf12ms_whu_gf2", "hrc_whu"]: + return [79, 253, 199, 77, 2, 115] + raise Exception("dataset_name not supported") + + +if __name__ == '__main__': + title = 'Cloud Segmentation for Remote Sensing Images' + custom_css = """ +h1 { + text-align: center; + font-size: 24px; + font-weight: bold; + margin-bottom: 20px; +} +""" + hrc_whu_examples = glob("example_inputs/hrc_whu/*") + gf1_examples = glob("example_inputs/gf1/*") + gf2_examples = glob("example_inputs/gf2/*") + l1c_examples = glob("example_inputs/l1c/*") + l2a_examples = glob("example_inputs/l2a/*") + l8_examples = glob("example_inputs/l8/*") + + device = "cuda:0" if torch.cuda.is_available() else "cpu" + with gr.Blocks(analytics_enabled=False, title=title,css=custom_css) as demo: + gr.Markdown(f'# {title}') + with gr.Tabs(): + with gr.TabItem('Google Earth'): + CloudAdapterGradio( + device=device, + example_inputs=hrc_whu_examples, + num_classes=2, + palette=get_palette("hrc_whu"), + other_model_weight_path="checkpoints/hrc_whu" + ) + with gr.TabItem('Gaofen-1'): + CloudAdapterGradio( + device=device, + example_inputs=gf1_examples, + num_classes=2, + palette=get_palette("gf12ms_whu_gf1"), + other_model_weight_path="checkpoints/gf12ms_whu_gf1" + ) + with gr.TabItem('Gaofen-2'): + CloudAdapterGradio( + device=device, + example_inputs=gf2_examples, + num_classes=2, + palette=get_palette("gf12ms_whu_gf2"), + other_model_weight_path="checkpoints/gf12ms_whu_gf2" + ) + + with gr.TabItem('Sentinel-2 (L1C)'): + CloudAdapterGradio( + device=device, + example_inputs=l1c_examples, + num_classes=4, + palette=get_palette("cloudsen12_high_l1c"), + other_model_weight_path="checkpoints/cloudsen12_high_l1c" + ) + with gr.TabItem('Sentinel-2 (L2A)'): + CloudAdapterGradio( + device=device, + example_inputs=l2a_examples, + num_classes=4, + palette=get_palette("cloudsen12_high_l2a"), + other_model_weight_path="checkpoints/cloudsen12_high_l2a" + ) + with gr.TabItem('Landsat-8'): + CloudAdapterGradio( + device=device, + example_inputs=l8_examples, + num_classes=4, + palette=get_palette("l8_biome"), + other_model_weight_path="checkpoints/l8_biome" + ) + + demo.launch(share=True, debug=True) \ No newline at end of file diff --git a/checkpoints/cloudsen12_high_l1c/cdnetv1.bin b/checkpoints/cloudsen12_high_l1c/cdnetv1.bin new file mode 100644 index 0000000000000000000000000000000000000000..daf034367894e11140b04340e973eda05f306b92 --- /dev/null +++ b/checkpoints/cloudsen12_high_l1c/cdnetv1.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b6d53b129196c94f1f058a24b5c25df4f1f7563c9acb29fd49758294d1f6cceb +size 186527237 diff --git a/checkpoints/cloudsen12_high_l1c/cdnetv2.bin b/checkpoints/cloudsen12_high_l1c/cdnetv2.bin new file mode 100644 index 0000000000000000000000000000000000000000..8fdd192d810592e942c2acb50c93871f8ba5a615 --- /dev/null +++ b/checkpoints/cloudsen12_high_l1c/cdnetv2.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6a9c38d7c9468b2f5b4ffc5dfe0b7d316003c6f8a1cebcc21f694d84dad97b2d +size 270986885 diff --git a/checkpoints/cloudsen12_high_l1c/dbnet.bin b/checkpoints/cloudsen12_high_l1c/dbnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..8d4a26fa281c3fb0c7270825184ab278058c3cf3 --- /dev/null +++ b/checkpoints/cloudsen12_high_l1c/dbnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cef06e9aa7ff8b667a1cede0d78a1a4a2ce08ce8aedba748ecf0980c4362f375 +size 381851304 diff --git a/checkpoints/cloudsen12_high_l1c/hrcloudnet.bin b/checkpoints/cloudsen12_high_l1c/hrcloudnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..a4189984e83e55bdc37a6d682b7e9fbfffc06767 --- /dev/null +++ b/checkpoints/cloudsen12_high_l1c/hrcloudnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3fb5499f03116b9fb2b5ea3599ddfcd7c09ef9841426cef10809baf81f0cfc9c +size 299109140 diff --git a/checkpoints/cloudsen12_high_l1c/kappamask.bin b/checkpoints/cloudsen12_high_l1c/kappamask.bin new file mode 100644 index 0000000000000000000000000000000000000000..24e47933726627d3dd1c0197913ecc3c898961f1 --- /dev/null +++ b/checkpoints/cloudsen12_high_l1c/kappamask.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ae76545a2ed8a2aaf7ab01f0fb2b9acec92892786df87840690eb8233c05b3d9 +size 124145209 diff --git a/checkpoints/cloudsen12_high_l1c/mcdnet.bin b/checkpoints/cloudsen12_high_l1c/mcdnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..d8b418e7b2bc615d7a2dfe33a559537925ef7d1c --- /dev/null +++ b/checkpoints/cloudsen12_high_l1c/mcdnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:036a018532dc2963141728dfd4312303d97d575dec8ef3c9cd39992cfc41e6ad +size 52544517 diff --git a/checkpoints/cloudsen12_high_l1c/scnn.bin b/checkpoints/cloudsen12_high_l1c/scnn.bin new file mode 100644 index 0000000000000000000000000000000000000000..bbad83e132c601c8ebe24945f7c649fc07f2a8c8 --- /dev/null +++ b/checkpoints/cloudsen12_high_l1c/scnn.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:72b9744cc640d6b4f01c7a1a3ec41cd6cbd23d390179df3330d7b282964cf3da +size 4711 diff --git a/checkpoints/cloudsen12_high_l1c/unetmobv2.bin b/checkpoints/cloudsen12_high_l1c/unetmobv2.bin new file mode 100644 index 0000000000000000000000000000000000000000..183d2654238a753336a25124757b8384a9dfe9e9 --- /dev/null +++ b/checkpoints/cloudsen12_high_l1c/unetmobv2.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b51e7cea6e595f397609b9c7c469a8c09134ff7a059be5c35648671d60431e4d +size 26781125 diff --git a/checkpoints/cloudsen12_high_l2a/cdnetv1.bin b/checkpoints/cloudsen12_high_l2a/cdnetv1.bin new file mode 100644 index 0000000000000000000000000000000000000000..cbf11ca839fc8b81b3b2cfb0a5a998d6d9e816be --- /dev/null +++ b/checkpoints/cloudsen12_high_l2a/cdnetv1.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:950e9000d57d91cd04b83879072dc53e889383186203f65d0eae627cfac51550 +size 186527237 diff --git a/checkpoints/cloudsen12_high_l2a/cdnetv2.bin b/checkpoints/cloudsen12_high_l2a/cdnetv2.bin new file mode 100644 index 0000000000000000000000000000000000000000..30edab282ad89ba9e23d12dd216ebde950e10ce7 --- /dev/null +++ b/checkpoints/cloudsen12_high_l2a/cdnetv2.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b10bf239641196801866b1946466f3dbfa04a895f9a54f126a528eae200e169c +size 270986885 diff --git a/checkpoints/cloudsen12_high_l2a/dbnet.bin b/checkpoints/cloudsen12_high_l2a/dbnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..026d8c2a614959c08b607973f57578e1f261f93a --- /dev/null +++ b/checkpoints/cloudsen12_high_l2a/dbnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:08bbb910c142844a7a143f9377cd4c2377f2d6cbac3846d7b065939ad217b06d +size 381851304 diff --git a/checkpoints/cloudsen12_high_l2a/hrcloudnet.bin b/checkpoints/cloudsen12_high_l2a/hrcloudnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..20ad319dc43c6dd4717ae028593e7b6241535c3c --- /dev/null +++ b/checkpoints/cloudsen12_high_l2a/hrcloudnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:86fe9ebb28eef69661f801b3c576c53c92d8a9c634d7d4e4178fb63a153d8e5e +size 299109140 diff --git a/checkpoints/cloudsen12_high_l2a/kappamask.bin b/checkpoints/cloudsen12_high_l2a/kappamask.bin new file mode 100644 index 0000000000000000000000000000000000000000..b1642a29fd31fe77a2434242afb32127be31d2fe --- /dev/null +++ b/checkpoints/cloudsen12_high_l2a/kappamask.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f329219f7b84c67e3515fcbfd69bfb1a5b29ef893cbfe80d7a3e96a9c4dd11b7 +size 124145209 diff --git a/checkpoints/cloudsen12_high_l2a/mcdnet.bin b/checkpoints/cloudsen12_high_l2a/mcdnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..8f19637a103fd3400fc141ec3ee3e774201d6eef --- /dev/null +++ b/checkpoints/cloudsen12_high_l2a/mcdnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:81334d65167b113cdcbe5d9f009c1f922556f5793a8d117b435fabb1ae097282 +size 52544517 diff --git a/checkpoints/cloudsen12_high_l2a/scnn.bin b/checkpoints/cloudsen12_high_l2a/scnn.bin new file mode 100644 index 0000000000000000000000000000000000000000..fafe9a465f6209a590fe6ddee79ffa7f4ad772df --- /dev/null +++ b/checkpoints/cloudsen12_high_l2a/scnn.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a498440a79b3a5a1dbe9b2017123dac1658ec3ca2edcc000cf469010dd7f930d +size 4711 diff --git a/checkpoints/cloudsen12_high_l2a/unetmobv2.bin b/checkpoints/cloudsen12_high_l2a/unetmobv2.bin new file mode 100644 index 0000000000000000000000000000000000000000..8585b47fcc09850e2cfcc7051b06ca19811fb883 --- /dev/null +++ b/checkpoints/cloudsen12_high_l2a/unetmobv2.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0751cd214121123687858defc4800cb2a3ae18944a8b07a08fc7b66b78192437 +size 26781125 diff --git a/checkpoints/gf12ms_whu_gf1/cdnetv1.bin b/checkpoints/gf12ms_whu_gf1/cdnetv1.bin new file mode 100644 index 0000000000000000000000000000000000000000..95a300e8209284d1c18e2532ed194320caa256dd --- /dev/null +++ b/checkpoints/gf12ms_whu_gf1/cdnetv1.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:48f0f7dea20d20df222644c5a0633579afd2f53de39fb4afaab88aa621e019c5 +size 186375173 diff --git a/checkpoints/gf12ms_whu_gf1/cdnetv2.bin b/checkpoints/gf12ms_whu_gf1/cdnetv2.bin new file mode 100644 index 0000000000000000000000000000000000000000..fd4972ef297764bbc443ba445aad9011f3e9c823 --- /dev/null +++ b/checkpoints/gf12ms_whu_gf1/cdnetv2.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:16b2c501c9647ccec87565391946fc95beba4459992ed6a598c8650e9a85128e +size 270982789 diff --git a/checkpoints/gf12ms_whu_gf1/dbnet.bin b/checkpoints/gf12ms_whu_gf1/dbnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..f14935d686f1ffb4b94eb9bea908d671fd4bfc99 --- /dev/null +++ b/checkpoints/gf12ms_whu_gf1/dbnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:87db726ffb9bc6d69822c0b5e8167391c50bdfbc7c4a023b2700bbf3808960b8 +size 381851176 diff --git a/checkpoints/gf12ms_whu_gf1/hrcloudnet.bin b/checkpoints/gf12ms_whu_gf1/hrcloudnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..492a393a0d4f309f7c687edcb192367ed8ebb731 --- /dev/null +++ b/checkpoints/gf12ms_whu_gf1/hrcloudnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:69c3ba3f272c55e667d9867d50d8cf43cc7c26aac8cf1b365f4d7b62d35b9d9d +size 299099284 diff --git a/checkpoints/gf12ms_whu_gf1/kappamask.bin b/checkpoints/gf12ms_whu_gf1/kappamask.bin new file mode 100644 index 0000000000000000000000000000000000000000..7d26078d1b31aa1c9f7c09c39efe1b23dc1b383c --- /dev/null +++ b/checkpoints/gf12ms_whu_gf1/kappamask.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:460f292b64632da5591e829f7b208ae68e36e79682230bebf20415b7def126e7 +size 124145209 diff --git a/checkpoints/gf12ms_whu_gf1/mcdnet.bin b/checkpoints/gf12ms_whu_gf1/mcdnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..61b44908cb86ec94b17b5c6c7c528f424dfce9d7 --- /dev/null +++ b/checkpoints/gf12ms_whu_gf1/mcdnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a2e81bb99369297eb73a00aacfe825f851967e559fc1bd01365326c410a39d07 +size 52542213 diff --git a/checkpoints/gf12ms_whu_gf1/scnn.bin b/checkpoints/gf12ms_whu_gf1/scnn.bin new file mode 100644 index 0000000000000000000000000000000000000000..bc2f53ef4530189d75fb5debf0ba6e23e8130a08 --- /dev/null +++ b/checkpoints/gf12ms_whu_gf1/scnn.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e21a3c86223760d50b43cf0a3f9c808ab263a5b1b074ce7b7249470c3da21335 +size 3751 diff --git a/checkpoints/gf12ms_whu_gf1/unetmobv2.bin b/checkpoints/gf12ms_whu_gf1/unetmobv2.bin new file mode 100644 index 0000000000000000000000000000000000000000..83ea3fcb4111d939452d93076bbd1982e266d1d3 --- /dev/null +++ b/checkpoints/gf12ms_whu_gf1/unetmobv2.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a55cfaf62350ce09ada9925b86529a0361a5143a610360c1682b2adab42543c5 +size 26779973 diff --git a/checkpoints/gf12ms_whu_gf2/cdnetv1.bin b/checkpoints/gf12ms_whu_gf2/cdnetv1.bin new file mode 100644 index 0000000000000000000000000000000000000000..fb421399db0b9243083702af8f30425c119ac53d --- /dev/null +++ b/checkpoints/gf12ms_whu_gf2/cdnetv1.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:97ff37beaabd1a85ca52682a7a73c1660fb285c4f6bbe6c6de42e9a5ace7e5cd +size 186375173 diff --git a/checkpoints/gf12ms_whu_gf2/cdnetv2.bin b/checkpoints/gf12ms_whu_gf2/cdnetv2.bin new file mode 100644 index 0000000000000000000000000000000000000000..48b22ff6032e751ed105c5baa26cb6713db1ced5 --- /dev/null +++ b/checkpoints/gf12ms_whu_gf2/cdnetv2.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:710c1d7a4e7ecf011aed9dfb659c1acc09388b68b0264389cd0fae72840ce9c3 +size 270982789 diff --git a/checkpoints/gf12ms_whu_gf2/dbnet.bin b/checkpoints/gf12ms_whu_gf2/dbnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..bf6978f1a1c6bb98a5d179e5804405abc3363777 --- /dev/null +++ b/checkpoints/gf12ms_whu_gf2/dbnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8a61516dc057fa5100f5222570654067406c8f8dee1fae3e3a41a99d13cec1fe +size 381851176 diff --git a/checkpoints/gf12ms_whu_gf2/hrcloudnet.bin b/checkpoints/gf12ms_whu_gf2/hrcloudnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..88d94121b126fdea2eb7be398cb0d50693bb97fe --- /dev/null +++ b/checkpoints/gf12ms_whu_gf2/hrcloudnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ecfc004b7af6d105d4e73d074de0c9e3c5629d59347a70aaa1f36938bc3b7377 +size 299099284 diff --git a/checkpoints/gf12ms_whu_gf2/kappamask.bin b/checkpoints/gf12ms_whu_gf2/kappamask.bin new file mode 100644 index 0000000000000000000000000000000000000000..38da8a007d84a6e7cb82bf8dccc602f00a349b21 --- /dev/null +++ b/checkpoints/gf12ms_whu_gf2/kappamask.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c8e5194faafe0a43fd2348924c0f0217e322801ad41f0a34506f26ed87c028d8 +size 124145209 diff --git a/checkpoints/gf12ms_whu_gf2/mcdnet.bin b/checkpoints/gf12ms_whu_gf2/mcdnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..e7a20b5cae73fda3a5b6e22b2e8eb30957faae74 --- /dev/null +++ b/checkpoints/gf12ms_whu_gf2/mcdnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f8427fb5f8ac496c681404ff49f101527776b34950963847f4233ad732e3266d +size 52542213 diff --git a/checkpoints/gf12ms_whu_gf2/scnn.bin b/checkpoints/gf12ms_whu_gf2/scnn.bin new file mode 100644 index 0000000000000000000000000000000000000000..94607b1e58a4728bdc8381992651d09140d213b2 --- /dev/null +++ b/checkpoints/gf12ms_whu_gf2/scnn.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2a979b772223382952c33f2456a259d20f3f89ceb8f6896c1c4de0232c2e9158 +size 3751 diff --git a/checkpoints/gf12ms_whu_gf2/unetmobv2.bin b/checkpoints/gf12ms_whu_gf2/unetmobv2.bin new file mode 100644 index 0000000000000000000000000000000000000000..897a77d4a9ccfa64e5d08dcddc05528177e7ce17 --- /dev/null +++ b/checkpoints/gf12ms_whu_gf2/unetmobv2.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a24fc14f9dc984af1353b8271d6f2bcfc56b9dae662dca5de550496bde761a55 +size 26779973 diff --git a/checkpoints/hrc_whu/cdnetv1.bin b/checkpoints/hrc_whu/cdnetv1.bin new file mode 100644 index 0000000000000000000000000000000000000000..d4fdf6394743d01a2ba2dfa13f241631ae9a5b28 --- /dev/null +++ b/checkpoints/hrc_whu/cdnetv1.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3109f560dd139e1f363872c723eba1c870ece000faeb18a76ad5988c07f5a583 +size 186375173 diff --git a/checkpoints/hrc_whu/cdnetv2.bin b/checkpoints/hrc_whu/cdnetv2.bin new file mode 100644 index 0000000000000000000000000000000000000000..9a143ff91237b26cb819a07d5ecd8870962bd843 --- /dev/null +++ b/checkpoints/hrc_whu/cdnetv2.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:026672dd7d45cac9c9f5ad6642cdcf8ce136b76c224cca4ecceca8c1af80aa5e +size 270982789 diff --git a/checkpoints/hrc_whu/dbnet.bin b/checkpoints/hrc_whu/dbnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..e2bb0210b829dec710120e3462405f80b7497a98 --- /dev/null +++ b/checkpoints/hrc_whu/dbnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f3153c1a03c0a1cba012afde20dad7500ef3cd230247320b092b5fce22d44292 +size 381851176 diff --git a/checkpoints/hrc_whu/hrcloudnet.bin b/checkpoints/hrc_whu/hrcloudnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..f6e2596be9c82ab2b0f008008ebf3cd463855028 --- /dev/null +++ b/checkpoints/hrc_whu/hrcloudnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:35616fcd0ec286a1b3eafefc62938d2abfbdd2b4df358094e597dc87839e6cc2 +size 299099284 diff --git a/checkpoints/hrc_whu/kappamask.bin b/checkpoints/hrc_whu/kappamask.bin new file mode 100644 index 0000000000000000000000000000000000000000..12048e2d2bf368cb938a14992c6767b4e097a18c --- /dev/null +++ b/checkpoints/hrc_whu/kappamask.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8ede3aedec3c8ffaac92762e21e01a81f63637bdf3250d331cb7ca2abaed7427 +size 124145209 diff --git a/checkpoints/hrc_whu/mcdnet.bin b/checkpoints/hrc_whu/mcdnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..56e0caae21650d301e890eaf6312ca6298267a17 --- /dev/null +++ b/checkpoints/hrc_whu/mcdnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:37fee9b47ffd935663f2f083aa1262b411a75f9baed6762a9471c5231418bd20 +size 52542213 diff --git a/checkpoints/hrc_whu/scnn.bin b/checkpoints/hrc_whu/scnn.bin new file mode 100644 index 0000000000000000000000000000000000000000..7f12ce9e5514be27205e14ee26d9b541cd3732be --- /dev/null +++ b/checkpoints/hrc_whu/scnn.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8c44d5fa44e25725a35cc8310b17b75cafc3bf0df4820dffe13c7565fe6b3fe3 +size 3751 diff --git a/checkpoints/hrc_whu/unetmobv2.bin b/checkpoints/hrc_whu/unetmobv2.bin new file mode 100644 index 0000000000000000000000000000000000000000..6c892073c36129e62dd38ccee5bc93ccb2106cce --- /dev/null +++ b/checkpoints/hrc_whu/unetmobv2.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0c43160542fdf8a06cab3b361bb590c6f610689729cb92a90d96cdd537f429be +size 26779973 diff --git a/checkpoints/l8_biome/cdnetv1.bin b/checkpoints/l8_biome/cdnetv1.bin new file mode 100644 index 0000000000000000000000000000000000000000..855367acc786a5e9936f1e6ffcf0cc1eed94dd4d --- /dev/null +++ b/checkpoints/l8_biome/cdnetv1.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1ad71461000af0742f3255ccaa2741e21f635525ac36ab8ca5b50cc46efd27f0 +size 186527707 diff --git a/checkpoints/l8_biome/cdnetv2.bin b/checkpoints/l8_biome/cdnetv2.bin new file mode 100644 index 0000000000000000000000000000000000000000..16258d9721d15176d1ac22629ce57f0d6f8439b9 --- /dev/null +++ b/checkpoints/l8_biome/cdnetv2.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:270d5879bf8323d4c2f57e5cd84e6fcbce5adafac70521e5c7fb58a8a4cc91bb +size 270987497 diff --git a/checkpoints/l8_biome/dbnet.bin b/checkpoints/l8_biome/dbnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..4c7e1ef88b5fe90bdc95bcbb016694d20f07e1a3 --- /dev/null +++ b/checkpoints/l8_biome/dbnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f8e89b13b1b111c93f19ceaf3281f1a1d30114bad86b1f4c5ba38ead5185a164 +size 381852469 diff --git a/checkpoints/l8_biome/hrcloudnet.bin b/checkpoints/l8_biome/hrcloudnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..06aa5549eddb0f245a2db9a2ce2e39bba3da3657 --- /dev/null +++ b/checkpoints/l8_biome/hrcloudnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:01aaad5c90314b6a7cca04309cc580e0957c35b29a78854a51ab0d14e28547fd +size 299109140 diff --git a/checkpoints/l8_biome/kappamask.bin b/checkpoints/l8_biome/kappamask.bin new file mode 100644 index 0000000000000000000000000000000000000000..b5ab140ea80b9618348c1b37c59a12cc0f3a3541 --- /dev/null +++ b/checkpoints/l8_biome/kappamask.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:07bc7288105ed0fccfc6ec54a86f016735b322f2b66761c36461e12f9c174f25 +size 124145259 diff --git a/checkpoints/l8_biome/mcdnet.bin b/checkpoints/l8_biome/mcdnet.bin new file mode 100644 index 0000000000000000000000000000000000000000..0dbb27c5199ec26954a5f8930113f46f398c33c3 --- /dev/null +++ b/checkpoints/l8_biome/mcdnet.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:47d9d9d4cb38976677a73e6b587dd96b35bb219ef300cac9bdb5433b385c4c39 +size 52544793 diff --git a/checkpoints/l8_biome/scnn.bin b/checkpoints/l8_biome/scnn.bin new file mode 100644 index 0000000000000000000000000000000000000000..60d38aa24a19635fb08403838dbc83f976f7d14b --- /dev/null +++ b/checkpoints/l8_biome/scnn.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:63dc57a34334c2ea9245c88f29753b1c4749639c950c32e9cb2b2823ebed3db1 +size 4719 diff --git a/checkpoints/l8_biome/unetmobv2.bin b/checkpoints/l8_biome/unetmobv2.bin new file mode 100644 index 0000000000000000000000000000000000000000..bdcfe5c572a32315c562c5db8139ee39d19ecefc --- /dev/null +++ b/checkpoints/l8_biome/unetmobv2.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f1ff694e2c97b7696e9415f9d56e0f9b9ca813ac3e1d7f773cafae38c16f9511 +size 26781501 diff --git a/example_inputs/gf1/11.png b/example_inputs/gf1/11.png new file mode 100644 index 0000000000000000000000000000000000000000..f140af36561e196074a0bf4adf84f95de8231d5d Binary files /dev/null and b/example_inputs/gf1/11.png differ diff --git a/example_inputs/gf1/48.png b/example_inputs/gf1/48.png new file mode 100644 index 0000000000000000000000000000000000000000..1a2ef6daa29f194bd961a07b4ab0709d2ee3ce47 Binary files /dev/null and b/example_inputs/gf1/48.png differ diff --git a/example_inputs/gf1/9.png b/example_inputs/gf1/9.png new file mode 100644 index 0000000000000000000000000000000000000000..a206feaeb4af1132deeacc565ab0a78b6ea39d23 Binary files /dev/null and b/example_inputs/gf1/9.png differ diff --git a/example_inputs/gf2/160.png b/example_inputs/gf2/160.png new file mode 100644 index 0000000000000000000000000000000000000000..3d5df50c6e3ab430549c781fd2bccca9b41dec23 Binary files /dev/null and b/example_inputs/gf2/160.png differ diff --git a/example_inputs/gf2/2.png b/example_inputs/gf2/2.png new file mode 100644 index 0000000000000000000000000000000000000000..b25bcb2b91ef66195e7c71a479b11237d80c6b5d Binary files /dev/null and b/example_inputs/gf2/2.png differ diff --git a/example_inputs/gf2/63.png b/example_inputs/gf2/63.png new file mode 100644 index 0000000000000000000000000000000000000000..ad831a8f2ec6af36edb7790d6a8f8c15126af61e Binary files /dev/null and b/example_inputs/gf2/63.png differ diff --git a/example_inputs/hrc_whu/barren_7.png b/example_inputs/hrc_whu/barren_7.png new file mode 100644 index 0000000000000000000000000000000000000000..d0fb526613808c2b0cc68a3b54c7fe7bf58fc8e9 Binary files /dev/null and b/example_inputs/hrc_whu/barren_7.png differ diff --git a/example_inputs/hrc_whu/snow_10.png b/example_inputs/hrc_whu/snow_10.png new file mode 100644 index 0000000000000000000000000000000000000000..af123dbc6ccfb9b135f8d94bd9750adbf76aec11 Binary files /dev/null and b/example_inputs/hrc_whu/snow_10.png differ diff --git a/example_inputs/hrc_whu/vegetation_21.png b/example_inputs/hrc_whu/vegetation_21.png new file mode 100644 index 0000000000000000000000000000000000000000..6ea92b4d80bde86c01d93d43d64cfad87d2e4a3c Binary files /dev/null and b/example_inputs/hrc_whu/vegetation_21.png differ diff --git a/example_inputs/l1c/1.png b/example_inputs/l1c/1.png new file mode 100644 index 0000000000000000000000000000000000000000..f36ecf60322837525f33ac5e4a54a051333b3b48 Binary files /dev/null and b/example_inputs/l1c/1.png differ diff --git a/example_inputs/l1c/27.png b/example_inputs/l1c/27.png new file mode 100644 index 0000000000000000000000000000000000000000..96a3bc291f8cc3c678ceea8b339807b8da9d2f31 Binary files /dev/null and b/example_inputs/l1c/27.png differ diff --git a/example_inputs/l1c/76.png b/example_inputs/l1c/76.png new file mode 100644 index 0000000000000000000000000000000000000000..c06811d59483a905ec5e179d0b2a0bb7fe6360fa Binary files /dev/null and b/example_inputs/l1c/76.png differ diff --git a/example_inputs/l2a/121.png b/example_inputs/l2a/121.png new file mode 100644 index 0000000000000000000000000000000000000000..7d9d4cbb45f6205c1684009133d0dd807c5dbcd3 Binary files /dev/null and b/example_inputs/l2a/121.png differ diff --git a/example_inputs/l2a/18.png b/example_inputs/l2a/18.png new file mode 100644 index 0000000000000000000000000000000000000000..903d2f11a8501566d2024c73ddfe39fe6d5cc08d Binary files /dev/null and b/example_inputs/l2a/18.png differ diff --git a/example_inputs/l2a/35.png b/example_inputs/l2a/35.png new file mode 100644 index 0000000000000000000000000000000000000000..9fa0521173620ef89f0d6aead16ec7d64348ebcf Binary files /dev/null and b/example_inputs/l2a/35.png differ diff --git a/example_inputs/l8/barren_LC81390292014135LGN00_patch_5632_512.png b/example_inputs/l8/barren_LC81390292014135LGN00_patch_5632_512.png new file mode 100644 index 0000000000000000000000000000000000000000..278dafd53a3fdf8bfc70313d14c4da4aded8d43d Binary files /dev/null and b/example_inputs/l8/barren_LC81390292014135LGN00_patch_5632_512.png differ diff --git a/example_inputs/l8/forest_LC80160502014041LGN00_patch_4608_4608.png b/example_inputs/l8/forest_LC80160502014041LGN00_patch_4608_4608.png new file mode 100644 index 0000000000000000000000000000000000000000..4ba447143ab2270d10e2357c564eb3219142cd72 Binary files /dev/null and b/example_inputs/l8/forest_LC80160502014041LGN00_patch_4608_4608.png differ diff --git a/example_inputs/l8/shrubland_LC81020802014100LGN00_patch_1024_3584.png b/example_inputs/l8/shrubland_LC81020802014100LGN00_patch_1024_3584.png new file mode 100644 index 0000000000000000000000000000000000000000..2d6043502426af800c994106fe27afc31885e77c Binary files /dev/null and b/example_inputs/l8/shrubland_LC81020802014100LGN00_patch_1024_3584.png differ diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/__pycache__/__init__.cpython-38.pyc b/models/__pycache__/__init__.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9acc30e067940f2769cafff711b329870a14a468 Binary files /dev/null and b/models/__pycache__/__init__.cpython-38.pyc differ diff --git a/models/__pycache__/cdnetv1.cpython-38.pyc b/models/__pycache__/cdnetv1.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e259043b9ead0070f625c08d7b8516ca09a57cc Binary files /dev/null and b/models/__pycache__/cdnetv1.cpython-38.pyc differ diff --git a/models/__pycache__/cdnetv2.cpython-38.pyc b/models/__pycache__/cdnetv2.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..00629d020798a3f26af0c84be5f380b0ff0af70c Binary files /dev/null and b/models/__pycache__/cdnetv2.cpython-38.pyc differ diff --git a/models/__pycache__/dbnet.cpython-38.pyc b/models/__pycache__/dbnet.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b01bfcab906b8ee0f9a56f618ce1f76e9026df3f Binary files /dev/null and b/models/__pycache__/dbnet.cpython-38.pyc differ diff --git a/models/__pycache__/hrcloudnet.cpython-38.pyc b/models/__pycache__/hrcloudnet.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84a0ab115ec545bbb126a7cde0d942c648ce28d4 Binary files /dev/null and b/models/__pycache__/hrcloudnet.cpython-38.pyc differ diff --git a/models/__pycache__/kappamask.cpython-38.pyc b/models/__pycache__/kappamask.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9fbd1b23003bc400d7464b47694fd7103b0d1dbe Binary files /dev/null and b/models/__pycache__/kappamask.cpython-38.pyc differ diff --git a/models/__pycache__/mcdnet.cpython-38.pyc b/models/__pycache__/mcdnet.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..117c7164b93302c43c59050ca74775e0c79c70fd Binary files /dev/null and b/models/__pycache__/mcdnet.cpython-38.pyc differ diff --git a/models/__pycache__/scnn.cpython-38.pyc b/models/__pycache__/scnn.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7d023706e056073502a48d81456fc98f7ff82207 Binary files /dev/null and b/models/__pycache__/scnn.cpython-38.pyc differ diff --git a/models/__pycache__/unetmobv2.cpython-38.pyc b/models/__pycache__/unetmobv2.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da710c90b1397033a86f31ac3cba96384cf84bc9 Binary files /dev/null and b/models/__pycache__/unetmobv2.cpython-38.pyc differ diff --git a/models/cdnetv1.py b/models/cdnetv1.py new file mode 100644 index 0000000000000000000000000000000000000000..081151455cbde1a92b028e7ada73a25cb88e0f09 --- /dev/null +++ b/models/cdnetv1.py @@ -0,0 +1,389 @@ +# -*- 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.nn.functional as F + +affine_par = True + + +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, in_channels=3,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(in_channels, 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]) \ No newline at end of file diff --git a/models/cdnetv2.py b/models/cdnetv2.py new file mode 100644 index 0000000000000000000000000000000000000000..e6fbdeecff5f3cd7d9d1fda9613909d4d85b5cc5 --- /dev/null +++ b/models/cdnetv2.py @@ -0,0 +1,693 @@ +# -*- 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 +# import torch.nn as nn +import torch.nn.functional as F +from torch import nn + +affine_par = True + + +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, in_channels=3,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(in_channels, 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 + 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]) \ No newline at end of file diff --git a/models/dbnet.py b/models/dbnet.py new file mode 100644 index 0000000000000000000000000000000000000000..b1dfa46e11f5befabfc4a0731fe8272bb0d5eb56 --- /dev/null +++ b/models/dbnet.py @@ -0,0 +1,680 @@ +# -*- coding: utf-8 -*- +# @Time : 2024/7/26 上午11:19 +# @Author : xiaoshun +# @Email : 3038523973@qq.com +# @File : dbnet.py +# @Software: PyCharm + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + + +# 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 DBNet(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 = DBNet(img_size=224, in_channels=3, num_classes=7).cuda() + y = model(x) + print(y.shape) + # torch.Size([1, 7, 224, 224]) \ No newline at end of file diff --git a/models/hrcloudnet.py b/models/hrcloudnet.py new file mode 100644 index 0000000000000000000000000000000000000000..fbb49a3cf808d16f3327ae8912beac955488682c --- /dev/null +++ b/models/hrcloudnet.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, in_channels=3,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(in_channels, 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, num_classes, 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]) \ No newline at end of file diff --git a/models/kappamask.py b/models/kappamask.py new file mode 100644 index 0000000000000000000000000000000000000000..57072c589abd757c0e53b4237e7531da4b3f37c3 --- /dev/null +++ b/models/kappamask.py @@ -0,0 +1,152 @@ +# -*- coding: utf-8 -*- +# @Time : 2024/8/7 下午3:51 +# @Author : xiaoshun +# @Email : 3038523973@qq.com +# @File : kappamask.py.py +# @Software: PyCharm + +import torch +from torch import nn as nn +from torch.nn import functional as F + + +class KappaMask(nn.Module): + def __init__(self, num_classes=2, in_channels=3): + super().__init__() + self.conv1 = nn.Sequential( + nn.Conv2d(in_channels, 64, 3, 1, 1), + nn.ReLU(inplace=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.ReLU(inplace=True), + ) + self.conv2 = nn.Sequential( + nn.Conv2d(64, 128, 3, 1, 1), + nn.ReLU(inplace=True), + nn.Conv2d(128, 128, 3, 1, 1), + nn.ReLU(inplace=True), + ) + self.conv3 = nn.Sequential( + nn.Conv2d(128, 256, 3, 1, 1), + nn.ReLU(inplace=True), + nn.Conv2d(256, 256, 3, 1, 1), + nn.ReLU(inplace=True), + ) + + self.conv4 = nn.Sequential( + nn.Conv2d(256, 512, 3, 1, 1), + nn.ReLU(inplace=True), + nn.Conv2d(512, 512, 3, 1, 1), + nn.ReLU(inplace=True), + ) + self.drop4 = nn.Dropout(0.5) + + self.conv5 = nn.Sequential( + nn.Conv2d(512, 1024, 3, 1, 1), + nn.ReLU(inplace=True), + nn.Conv2d(1024, 1024, 3, 1, 1), + nn.ReLU(inplace=True), + ) + self.drop5 = nn.Dropout(0.5) + + self.up6 = nn.Sequential( + nn.Upsample(scale_factor=2), + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(1024, 512, 2), + nn.ReLU(inplace=True) + ) + self.conv6 = nn.Sequential( + nn.Conv2d(1024, 512, 3, 1, 1), + nn.ReLU(inplace=True), + nn.Conv2d(512, 512, 3, 1, 1), + nn.ReLU(inplace=True), + ) + self.up7 = nn.Sequential( + nn.Upsample(scale_factor=2), + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(512, 256, 2), + nn.ReLU(inplace=True) + ) + self.conv7 = nn.Sequential( + nn.Conv2d(512, 256, 3, 1, 1), + nn.ReLU(inplace=True), + nn.Conv2d(256, 256, 3, 1, 1), + nn.ReLU(inplace=True), + ) + + self.up8 = nn.Sequential( + nn.Upsample(scale_factor=2), + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(256, 128, 2), + nn.ReLU(inplace=True) + ) + self.conv8 = nn.Sequential( + nn.Conv2d(256, 128, 3, 1, 1), + nn.ReLU(inplace=True), + nn.Conv2d(128, 128, 3, 1, 1), + nn.ReLU(inplace=True), + ) + + self.up9 = nn.Sequential( + nn.Upsample(scale_factor=2), + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(128, 64, 2), + nn.ReLU(inplace=True) + ) + self.conv9 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.ReLU(inplace=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.ReLU(inplace=True), + nn.Conv2d(64, 2, 3, 1, 1), + nn.ReLU(inplace=True), + ) + self.conv10 = nn.Conv2d(2, num_classes, 1) + self.__init_weights() + + def __init_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + + def forward(self, x): + conv1 = self.conv1(x) + pool1 = F.max_pool2d(conv1, 2, 2) + + conv2 = self.conv2(pool1) + pool2 = F.max_pool2d(conv2, 2, 2) + + conv3 = self.conv3(pool2) + pool3 = F.max_pool2d(conv3, 2, 2) + + conv4 = self.conv4(pool3) + drop4 = self.drop4(conv4) + pool4 = F.max_pool2d(drop4, 2, 2) + + conv5 = self.conv5(pool4) + drop5 = self.drop5(conv5) + + up6 = self.up6(drop5) + merge6 = torch.cat((drop4, up6), dim=1) + conv6 = self.conv6(merge6) + + up7 = self.up7(conv6) + merge7 = torch.cat((conv3, up7), dim=1) + conv7 = self.conv7(merge7) + + up8 = self.up8(conv7) + merge8 = torch.cat((conv2, up8), dim=1) + conv8 = self.conv8(merge8) + + up9 = self.up9(conv8) + merge9 = torch.cat((conv1, up9), dim=1) + conv9 = self.conv9(merge9) + + output = self.conv10(conv9) + return output + + +if __name__ == '__main__': + model = KappaMask(num_classes=2, in_channels=3) + fake_data = torch.rand(2, 3, 256, 256) + output = model(fake_data) + print(output.shape) \ No newline at end of file diff --git a/models/mcdnet.py b/models/mcdnet.py new file mode 100644 index 0000000000000000000000000000000000000000..93a1e0f8f944377e917da793ca7e73a5799ba704 --- /dev/null +++ b/models/mcdnet.py @@ -0,0 +1,435 @@ +# -*- coding: utf-8 -*- +# @Time : 2024/7/21 下午3:51 +# @Author : xiaoshun +# @Email : 3038523973@qq.com +# @File : mcdnet.py +# @Software: PyCharm +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().__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) + + @torch.no_grad() + def get_lr_data(self, x: torch.Tensor) -> torch.Tensor: + images = x.cpu().permute(0, 2, 3, 1).numpy() # b, h, w, c + batch_size = images.shape[0] + lr = [] + for i in range(batch_size): + lr_image = image_dehazer.remove_haze((images[i]*255).astype(np.uint8), showHazeTransmissionMap=False)[0] # h, w, c, numpy.array + lr_tensor = torch.from_numpy(lr_image).permute(2, 0, 1)/255. # c, h, w + lr.append(lr_tensor) + return torch.stack(lr, dim=0).to(x.device) # b, c, h, w + + 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 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 + + +if __name__ == "__main__": + num_classes = 2 + model = MCDNet() + # inp = torch.randn(size=(2, 3, 256, 256)) + # assert model(input).shape == (2, 2, 256, 256) \ No newline at end of file diff --git a/models/scnn.py b/models/scnn.py new file mode 100644 index 0000000000000000000000000000000000000000..b5bd5d6583196717241440e434415060d1ba8e12 --- /dev/null +++ b/models/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 SCNN(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 = SCNN(num_classes=7) + fake_img = torch.randn((2, 3, 224, 224)) + out = model(fake_img) + print(out.shape) + # torch.Size([2, 7, 224, 224]) \ No newline at end of file diff --git a/models/unetmobv2.py b/models/unetmobv2.py new file mode 100644 index 0000000000000000000000000000000000000000..3ae78246b4c7ceed76e2eb2c7c013e5300666fe0 --- /dev/null +++ b/models/unetmobv2.py @@ -0,0 +1,31 @@ +# -*- coding: utf-8 -*- +# @Time : 2024/8/6 下午3:44 +# @Author : xiaoshun +# @Email : 3038523973@qq.com +# @File : unetmobv2.py +# @Software: PyCharm +import segmentation_models_pytorch as smp +import torch +from torch import nn as nn + + +class UNetMobV2(nn.Module): + def __init__(self,num_classes,in_channels=3): + super().__init__() + self.backbone = smp.Unet( + encoder_name='mobilenet_v2', + encoder_weights=None, + in_channels=in_channels, + classes=num_classes, + ) + + def forward(self, x): + x = self.backbone(x) + return x + + +if __name__ == '__main__': + fake_image = torch.rand(1, 3, 224, 224) + model = UNetMobV2(num_classes=2) + output = model(fake_image) + print(output.size()) \ No newline at end of file