调试convnext推理功能

models/line_detect/111.py

@@ -150,7 +150,7 @@ class Trainer(BaseTrainer):
     def train(self, model, **kwargs):
         dataset_train = WirePointDataset(dataset_path=kwargs['io']['datadir'], dataset_type='train')
         train_sampler = torch.utils.data.RandomSampler(dataset_train)
-        train_batch_sampler = torch.utils.data.BatchSampler(train_sampler, batch_size=8, drop_last=True)
+        train_batch_sampler = torch.utils.data.BatchSampler(train_sampler, batch_size=2, drop_last=True)
         train_collate_fn = utils.collate_fn_wirepoint
         data_loader_train = torch.utils.data.DataLoader(
             dataset_train, batch_sampler=train_batch_sampler, num_workers=1, collate_fn=train_collate_fn
@@ -158,7 +158,7 @@ class Trainer(BaseTrainer):
 
         dataset_val = WirePointDataset(dataset_path=kwargs['io']['datadir'], dataset_type='val')
         val_sampler = torch.utils.data.RandomSampler(dataset_val)
-        val_batch_sampler = torch.utils.data.BatchSampler(val_sampler, batch_size=4, drop_last=True)
+        val_batch_sampler = torch.utils.data.BatchSampler(val_sampler, batch_size=1, drop_last=True)
         val_collate_fn = utils.collate_fn_wirepoint
         data_loader_val = torch.utils.data.DataLoader(
             dataset_val, batch_sampler=val_batch_sampler, num_workers=1, collate_fn=val_collate_fn
@@ -234,8 +234,9 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 if __name__ == '__main__':
     # model = LineNet('line_net.yaml')
     # model = linenet_resnet50_fpn().to(device)
+    model=linenet_resnet18_fpn().to(device)
     # model=get_line_net_efficientnetv2(2, pretrained_backbone=True).to(device)
-    model=get_line_net_convnext_fpn(num_classes=2).to(device)
+    # model=get_line_net_convnext_fpn(num_classes=2).to(device)
     # model=linenet_resnet18_fpn()
     trainer = Trainer()
     trainer.train_cfg(model,cfg='./train.yaml')

models/line_detect/predict2.py

@@ -1,158 +1,3 @@
-# import time
-#
-# from models.line_detect.postprocess import show_predict
-# import os
-#
-# import torch
-# from PIL import Image
-# import matplotlib.pyplot as plt
-# import matplotlib as mpl
-# import numpy as np
-# from models.line_detect.line_net import linenet_resnet50_fpn
-# from torchvision import transforms
-# from rtree import index
-# # from models.wirenet.postprocess import postprocess
-# from models.wirenet.postprocess import postprocess
-#
-# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-#
-#
-# def load_best_model(model, save_path, device):
-#     if os.path.exists(save_path):
-#         checkpoint = torch.load(save_path, map_location=device)
-#         model.load_state_dict(checkpoint['model_state_dict'])
-#         # if optimizer is not None:
-#         #     optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
-#         epoch = checkpoint['epoch']
-#         loss = checkpoint['loss']
-#         print(f"Loaded best model from {save_path} at epoch {epoch} with loss {loss:.4f}")
-#     else:
-#         print(f"No saved model found at {save_path}")
-#     return model
-#
-#
-# def box_line_optimized(pred):
-#     # 创建R-tree索引
-#     idx = index.Index()
-#
-#     # 将所有线段添加到R-tree中
-#     lines = pred[-1]['wires']['lines']  # 形状为[1, 2500, 2, 2]
-#     scores = pred[-1]['wires']['score'][0]  # 假设形状为[2500]
-#
-#     # 提取并处理所有线段
-#     for idx_line in range(lines.shape[1]):  # 遍历2500条线段
-#         line_tensor = lines[0, idx_line].cpu().numpy() / 128 * 512  # 转换为numpy数组并调整比例
-#         x_min = float(min(line_tensor[0][0], line_tensor[1][0]))
-#         y_min = float(min(line_tensor[0][1], line_tensor[1][1]))
-#         x_max = float(max(line_tensor[0][0], line_tensor[1][0]))
-#         y_max = float(max(line_tensor[0][1], line_tensor[1][1]))
-#         idx.insert(idx_line, (x_min, y_min, x_max, y_max))
-#
-#     for idx_box, box_ in enumerate(pred[0:-1]):
-#         box = box_['boxes'].cpu().numpy()  # 确保将张量转换为numpy数组
-#         line_ = []
-#         score_ = []
-#
-#         for i in box:
-#             score_max = 0.0
-#             tmp = [[0.0, 0.0], [0.0, 0.0]]
-#
-#             # 获取与当前box可能相交的所有线段
-#             possible_matches = list(idx.intersection((i[0], i[1], i[2], i[3])))
-#
-#             for j in possible_matches:
-#                 line_j = lines[0, j].cpu().numpy() / 128 * 512
-#                 if (line_j[0][1] >= i[0] and line_j[1][1] >= i[0] and  # 注意这里交换了x和y
-#                         line_j[0][1] <= i[2] and line_j[1][1] <= i[2] and
-#                         line_j[0][0] >= i[1] and line_j[1][0] >= i[1] and
-#                         line_j[0][0] <= i[3] and line_j[1][0] <= i[3]):
-#
-#                     if scores[j] > score_max:
-#                         tmp = line_j
-#                         score_max = scores[j]
-#
-#             line_.append(tmp)
-#             score_.append(score_max)
-#
-#         processed_list = torch.tensor(line_)
-#         pred[idx_box]['line'] = processed_list
-#
-#         processed_s_list = torch.tensor(score_)
-#         pred[idx_box]['line_score'] = processed_s_list
-#
-#     return pred
-#
-# # def box_line_(pred):
-# #     for idx, box_ in enumerate(pred[0:-1]):
-# #         box = box_['boxes']  # 是一个tensor
-# #         line = pred[-1]['wires']['lines'][idx].cpu().numpy() / 128 * 512
-# #         score = pred[-1]['wires']['score'][idx]
-# #         line_ = []
-# #         score_ = []
-# #
-# #         for i in box:
-# #             score_max = 0.0
-# #             tmp = [[0.0, 0.0], [0.0, 0.0]]
-# #
-# #             for j in range(len(line)):
-# #                 if (line[j][0][1] >= i[0] and line[j][1][1] >= i[0] and
-# #                         line[j][0][1] <= i[2] and line[j][1][1] <= i[2] and
-# #                         line[j][0][0] >= i[1] and line[j][1][0] >= i[1] and
-# #                         line[j][0][0] <= i[3] and line[j][1][0] <= i[3]):
-# #
-# #                     if score[j] > score_max:
-# #                         tmp = line[j]
-# #                         score_max = score[j]
-# #             line_.append(tmp)
-# #             score_.append(score_max)
-# #         processed_list = torch.tensor(line_)
-# #         pred[idx]['line'] = processed_list
-# #
-# #         processed_s_list = torch.tensor(score_)
-# #         pred[idx]['line_score'] = processed_s_list
-# #     return pred
-#
-#
-# def predict(pt_path, model, img):
-#     model = load_best_model(model, pt_path, device)
-#
-#     model.eval()
-#
-#     if isinstance(img, str):
-#         img = Image.open(img).convert("RGB")
-#
-#     transform = transforms.ToTensor()
-#     img_tensor = transform(img)
-#
-#     with torch.no_grad():
-#         t_start = time.time()
-#         predictions = model([img_tensor.to(device)])
-#         t_end=time.time()
-#         print(f'predict used:{t_end-t_start}')
-#         # print(f'predictions:{predictions}')
-#         boxes=predictions[0]['boxes'].shape
-#         lines=predictions[-1]['wires']['lines'].shape
-#         lines_scores=predictions[-1]['wires']['score'].shape
-#         print(f'predictions boxes:{boxes},lines:{lines},lines_scores:{lines_scores}')
-#     t_start=time.time()
-#     pred = box_line_optimized(predictions)
-#     t_end=time.time()
-#     print(f'matched boxes and lines used:{t_end - t_start}')
-#     # print(f'pred:{pred[0]}')
-#     show_predict(img_tensor, pred, t_start)
-#
-#
-# if __name__ == '__main__':
-#     t_start = time.time()
-#     print(f'start to predict:{t_start}')
-#     model = linenet_resnet50_fpn().to(device)
-#     pt_path = r"F:\BaiduNetdiskDownload\resnet50_best_e8.pth"
-#     img_path = r"I:\datasets\wirenet_1000\images\val\00035148_0.png"
-#     predict(pt_path, model, img_path)
-#     t_end = time.time()
-#     # print(f'predict used:{t_end - t_start}')
-
-
 import time
 
 import cv2
@@ -165,13 +10,12 @@ from PIL import Image
 import matplotlib.pyplot as plt
 import matplotlib as mpl
 import numpy as np
-from models.line_detect.line_net import linenet_resnet50_fpn
+from models.line_detect.line_net import linenet_resnet50_fpn, get_line_net_efficientnetv2, get_line_net_convnext_fpn
 from torchvision import transforms
 
 # from models.wirenet.postprocess import postprocess
 from models.wirenet.postprocess import postprocess
 from rtree import index
-import imageio
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
@@ -179,7 +23,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 def load_best_model(model, save_path, device):
     if os.path.exists(save_path):
         checkpoint = torch.load(save_path, map_location=device)
-        model.load_state_dict(checkpoint['model_state_dict'])
+        model.load_state_dict(checkpoint['model_state_dict'],strict=False)
         # if optimizer is not None:
         #     optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
         # epoch = checkpoint['epoch']
@@ -424,18 +268,11 @@ def predict(pt_path, model, img):
 
     model.eval()
 
-    # if isinstance(img, str):
-    #     img = Image.open(img).convert("RGB")
-    print(imageio.v3.imread(img_path).shape)
-    img = imageio.v3.imread(img_path).reshape(2114, 1332, 1)
-    img_3channel = np.zeros((2114, 1332, 3), dtype=img.dtype)
-    img_3channel[:, :, 2] = img[:, :, 0]
-    img = torch.from_numpy(img_3channel).permute(2, 0, 1)
+    if isinstance(img, str):
+        img = Image.open(img).convert("RGB")
 
-    img_tensor = img
-
-    # transform = transforms.ToTensor()
-    # img_tensor = transform(img)  # [3, 512, 512]
+    transform = transforms.ToTensor()
+    img_tensor = transform(img)  # [3, 512, 512]
 
     # 将图像调整为512x512大小
     t_start = time.time()
@@ -445,7 +282,8 @@ def predict(pt_path, model, img):
 
     im = img_tensor.permute(1, 2, 0)  # [H, W, 3]
     if im.shape != (512, 512, 3):
-        im = cv2.resize(im.cpu().numpy().astype(np.float32), (512, 512), interpolation=cv2.INTER_NEAREST)
+        # im = cv2.resize(im.cpu().numpy().astype(np.float32), (512, 512), interpolation=cv2.INTER_NEAREST)
+        im = cv2.resize(im.cpu().numpy().astype(np.float32), (512, 512), interpolation=cv2.INTER_NEAREST_EXACT)
     img_ = torch.tensor(im).permute(2, 0, 1)  # [3, 512, 512]
 
     t_end = time.time()
@@ -471,18 +309,18 @@ def predict(pt_path, model, img):
     #
     # show_predict(img_, pred, t_start)
 
-from models.line_detect.line_net import linenet_resnet50_fpn, LineNet, linenet_resnet18_fpn, get_line_net_efficientnetv2
 
 if __name__ == '__main__':
     t_start = time.time()
     print(f'start to predict:{t_start}')
     # model = linenet_resnet50_fpn().to(device)
-    model = get_line_net_efficientnetv2(2, pretrained_backbone=True).to(device)
+    # model = get_line_net_efficientnetv2(2, pretrained_backbone=True).to(device)
+    model=get_line_net_convnext_fpn(num_classes=2).to(device)
     # pt_path = r"C:\Users\m2337\Downloads\best_lmap代替x，训练24轮结果.pth"
     # pt_path = r"C:\Users\m2337\Downloads\best_lmap代替x，训练75轮.pth"
-    pt_path = r"\\192.168.50.222\share\lm\weight\20250510_155941\weights\best.pth"
+    pt_path = r"\\192.168.50.222\share\rlq\weights\convnext250514.pth"
     # pt_path = r"C:\Users\m2337\Downloads\best_e20.pth"
-    img_path = r"D:\python\PycharmProjects\20250214\cloud\新建文件夹\depth_map.tiff"
+    img_path = r"color_img.jpg"
     predict(pt_path, model, img_path)
     t_end = time.time()
-    print(f'predict used:{t_end - t_start}')
+    print(f'predict used:{t_end - t_start}')

models/line_detect/train.yaml

@@ -1,6 +1,6 @@
 io:
   logdir: logs/
-  datadir: \\192.168.50.222/share/zyh/513/a_dataset
+  datadir: \\192.168.50.222/share/zyh/513train/a_dataset
 #  datadir: D:\python\PycharmProjects\data_20250223\0423_
 #  datadir: I:\datasets\wirenet_1000
   resume_from: