all boxes is ok with only train arc

models/line_detect/line_dataset.py

@@ -546,39 +546,39 @@ def get_boxes_lines(objs, shape):
     for obj in objs:
 
         label = obj['label']
-        # 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]])
+        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':
@@ -601,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]])
 
-        if label == 'arc':
+        elif label == 'arc':
 
             arc_params = obj['params']
             arc_ends = obj['ends']

models/line_detect/trainer.py

@@ -5,6 +5,7 @@ from datetime import datetime
 import cv2
 import numpy as np
 import torch
+import torchvision
 from PIL.ImageDraw import ImageDraw
 from matplotlib import pyplot as plt
 from scipy.ndimage import gaussian_filter
@@ -17,6 +18,7 @@ from models.base.base_trainer import BaseTrainer
 from models.config.config_tool import read_yaml
 from models.line_detect.line_dataset import LineDataset
 import torch.nn.functional as F
+from torchvision.transforms import functional as TF
 
 from tools import utils
 
@@ -284,7 +286,7 @@ def draw_lines_with_scores(tensor_image, lines, scores, width=3, cmap='viridis')
     colors = (colormap(scores.cpu().numpy())[:, :3] * 255).astype('uint8')  # 去掉 alpha 通道
 
     # 转为 PIL 画图
-    image_pil = F.to_pil_image(tensor_image)
+    image_pil = TF.to_pil_image(tensor_image)
     draw = ImageDraw.Draw(image_pil)
 
     for line, color in zip(lines, colors):
@@ -292,7 +294,7 @@ def draw_lines_with_scores(tensor_image, lines, scores, width=3, cmap='viridis')
         end = tuple(map(float, line[1][:2].tolist()))
         draw.line([start, end], fill=tuple(color), width=width)
 
-    return (F.to_tensor(image_pil) * 255).to(torch.uint8)
+    return (torchvision.transforms.functional.to_tensor(image_pil) * 255).to(torch.uint8)
 
 
 class Trainer(BaseTrainer):