xiuzheng

libs/vision_libs/models/detection/transform.py

@@ -185,6 +185,7 @@ class GeneralizedRCNNTransform(nn.Module):
         target: Optional[Dict[str, Tensor]] = None,
     ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
         h, w = image.shape[-2:]
+        print(f'h1:{h}, w1:{w}, image.shape[-2:]:{image.shape[-2:]}')
         if self.training:
             if self._skip_resize:
                 return image, target
@@ -200,6 +201,7 @@ class GeneralizedRCNNTransform(nn.Module):
         print(f'bbox:{bbox}')
         print(f'image.shape[-2:]:{image.shape},,,{image.shape[-2:]}')
         bbox = resize_boxes(bbox, (h, w), image.shape[-2:])
+        print(f'bbox after resize:{bbox}')
 
         target["boxes"] = bbox
 
@@ -231,16 +233,13 @@ class GeneralizedRCNNTransform(nn.Module):
 
         if "mask_ends" in target:
             arc_mask = target["mask_ends"]
-            # print(f'arc_mask in 1:{arc_mask}')
             arc_mask = resize_keypoints(arc_mask, (h, w), image.shape[-2:])
-            # print(f'arc_mask in 2:{arc_mask}')
             target["mask_ends"] = arc_mask
+
         if "mask_params" in target:
             arc_mask = target["mask_params"]
-            # print(f'arc_mask in 3:{arc_mask}')
             arc_mask = resize_keypoints(arc_mask, (h, w), image.shape[-2:])
             arc_mask[:,2:4] = resize_keypoints(arc_mask[:,2:4], (h, w), image.shape[-2:])
-            # print(f'arc_mask in 4:{arc_mask}')
             target["mask_params"] = arc_mask
         return image, target
 

models/base/base_detection_net.py

@@ -116,11 +116,8 @@ 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/line_detect/heads/arc/arc_heads.py

