add high_reso_resnet18fpn

libs/vision_libs/models/detection/rpn.py

@@ -357,7 +357,9 @@ class RegionProposalNetwork(torch.nn.Module):
         """
         # RPN uses all feature maps that are available
         features = list(features.values())
+
         objectness, pred_bbox_deltas = self.head(features)
+
         anchors = self.anchor_generator(images, features)
 
         num_images = len(anchors)
@@ -385,4 +387,5 @@ class RegionProposalNetwork(torch.nn.Module):
                 "loss_objectness": loss_objectness,
                 "loss_rpn_box_reg": loss_rpn_box_reg,
             }
+        # print(f'boxes:{boxes[0].shape}')
         return boxes, losses

models/base/base_detection_net.py

@@ -114,6 +114,11 @@ class BaseDetectionNet(BaseModel):
             features = OrderedDict([("0", features)])
         proposals, proposal_losses = self.rpn(images, features, targets)
 
+
+        # print(f'proposals:{proposals[0].shape}')
+
+
+
         detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)
         detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)  # type: ignore[operator]
 

models/base/resnet50fpn.py → models/base/high_reso_resnet.py

@@ -88,8 +88,29 @@ class Bottleneck(nn.Module):
 # ResNet 主类
 # ----------------------------
 
+def resnet18fpn(out_channels=256):
+    backbone = ResNet(Bottleneck,[2,2,2])
+    return_layers = {
+        'encoder0': '0',
+        'encoder1': '1',
+        'encoder2': '2',
+        'encoder3': '3',
+        # 'encoder4': '5'
+    }
+
+    # in_channels_list = [self.inplanes, 64, 128, 256, 512]
+    # in_channels_list = [64, 256, 512, 1024, 2048]
+    in_channels_list = [64, 256, 512, 1024]
+
+    return BackboneWithFPN(
+        backbone,
+        return_layers=return_layers,
+        in_channels_list=in_channels_list,
+        out_channels=out_channels,
+    )
+
 def resnet50fpn(out_channels=256):
-    backbone = ResNet(Bottleneck)
+    backbone = ResNet(Bottleneck,[3,4,6])
     return_layers = {
         'encoder0': '0',
         'encoder1': '1',
@@ -111,7 +132,7 @@ def resnet50fpn(out_channels=256):
 
 
 class ResNet(nn.Module):
-    def __init__(self, block: Type[Union[Bottleneck]],):
+    def __init__(self, block: Type[Union[Bottleneck]], layers: List[int],):
         super(ResNet, self).__init__()
         self._norm_layer = nn.BatchNorm2d
         self.inplanes = 64
@@ -126,9 +147,9 @@ class ResNet(nn.Module):
             nn.ReLU(inplace=True),
             nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
         )
-        self.encoder1 = self._make_layer(block, 64, 3,stride=2)
-        self.encoder2 = self._make_layer(block, 128, 4, stride=2)
-        self.encoder3 = self._make_layer(block, 256, 6, stride=2)
+        self.encoder1 = self._make_layer(block, 64, layers[0],stride=2)
+        self.encoder2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.encoder3 = self._make_layer(block, 256, layers[2], stride=2)
         # self.encoder4 = self._make_layer(block, 512, 3, stride=2)
         # self.encoder5 = self._make_layer(block, 512, 3, stride=2)
         # self.body = nn.ModuleDict({
@@ -253,7 +274,7 @@ if __name__ == "__main__":
     input_tensor = torch.randn(1, 3, 512, 512).to(device)
     output_tensor = model(input_tensor)
 
-    backbone = ResNet(Bottleneck).to(device)
+    backbone = ResNet(Bottleneck,[3,4,6]).to(device)
     features = backbone(input_tensor)
     print("Raw backbone output:", list(features.keys()))
     print(f"Input shape: {input_tensor.shape}")

models/line_detect/line_net.py

@@ -34,7 +34,7 @@ from ..base.base_detection_net import BaseDetectionNet
 import torch.nn.functional as F
 
 from .predict import Predict1, Predict
-from ..base.resnet50fpn import resnet50fpn
+from ..base.high_reso_resnet import resnet50fpn, resnet18fpn
 
 from ..config.config_tool import read_yaml
 
@@ -526,8 +526,60 @@ class LineNet_MobileNet_V3_Large_320_FPN_Weights(WeightsEnum):
         },
     )
     DEFAULT = COCO_V1
+def linenet_newresnet18fpn(
+        *,
+        weights: Optional[LineNet_ResNet50_FPN_Weights] = None,
+        progress: bool = True,
+        num_classes: Optional[int] = None,
+        weights_backbone: Optional[ResNet18_Weights] = ResNet18_Weights.IMAGENET1K_V1,
+        trainable_backbone_layers: Optional[int] = None,
+        **kwargs: Any,
+) -> LineNet:
+    # weights = LineNet_ResNet50_FPN_Weights.verify(weights)
+    # weights_backbone = ResNet50_Weights.verify(weights_backbone)
+
+    if weights is not None:
+        weights_backbone = None
+        num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"]))
+    elif num_classes is None:
+        num_classes = 91
+    if weights_backbone is not None:
+        print(f'resnet50 weights is not None')
+
+    is_trained = weights is not None or weights_backbone is not None
+    trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3)
+    norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d
+
+    backbone =resnet18fpn()
+    featmap_names=['0', '1', '2', '3','pool']
+    # print(f'featmap_names:{featmap_names}')
+    roi_pooler = MultiScaleRoIAlign(
+        featmap_names=featmap_names,
+        output_size=7,
+        sampling_ratio=2
+    )
+    num_features=len(featmap_names)
+    anchor_sizes = tuple((int(16 * 2 ** i),) for i in range(num_features))  # 自动生成不同大小
+    # print(f'anchor_sizes:{anchor_sizes}')
+    aspect_ratios = ((0.5, 1.0, 2.0),) * num_features
+    # print(f'aspect_ratios:{aspect_ratios}')
 
 
+    anchor_generator =  AnchorGenerator(sizes=anchor_sizes, aspect_ratios=aspect_ratios)
+    # anchors = anchor_generator.generate_anchors()
+    # print("Number of anchor sizes:", len(anchor_generator.sizes))  # 应为 5
+    model = LineNet(backbone, num_classes=num_classes,anchor_generator=anchor_generator,
+
+                    box_roi_pool=roi_pooler,
+                    **kwargs)
+
+    if weights is not None:
+        model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True))
+        if weights == LineNet_ResNet50_FPN_Weights.COCO_V1:
+            overwrite_eps(model, 0.0)
+
+    return model
+
 def linenet_newresnet50fpn(
         *,
         weights: Optional[LineNet_ResNet50_FPN_Weights] = None,

models/line_detect/line_predictor.py

@@ -106,11 +106,11 @@ class LineRCNNPredictor(nn.Module):
         else:
             self.pooling = nn.MaxPool1d(scale_factor, scale_factor)
             self.fc2 = nn.Sequential(
-                nn.Linear(self.dim_loi * self.n_pts1 + FEATURE_DIM, self.dim_fc*16),
+                nn.Linear(self.dim_loi * self.n_pts1 + FEATURE_DIM, self.dim_fc*2),
                 nn.ReLU(inplace=True),
-                nn.Linear(self.dim_fc*16, self.dim_fc*8),
+                nn.Linear(self.dim_fc*2, self.dim_fc*2),
                 nn.ReLU(inplace=True),
-                nn.Linear(self.dim_fc*8, 1),
+                nn.Linear(self.dim_fc*2, 1),
             )
         self.loss = nn.BCEWithLogitsLoss(reduction="none")
 

models/line_detect/train.yaml

@@ -1,6 +1,6 @@
 io:
   logdir: train_results
-  datadir: /data/share/lm/Dataset_all
+  datadir: \\192.168.50.222/share/lm/Dataset_all
 #  datadir: D:\python\PycharmProjects\data_20250223\0423_
 #  datadir: I:\datasets\wirenet_1000
 

models/line_detect/train_demo.py

@@ -1,7 +1,7 @@
 import torch
 
 from models.line_detect.line_net import linenet_resnet50_fpn, LineNet, linenet_resnet18_fpn, linenet_newresnet50fpn, \
-    get_line_net_convnext_fpn
+    get_line_net_convnext_fpn, linenet_newresnet18fpn
 from models.line_detect.trainer import Trainer
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -11,8 +11,9 @@ if __name__ == '__main__':
     # model=linenet_resnet50_fpn()
     # model = linenet_resnet18_fpn()
     # model=get_line_net_convnext_fpn(num_classes=2).to(device)
-    model=linenet_newresnet50fpn()
-    # model.load_best_model('train_results/20250622_143530/weights/best_val.pth')
+    # model=linenet_newresnet50fpn()
+    model = linenet_newresnet18fpn()
+    model.load_best_model('train_results/20250622_140412/weights/best_val.pth')
     # trainer = Trainer()
     # trainer.train_cfg(model,cfg='./train.yaml')
     model.start_train(cfg='train.yaml')