调整神经网络输入为4通道

models/ins_detect/mask_rcnn.py

@@ -16,7 +16,7 @@ from .faster_rcnn import _default_anchorgen, FasterRCNN, FastRCNNConvFCHead, RPN
 
 
 __all__ = [
-    "MaskRCNN",
+    "InsDetectNet",
     "MaskRCNN_ResNet50_FPN_Weights",
     "MaskRCNN_ResNet50_FPN_V2_Weights",
     "maskrcnn_resnet50_fpn",
@@ -24,7 +24,7 @@ __all__ = [
 ]
 
 
-class MaskRCNN(FasterRCNN):
+class InsDetectNet(FasterRCNN):
     """
     Implements Mask R-CNN.
 
@@ -413,7 +413,7 @@ def maskrcnn_resnet50_fpn(
     weights_backbone: Optional[ResNet50_Weights] = ResNet50_Weights.IMAGENET1K_V1,
     trainable_backbone_layers: Optional[int] = None,
     **kwargs: Any,
-) -> MaskRCNN:
+) -> InsDetectNet:
     """Mask R-CNN model with a ResNet-50-FPN backbone from the `Mask R-CNN
     <https://arxiv.org/abs/1703.06870>`_ paper.
 
@@ -498,7 +498,7 @@ def maskrcnn_resnet50_fpn(
 
     backbone = resnet50(weights=weights_backbone, progress=progress, norm_layer=norm_layer)
     backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers)
-    model = MaskRCNN(backbone, num_classes=num_classes, **kwargs)
+    model = InsDetectNet(backbone, num_classes=num_classes, **kwargs)
 
     if weights is not None:
         model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True))
@@ -521,7 +521,7 @@ def maskrcnn_resnet50_fpn_v2(
     weights_backbone: Optional[ResNet50_Weights] = None,
     trainable_backbone_layers: Optional[int] = None,
     **kwargs: Any,
-) -> MaskRCNN:
+) -> InsDetectNet:
     """Improved Mask R-CNN model with a ResNet-50-FPN backbone from the `Benchmarking Detection Transfer
     Learning with Vision Transformers <https://arxiv.org/abs/2111.11429>`_ paper.
 
@@ -571,7 +571,7 @@ def maskrcnn_resnet50_fpn_v2(
         (backbone.out_channels, 7, 7), [256, 256, 256, 256], [1024], norm_layer=nn.BatchNorm2d
     )
     mask_head = MaskRCNNHeads(backbone.out_channels, [256, 256, 256, 256], 1, norm_layer=nn.BatchNorm2d)
-    model = MaskRCNN(
+    model = InsDetectNet(
         backbone,
         num_classes=num_classes,
         rpn_anchor_generator=rpn_anchor_generator,

models/line_detect/dataset_LD.py

@@ -84,6 +84,9 @@ class WirePointDataset(BaseDataset):
         else:
             img = self.default_transform(img)
 
+
+        # new_channel = torch.zeros(1, 512, 512)
+        # img=torch.cat((img,new_channel),dim=0)
         # print(f'img:{img.shape}')
         return img, target
 

models/line_detect/line_net.py

@@ -133,6 +133,15 @@ class LineNet(BaseDetectionNet):
                 f"box_roi_pool should be of type MultiScaleRoIAlign or None instead of {type(box_roi_pool)}"
             )
 
+        # 修改第一个卷积层，将 in_channels 从 3 改为 4
+        backbone.body.conv1 = nn.Conv2d(
+            in_channels=4,
+            out_channels=64,
+            kernel_size=7,
+            stride=2,
+            padding=3,
+            bias=False
+        )
         if num_classes is not None:
             if box_predictor is not None:
                 raise ValueError("num_classes should be None when box_predictor is specified")
@@ -203,9 +212,12 @@ class LineNet(BaseDetectionNet):
         )
 
         if image_mean is None:
-            image_mean = [0.485, 0.456, 0.406]
+            # image_mean = [0.485, 0.456, 0.406]
+            image_mean = [0.485, 0.456, 0.406, 0.5]  # 假设你新加的通道均值为0.5
+
         if image_std is None:
-            image_std = [0.229, 0.224, 0.225]
+            # image_std = [0.229, 0.224, 0.225]
+            image_std = [0.229, 0.224, 0.225, 0.2]  # 标准差也补一个值
         transform = GeneralizedRCNNTransform(min_size, max_size, image_mean, image_std, **kwargs)
 
         super().__init__(backbone, rpn, roi_heads, transform)

models/line_detect/roi_heads.py

@@ -1007,7 +1007,7 @@ class RoIHeads(nn.Module):
         # else:
         #     self.training = False
         #     # print(f'targets is None')
-        print(f'roihead forward!!!')
+        # print(f'roihead forward!!!')
         if targets is not None:
             for t in targets:
                 # TODO: https://github.com/pytorch/pytorch/issues/26731

models/line_detect/train.yaml

@@ -1,6 +1,6 @@
 io:
   logdir: train_results
-  datadir: \\192.168.50.222/share/zyh/513train/a_dataset
+  datadir: \\192.168.50.222/share/zyh/5月彩色钢板数据汇总/zjh/a_dataset
 #  datadir: D:\python\PycharmProjects\data_20250223\0423_
 #  datadir: I:\datasets\wirenet_1000