@@ -84,7 +84,6 @@ class ArcEquationHead(nn.Module):
         feature_logits: [N, 1, H, W]
         """
         N, _, H, W = feature_logits.shape
-        print(f'N:{N}, H:{H}, W:{W}')
 
         # --------------------------------------------
         # Global average pooling

models/line_detect/heads/head_losses.py

@@ -878,20 +878,23 @@ def compute_ins_loss(feature_logits, proposals, gt_, pos_matched_idxs):
 def align_masks(keypoints, rois, heatmap_size):
     print(f'rois:{rois.shape}')
     print(f'heatmap_size:{heatmap_size}')
-
     print(f'keypoints.shape:{keypoints.shape}')
-    # batch_size, num_keypoints, _ = keypoints.shape
-    t_h, t_w = keypoints.shape[-2:]
-    scale=heatmap_size/t_w
-    print(f'scale:{scale}')
-    x = keypoints[..., 0]*scale
-    y = keypoints[..., 1]*scale
 
-    x = x.unsqueeze(1)
-    y = y.unsqueeze(1)
+    # t_h, t_w = keypoints.shape[-2:]
+    # scale=heatmap_size/t_w
+    # print(f'scale:{scale}')
+
+    # x = keypoints[..., 0]*scale
+    # y = keypoints[..., 1]*scale
+    #
+    # x = x.unsqueeze(1)
+    # y = y.unsqueeze(1)
+    #
+    # num_points=x.shape[2]
+    # print(f'num_points:{num_points}')
 
-    num_points=x.shape[2]
-    print(f'num_points:{num_points}')
+    # plt.imshow(keypoints[0].cpu())
+    # plt.show()
     mask_4d = keypoints.unsqueeze(1).float()
     resized_mask = F.interpolate(
         mask_4d,

models/line_detect/line_dataset.py

@@ -98,7 +98,6 @@ class LineDataset(BaseDataset):
 
         target["image_id"] = torch.tensor(item)
 
-        #boxes, line_point_pairs, points, labels, mask_ends, mask_params
         boxes, lines, points, labels, arc_ends, arc_params = get_boxes_lines(objs, shape)
 
         if lines is not None:
@@ -120,16 +119,13 @@ class LineDataset(BaseDataset):
 
         print(f'target[circle_masks]:{target["circle_masks"].shape}')
 
-
         target["boxes"] = boxes
         target["labels"] = labels
         target["img_size"] = shape
 
-        # validate_keypoints(lines, shape[0], shape[1])
         return target
 
 
-
     def show(self, idx, show_type='all'):
         image, target = self.__getitem__(idx)
 
@@ -141,6 +137,23 @@ class LineDataset(BaseDataset):
         img = image
         # print(f'img:{img.shape}')
 
+        if show_type == 'Original_mask':
+
+            image = image.permute(1, 2, 0).cpu().numpy()  # (3,H,W) -> (H,W,3)
+            arc_mask = target['circle_masks'][-1]
+
+            plt.figure(figsize=(10, 5))
+            plt.subplot(1, 2, 1)
+            plt.imshow(image)
+            plt.title('Original')
+
+            # 叠加 mask（用红色轮廓或填充）
+            plt.subplot(1, 2, 2)
+            plt.imshow(image)
+            plt.imshow(arc_mask, cmap='Reds', alpha=0.5)  # 半透明红色
+            plt.title('Image with Mask')
+            plt.show()
+
         if show_type == 'arc_yuan_point_ellipse':
             arc_ends = target['mask_ends']
             arc_params = target['mask_params']
@@ -171,10 +184,7 @@ class LineDataset(BaseDataset):
 
         if show_type == 'circle_masks':
             arc_mask = target['circle_masks']
-            # print(f'taget circle:{arc.shape}')
             print(f'target circle_masks:{arc_mask.shape}')
-            combined = torch.cat(list(arc_mask), dim=1)
-            print(f'combine:{combined.shape}')
             plt.imshow(arc_mask[-1])
             plt.show()
 
@@ -314,7 +324,6 @@ def arc_to_mask_safe(arc_param, arc_end, shape, line_width=5, debug=True, idx=-1
     return torch.tensor(mask, dtype=torch.float32)
 
 
-
 def draw_el(all):
     # 解析椭圆参数
     if isinstance(all, torch.Tensor):
@@ -535,47 +544,41 @@ def get_boxes_lines(objs, shape):
     mask_params = []
 
     for obj in objs:
-        # plt.plot([a[1], b[1]], [a[0], b[0]], c="red", linewidth=1)  # a[1], b[1]无明确大小
 
-        # print(f"points:{obj['points']}")
         label = obj['label']
-        if label == 'line' or label == 'dseam1':
-            a, b = obj['points'][0], obj['points'][1]
-
-            # line_point_pairs.append(a)
-            # line_point_pairs.append(b)
-            line_point_pairs.append([a, b])
-
-            xmin = max(0, (min(a[0], b[0]) - 6))
-            xmax = min(w, (max(a[0], b[0]) + 6))
-            ymin = max(0, (min(a[1], b[1]) - 6))
-            ymax = min(h, (max(a[1], b[1]) + 6))
-
-            boxes.append([xmin, ymin, xmax, ymax])
-            labels.append(torch.tensor(2))
-
-            points.append(torch.tensor([0.0]))
-            mask_ends.append([[0, 0], [0, 0]])
-            mask_params.append([0, 0, 0, 0, 0])
-            # circle_4points.append([[0, 0], [0, 0], [0, 0], [0, 0]])
-
-
-
-        elif label == 'point':
-            p = obj['points'][0]
-            xmin = max(0, p[0] - 12)
-            xmax = min(w, p[0] + 12)
-            ymin = max(0, p[1] - 12)
-            ymax = min(h, p[1] + 12)
-
-            points.append(p)
-            labels.append(torch.tensor(1))
-            boxes.append([xmin, ymin, xmax, ymax])
-
-            line_point_pairs.append([[0, 0], [0, 0]])
-            mask_ends.append([[0, 0], [0, 0]])
-            mask_params.append([0, 0, 0, 0, 0])
-            # circle_4points.append([[0, 0], [0, 0], [0, 0], [0, 0]])
+        # if label == 'line' or label == 'dseam1':
+        #     a, b = obj['points'][0], obj['points'][1]
+        #     line_point_pairs.append([a, b])
+        #
+        #     xmin = max(0, (min(a[0], b[0]) - 6))
+        #     xmax = min(w, (max(a[0], b[0]) + 6))
+        #     ymin = max(0, (min(a[1], b[1]) - 6))
+        #     ymax = min(h, (max(a[1], b[1]) + 6))
+        #
+        #     boxes.append([xmin, ymin, xmax, ymax])
+        #     labels.append(torch.tensor(2))
+        #
+        #     points.append(torch.tensor([0.0]))
+        #     mask_ends.append([[0, 0], [0, 0]])
+        #     mask_params.append([0, 0, 0, 0, 0])
+        #     # circle_4points.append([[0, 0], [0, 0], [0, 0], [0, 0]])
+        #
+        #
+        # elif label == 'point':
+        #     p = obj['points'][0]
+        #     xmin = max(0, p[0] - 12)
+        #     xmax = min(w, p[0] + 12)
+        #     ymin = max(0, p[1] - 12)
+        #     ymax = min(h, p[1] + 12)
+        #
+        #     points.append(p)
+        #     labels.append(torch.tensor(1))
+        #     boxes.append([xmin, ymin, xmax, ymax])
+        #
+        #     line_point_pairs.append([[0, 0], [0, 0]])
+        #     mask_ends.append([[0, 0], [0, 0]])
+        #     mask_params.append([0, 0, 0, 0, 0])
+        #     # circle_4points.append([[0, 0], [0, 0], [0, 0], [0, 0]])
 
 
         # elif label == 'arc':
@@ -598,7 +601,7 @@ def get_boxes_lines(objs, shape):
         #     line_point_pairs.append([[0, 0], [0, 0]])
         #     circle_4points.append([[0, 0], [0, 0], [0, 0], [0, 0]])
 
-        elif label == 'arc':
+        if label == 'arc':
 
             arc_params = obj['params']
             arc_ends = obj['ends']
@@ -653,4 +656,5 @@ def get_boxes_lines(objs, shape):
 if __name__ == '__main__':
     path = r'/data/share/zyh/master_dataset/dataset_net/pokou_251115_251121/a_dataset'
     dataset = LineDataset(dataset_path=path, dataset_type='train', augmentation=False, data_type='jpg')
-    dataset.show(19, show_type='circle_masks')
+    for i in range(100):
+         dataset.show(i, show_type='Original_mask')

models/line_detect/loi_heads.py

@@ -1069,6 +1069,7 @@ class RoIHeads(nn.Module):
                 losses.update(loss_line)
                 losses.update(loss_line_iou)
                 print(f'losses:{losses}')
+
         if self.has_point() and self.detect_point:
             print(f'roi_heads forward has_point()!!!!')
             # print(f'labels:{labels}')
@@ -1188,7 +1189,6 @@ class RoIHeads(nn.Module):
                 losses.update(loss_point)
                 print(f'losses:{losses}')
 
-
         if self.has_arc() and self.detect_arc:
             print(f'roi_heads forward has_arc()!!!!')
             # print(f'labels:{labels}')
@@ -1320,7 +1320,6 @@ class RoIHeads(nn.Module):
             # print(f'ins_proposals11:{ins_proposals}')
             print(f'boxes_proposals:{len(ins_proposals)}')
 
-
             if self.training:
                 # during training, only focus on positive boxes
                 num_images = len(proposals)
@@ -1331,9 +1330,6 @@ class RoIHeads(nn.Module):
                     raise ValueError("if in trainning, matched_idxs should not be None")
                 for img_id in range(num_images):
 
-                    # circle_pos = torch.where(labels[img_id] == 4)[0]
-                    # ins_proposals.append(proposals[img_id][circle_pos])
-                    # ins_pos_matched_idxs.append(matched_idxs[img_id][circle_pos])
                     circle_pos = torch.where(labels[img_id] == 4)[0]
                     circle_pos = circle_pos.flatten()
                     idxs = circle_pos.detach().cpu().tolist()
@@ -1344,12 +1340,8 @@ class RoIHeads(nn.Module):
                                 f"Index out of bounds: circle_pos={idx}, but proposals len={num_prop}, "
                                 f"img_id={img_id}"
                             )
-                    ins_proposals.append(
-                        proposals[img_id][idxs]
-                    )
-                    ins_pos_matched_idxs.append(
-                        matched_idxs[img_id][idxs]
-                    )
+                    ins_proposals.append(proposals[img_id][idxs])
+                    ins_pos_matched_idxs.append(matched_idxs[img_id][idxs])
                     print(f'ins_proposals in train:{ins_proposals}')
             else:
                 if targets is not None:
@@ -1363,9 +1355,6 @@ class RoIHeads(nn.Module):
                         raise ValueError("if in trainning, matched_idxs should not be None")
 
                     for img_id in range(num_images):
-                        # circle_pos = torch.where(labels[img_id] == 4)[0]
-                        # ins_proposals.append(proposals[img_id][circle_pos])
-                        # ins_pos_matched_idxs.append(matched_idxs[img_id][circle_pos])
                         circle_pos = torch.where(labels[img_id] == 4)[0]
                         circle_pos = circle_pos.flatten()
                         idxs = circle_pos.detach().cpu().tolist()
@@ -1376,12 +1365,8 @@ class RoIHeads(nn.Module):
                                     f"Index out of bounds: circle_pos={idx}, but proposals len={num_prop}, "
                                     f"img_id={img_id}"
                                 )
-                        ins_proposals.append(
-                            proposals[img_id][idxs]
-                        )
-                        ins_pos_matched_idxs.append(
-                            matched_idxs[img_id][idxs]
-                        )
+                        ins_proposals.append(proposals[img_id][idxs])
+                        ins_pos_matched_idxs.append(matched_idxs[img_id][idxs])
                         print(f'ins_proposals in val:{ins_proposals}')
                 else:
                     pos_matched_idxs = None
@@ -1389,17 +1374,21 @@ class RoIHeads(nn.Module):
             # circle_proposals_tensor=torch.cat(circle_proposals)
 
             ins_proposals_valid = self.check_proposals(ins_proposals)
-            print(f"self.train{self.training}")
-            print(f"self.val{ins_proposals_valid}")
-            if  ins_proposals_valid:
+            print(f"self.train:{self.training}")
+            print(f"proposals_valid:{ins_proposals_valid}")
 
+            if  ins_proposals_valid:
                 print(f'features from backbone:{features['0'].shape}')
                 print(f'ins_proposals in ins_forward1:{ins_proposals}')
                 feature_logits = self.ins_forward1(features, image_shapes, ins_proposals)
                 # ins_masks, ins_scores, circle_points = ins_inference(feature_logits,
                 #                                                      ins_proposals, th=0)
 
-                arc_equation = self.arc_equation_head(feature_logits)  # [proposal和，7]
+                # plt.title('feature_logits')
+                # plt.imshow(feature_logits[0][0].cpu().detach().numpy())
+                # plt.show()
+
+                arc_equation = self.arc_equation_head(feature_logits)  # [proposal和，9]
 
                 loss_ins = None
                 loss_ins_extra=None
@@ -1419,9 +1408,8 @@ class RoIHeads(nn.Module):
                     gt_mask_ends = [t["mask_ends"] for t in targets if "mask_ends" in t]
                     gt_mask_params = [t["mask_params"] for t in targets if "mask_params" in t]
 
-                    # print(f'gt_ins:{gt_inses[0].shape}')
-                    h, w = targets[0]["img_size"]
-                    img_size = h
+                    # h, w = targets[0]["img_size"]
+                    # img_size = h
 
                     gt_ins_tensor = torch.zeros(0, 0)
                     if len(gt_inses) > 0:
@@ -1431,7 +1419,6 @@ class RoIHeads(nn.Module):
 
                     if gt_ins_tensor.shape[0] > 0:
                         print(f'start to compute circle_loss')
-
                         loss_ins = compute_ins_loss(feature_logits, ins_proposals, gt_inses,ins_pos_matched_idxs)
                         # total_loss, loss_arc_equation, loss_arc_ends = compute_arc_equation_loss(arc_equation,ins_proposals,gt_mask_ends,gt_mask_params,ins_pos_matched_idxs,labels)
                         # loss_arc_ends = loss_arc_ends
@@ -1447,20 +1434,19 @@ class RoIHeads(nn.Module):
                         print(f'loss_ins is None111')
                         loss_ins = torch.tensor(0.0, device=device)
 
-                    if loss_ins_extra is None:
-                        print(f'loss_ins_extra is None111')
-                        loss_ins_extra = torch.tensor(0.0, device=device)
+                    # if loss_ins_extra is None:
+                    #     print(f'loss_ins_extra is None111')
+                    #     loss_ins_extra = torch.tensor(0.0, device=device)
 
                     loss_ins = {"loss_ins": loss_ins}
                     # loss_ins_extra = {"loss_ins_extra": loss_ins_extra}
                     # loss_arc_equation = {"loss_arc_equation": loss_arc_equation}
                     # loss_arc_ends = {"loss_arc_ends": loss_arc_ends}
 
-
                 else:
                     if targets is not None:
-                        h, w = targets[0]["img_size"]
-                        img_size = h
+                        # h, w = targets[0]["img_size"]
+                        # img_size = h
                         gt_inses = [t["circle_masks"] for t in targets if "circle_masks" in t]
                         gt_labels = [t["labels"] for t in targets]
                         gt_mask_ends = [t["mask_ends"] for t in targets if "mask_ends" in t]
@@ -1473,17 +1459,16 @@ class RoIHeads(nn.Module):
                         if gt_ins_tensor.shape[0] > 0:
                             print(f'start to compute circle_loss')
 
-                            loss_ins = compute_ins_loss(feature_logits, ins_proposals, gt_inses,
-                                                           ins_pos_matched_idxs)
+                            loss_ins = compute_ins_loss(feature_logits, ins_proposals, gt_inses, ins_pos_matched_idxs)
                             # total_loss, loss_arc_equation, loss_arc_ends = compute_arc_equation_loss(arc_equation,ins_proposals,gt_mask_ends,gt_mask_params,ins_pos_matched_idxs,labels)
                             #
                             # loss_arc_ends = loss_arc_ends
 
                             # loss_ins_extra = compute_circle_extra_losses(feature_logits, circle_proposals, gt_circles,circle_pos_matched_idxs)
 
-                        # if loss_ins is None:
-                        #     print(f'loss_ins is None111')
-                        #     loss_ins = torch.tensor(0.0, device=device)
+                        if loss_ins is None:
+                            print(f'loss_ins is None111')
+                            loss_ins = torch.tensor(0.0, device=device)
 
                         # if loss_ins_extra is None:
                         #     print(f'loss_ins_extra is None111')
@@ -1497,21 +1482,11 @@ class RoIHeads(nn.Module):
                         #     print(f'loss_arc_ends is None')
                         #     loss_arc_ends = torch.tensor(0.0, device=device)
 
-                        if loss_ins is None:
-                            print(f'loss_ins is None111')
-                            loss_ins = torch.tensor(0.0, device=device)
-
-                        # if loss_ins_extra is None:
-                        #     print(f'loss_ins_extra is None111')
-                        #     loss_ins_extra = torch.tensor(0.0, device=device)
-
                         loss_ins = {"loss_ins": loss_ins}
                         # loss_ins_extra = {"loss_ins_extra": loss_ins_extra}
                         # loss_arc_equation = {"loss_arc_equation": loss_arc_equation}
                         # loss_arc_ends = {"loss_arc_ends": loss_arc_ends}
 
-
-
                     else:
                         loss_ins = {}
                         # loss_ins_extra = {}
@@ -1551,7 +1526,7 @@ class RoIHeads(nn.Module):
 
 
                 print(f'loss_ins:{loss_ins}')
-                print(f'loss_ins_extra:{loss_ins_extra}')
+                # print(f'loss_ins_extra:{loss_ins_extra}')
                 losses.update(loss_ins)
                 # losses.update(loss_ins_extra)
                 # losses.update(loss_arc_equation)