Upload 13 files

model/net/super_retina.py

@@ -0,0 +1,47 @@
+# 这部分是解耦之后重构的SuperRetina模型
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from .super_retina_blocks.SharedEncoder import SharedEncoder
+from .super_retina_blocks.DetectorHead import DetectorHead
+from .super_retina_blocks.DescriptorHead import DescriptorHead
+
+
+class SuperRetina(nn.Module):
+    def __init__(self, config=None, encoder_device='cpu', detector_device='cpu', descriptor_device='cpu', n_class=1):
+        super().__init__()
+
+        self.channel_list = config["channel_list"]
+
+        self.encoder = SharedEncoder(self.channel_list)
+        self.desc_head = DescriptorHead(self.channel_list)
+        self.det_head = DetectorHead(self.channel_list, n_class)
+        self.encoder_device = encoder_device
+        self.detector_device = detector_device
+        self.descriptor_device = descriptor_device
+
+        # 移动到对应设备上
+        self.encoder.to(self.encoder_device)
+        self.det_head.to(self.detector_device)
+        self.desc_head.to(self.descriptor_device)
+
+        if config is not None:
+            self.config = config    # 传递config参数
+            self.nms_size = config['nms_size']
+            self.nms_thresh = config['nms_thresh']
+            self.scale = 8
+            # self.dice = DiceLoss()
+            # self.kernel = get_gaussian_kernel(kernlen=config['gaussian_kernel_size'], nsig=config['gaussian_sigma']).to(device)
+
+    def forward(self, x):   # 定义模型的前向传播过程
+        # 特征编码
+        enc_out, convs = self.encoder(x.to(self.encoder_device))
+
+        # 关键点检测
+        semi = self.det_head(enc_out.to(self.detector_device), [item.to(self.detector_device) for item in convs])
+
+        # 描述符生成
+        desc = self.desc_head(enc_out.to(self.descriptor_device))
+
+        return semi, desc

model/net/super_retina_blocks/DescriptorHead.py

@@ -0,0 +1,32 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DescriptorHead(nn.Module):
+    """描述符生成头"""
+    def __init__(self, channel_list):
+        super().__init__()
+
+        c1, c2, c3, c4, c5, d1, d2 = channel_list
+
+        self.relu = nn.ReLU(inplace=True)
+        self.convDa = torch.nn.Conv2d(c4, c5, kernel_size=3, stride=1, padding=1)   # 尺寸不变
+        # ---------------------------------------------------------------------------------------
+        # 原本的代码这里尺寸有点问题，96下采样过后是47不是48，需要padding=1
+        # self.convDb = torch.nn.Conv2d(c5, d1, kernel_size=4, stride=2, padding=0)   # 尺寸缩小1/2
+        self.convDb = torch.nn.Conv2d(c5, d1, kernel_size=4, stride=2, padding=1)   # 尺寸缩小1/2
+        # ---------------------------------------------------------------------------------------
+        self.convDc = torch.nn.Conv2d(d1, d2, kernel_size=1, stride=1, padding=0)   # 尺寸不变
+        self.trans_conv = nn.ConvTranspose2d(d1, d2, 2, stride=2)                   # 尺寸放大2倍
+
+    def forward(self, x):
+        cDa = self.relu(self.convDa(x))
+        cDb = self.relu(self.convDb(cDa))
+        desc = self.convDc(cDb)
+        
+        # 归一化描述符
+        dn = torch.norm(desc, p=2, dim=1)
+        desc = desc.div(torch.unsqueeze(dn, 1))
+        
+        return self.trans_conv(desc)

model/net/super_retina_blocks/DescriptorHead_9_0.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DescriptorHead(nn.Module):
+    """描述符生成头"""
+    def __init__(self, channel_list):
+        super().__init__()
+
+        c1, c2, c3, c4, c5, d1, d2 = channel_list
+
+        self.relu = nn.ReLU(inplace=True)
+        self.convDa = torch.nn.Conv2d(c4, c5, kernel_size=3, stride=1, padding=1)   # 尺寸不变
+        # ---------------------------------------------------------------------------------------
+        # 原本的代码这里尺寸有点问题，96下采样过后是47不是48，需要padding=1
+        # self.convDb = torch.nn.Conv2d(c5, d1, kernel_size=4, stride=2, padding=0)   # 尺寸缩小1/2
+        self.convDb = torch.nn.Conv2d(c5, d1, kernel_size=4, stride=2, padding=1)   # 尺寸缩小1/2
+        # ---------------------------------------------------------------------------------------
+        self.convDc = torch.nn.Conv2d(d1, d2, kernel_size=1, stride=1, padding=0)   # 尺寸不变
+        # self.trans_conv = nn.ConvTranspose2d(d1, d2, 2, stride=2)                   # 尺寸放大2倍
+
+    def forward(self, x):
+        cDa = self.relu(self.convDa(x))
+        cDb = self.relu(self.convDb(cDa))
+        desc = self.convDc(cDb)
+        
+        # 归一化描述符
+        dn = torch.norm(desc, p=2, dim=1)
+        desc = desc.div(torch.unsqueeze(dn, 1))
+        
+        # return self.trans_conv(desc)
+        return desc

