RenLiqiang hai 3 meses
pai
achega
cc80e50f53
Modificáronse 2 ficheiros con 40 adicións e 24 borrados
  1. 5 6
      models/line_detect/line_net.py
  2. 35 18
      models/line_detect/line_predictor.py

+ 5 - 6
models/line_detect/line_net.py

@@ -68,7 +68,6 @@ class LineNet(BaseDetectionNet):
     #
     #     self.__construct__(backbone=backbone, num_classes=num_classes, **kwargs)
 
-
     def __init__(
             self,
             backbone,
@@ -202,17 +201,15 @@ class LineNet(BaseDetectionNet):
 
         super().__init__(backbone, rpn, roi_heads, transform)
 
-
         self.roi_heads = roi_heads
 
-        self.roi_heads.line_head = line_head
-        self.roi_heads.line_predictor = line_predictor
+        # self.roi_heads.line_head = line_head
+        # self.roi_heads.line_predictor = line_predictor
 
     def train_by_cfg(self, cfg):
         # cfg = read_yaml(cfg)
         self.trainer = Trainer()
-        self.trainer.train_cfg(model=self,cfg=cfg)
-
+        self.trainer.train_cfg(model=self, cfg=cfg)
 
 
 class TwoMLPHead(nn.Module):
@@ -496,6 +493,8 @@ def linenet_resnet50_fpn(
         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)

+ 35 - 18
models/line_detect/line_predictor.py

@@ -47,7 +47,24 @@ class Bottleneck1D(nn.Module):
         return x + self.op(x)
 
 class LineRCNNPredictor(nn.Module):
-    def __init__(self,**kwargs):
+    def __init__(self,n_pts0 = 32,
+                 n_pts1 = 8,
+                 n_stc_posl =300,
+                 dim_loi = 128,
+                 use_conv = 0,
+                 dim_fc = 1024,
+                 n_out_line = 2500,
+                 n_out_junc =250,
+                 n_dyn_junc = 300,
+                 eval_junc_thres = 0.008,
+                 n_dyn_posl =300,
+                 n_dyn_negl =80,
+                 n_dyn_othr = 600,
+                 use_cood = 0,
+                 use_slop = 0,
+                 n_stc_negl = 40,
+                 head_size = [[2], [1], [2]] ,
+                 **kwargs):
         super().__init__()
         # self.backbone = backbone
         # self.cfg = read_yaml(cfg)
@@ -76,24 +93,24 @@ class LineRCNNPredictor(nn.Module):
         # self.head_size = self.cfg['head_size']
 
 
-        self.n_pts0 = 32
-        self.n_pts1 = 8
-        self.n_stc_posl =300
-        self.dim_loi = 128
-        self.use_conv = 0
-        self.dim_fc = 1024
-        self.n_out_line = 2500
-        self.n_out_junc =250
+        self.n_pts0 = n_pts0
+        self.n_pts1 = n_pts1
+        self.n_stc_posl =n_stc_posl
+        self.dim_loi = dim_loi
+        self.use_conv = use_conv
+        self.dim_fc = dim_fc
+        self.n_out_line = n_out_line
+        self.n_out_junc =n_out_junc
         # self.loss_weight =
-        self.n_dyn_junc = 300
-        self.eval_junc_thres = 0.008
-        self.n_dyn_posl =300
-        self.n_dyn_negl =80
-        self.n_dyn_othr = 600
-        self.use_cood = 0
-        self.use_slop = 0
-        self.n_stc_negl = 80
-        self.head_size = [[2], [1], [2]]
+        self.n_dyn_junc = n_dyn_junc
+        self.eval_junc_thres = eval_junc_thres
+        self.n_dyn_posl =n_dyn_posl
+        self.n_dyn_negl = n_dyn_negl
+        self.n_dyn_othr = n_dyn_othr
+        self.use_cood = use_cood
+        self.use_slop = use_slop
+        self.n_stc_negl = n_stc_negl
+        self.head_size = head_size
 
         self.num_class = sum(sum(self.head_size, []))
         self.head_off = np.cumsum([sum(h) for h in self.head_size])