소스 검색

修改line_datasets 数据集格式,去掉json中的label

RenLiqiang 3 달 전
부모
커밋
b9ef1d843b
3개의 변경된 파일13개의 추가작업 그리고 6개의 파일을 삭제
  1. 1 0
      models/dataset_tool.py
  2. 10 5
      models/line_detect/dataset_LD.py
  3. 2 1
      models/line_detect/train.yaml

+ 1 - 0
models/dataset_tool.py

@@ -269,6 +269,7 @@ def read_polygon_points_wire(lbl_path, shape):
 def read_masks_from_txt_wire(label_path, shape):
     polygon_points = read_polygon_points_wire(label_path, shape)
     masks = create_masks_from_polygons(polygon_points, shape)
+
     labels = [torch.tensor(item[0]) for item in polygon_points]
 
     return labels, masks

+ 10 - 5
models/line_detect/dataset_LD.py

@@ -105,18 +105,23 @@ class WirePointDataset(BaseDataset):
         }
 
         labels = []
-        if self.target_type == 'polygon':
-            labels, masks = read_masks_from_txt_wire(lbl_path, shape)
-        elif self.target_type == 'pixel':
-            labels = read_masks_from_pixels_wire(lbl_path, shape)
+        #
+        # if self.target_type == 'polygon':
+        #     labels, masks = read_masks_from_txt_wire(lbl_path, shape)
+        # elif self.target_type == 'pixel':
+        #     labels = read_masks_from_pixels_wire(lbl_path, shape)
 
         # print(torch.stack(masks).shape)    # [线段数, 512, 512]
         target = {}
-        target["labels"] = torch.stack(labels)
+        # target["labels"] = torch.stack(labels)
+
+
         target["image_id"] = torch.tensor(item)
         # return wire_labels, target
         target["wires"] = wire_labels
         target["boxes"] = line_boxes(target)
+        target["labels"]= torch.ones(len(target["boxes"]),dtype=torch.int64)
+        # print(f'target["labels"]:{ target["labels"]}')
         # print(f'boxes:{target["boxes"].shape}')
         return target
 

+ 2 - 1
models/line_detect/train.yaml

@@ -1,6 +1,7 @@
 io:
   logdir: logs/
-  datadir: I:/datasets/wirenet_1000
+  datadir: J:/datasets/wirenet_rgb_gray
+#  datadir: I:\datasets\wirenet_1000
   resume_from:
   num_workers: 8
   tensorboard_port: 6000