model/net/super_retina_blocks/DetectorHead.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .DoubleConv import DoubleConv
+
+class DetectorHead(nn.Module):
+    """关键点检测头部"""
+    def __init__(self, channel_list, n_class):
+        """
+        in_channels_list: 各层输入通道数 [c3+c4, c2+c3, c1+c2]
+        """
+        super().__init__()
+
+        c1, c2, c3, c4, c5, d1, d2 = channel_list
+
+        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        # self.upsample = nn.Upsample(scale_factor=2, mode='bicubic', align_corners=True)
+
+        # 上采样路径
+        self.dconv_up3 = DoubleConv(c3 + c4, c3)
+        self.dconv_up2 = DoubleConv(c2 + c3, c2)
+        self.dconv_up1 = DoubleConv(c1 + c2, c1)
+        self.conv_last = nn.Conv2d(c1, n_class, kernel_size=1)
+    
+    def forward(self, x, convs):
+        """
+        x: 编码器输出特征
+        convs: 中间特征 [conv1, conv2, conv3]
+        """
+        # 上采样阶段1
+        cPa = self.upsample(x)
+        cPa = torch.cat([cPa, convs[2]], dim=1)  # 连接conv3
+        cPa = self.dconv_up3(cPa)
+        
+        # 上采样阶段2
+        cPa = self.upsample(cPa)
+        cPa = torch.cat([cPa, convs[1]], dim=1)  # 连接conv2
+        cPa = self.dconv_up2(cPa)
+        
+        # 上采样阶段3
+        cPa = self.upsample(cPa)
+        cPa = torch.cat([cPa, convs[0]], dim=1)  # 连接conv1
+        cPa = self.dconv_up1(cPa)
+        
+        # 最终输出
+        semi = self.conv_last(cPa)
+        return torch.sigmoid(semi)

model/net/super_retina_blocks/DoubleConv.py

@@ -0,0 +1,17 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class DoubleConv(nn.Module):
+    """双卷积模块"""
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, 3, padding=1),
+            nn.ReLU(inplace=True)
+        )
+    
+    def forward(self, x):
+        return self.conv(x)

model/net/super_retina_blocks/SharedEncoder.py

@@ -0,0 +1,41 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SharedEncoder(nn.Module):
+    """共享特征编码器"""
+    def __init__(self, channel_list):
+        super().__init__()
+
+        c1, c2, c3, c4, c5, d1, d2 = channel_list
+
+        self.relu = nn.ReLU(inplace=True)
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.conv1a = torch.nn.Conv2d(1, c1, kernel_size=3, stride=1, padding=1)
+        self.conv1b = torch.nn.Conv2d(c1, c1, kernel_size=3, stride=1, padding=1)
+        self.conv2a = torch.nn.Conv2d(c1, c2, kernel_size=3, stride=1, padding=1)
+        self.conv2b = torch.nn.Conv2d(c2, c2, kernel_size=3, stride=1, padding=1)
+        self.conv3a = torch.nn.Conv2d(c2, c3, kernel_size=3, stride=1, padding=1)
+        self.conv3b = torch.nn.Conv2d(c3, c3, kernel_size=3, stride=1, padding=1)
+        self.conv4a = torch.nn.Conv2d(c3, c4, kernel_size=3, stride=1, padding=1)
+        self.conv4b = torch.nn.Conv2d(c4, c4, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        # 第一层
+        x = self.relu(self.conv1a(x))
+        conv1 = self.relu(self.conv1b(x))
+        x = self.pool(conv1)
+        # 第二层
+        x = self.relu(self.conv2a(x))
+        conv2 = self.relu(self.conv2b(x))
+        x = self.pool(conv2)
+        # 第三层
+        x = self.relu(self.conv3a(x))
+        conv3 = self.relu(self.conv3b(x))
+        x = self.pool(conv3)
+        # 第四层
+        x = self.relu(self.conv4a(x))
+        x = self.relu(self.conv4b(x))
+        
+        return x, [conv1, conv2, conv3]