1-N实践yolov7目标检测算法-基于已有模型训练算法实现物体检测#

1
# 划分数据集
2
import os
3
import random
4

5
def split_files(folder_path, ratio):
6
    # 获取文件夹中的文件列表
7
    file_list = os.listdir(folder_path)
8

9
    # 计算划分的文件数量
10
    total_files = len(file_list)
11
    num_files_part1 = int(total_files * ratio)
12
    num_files_part2 = total_files - num_files_part1
13

14
    # 随机划分文件
15
    random.shuffle(file_list)
16
    part1_files = file_list[:num_files_part1]
17
    part2_files = file_list[num_files_part1:]
18

19
    # 将划分结果写入 txt 文件
20
    with open("./data/fruit/train_list.txt", "w") as f:
21
        for filename in part1_files:
22
            abs_path = os.path.join(folder_path, filename)
23
            f.write(abs_path + "\n")
24

25
    with open("./data/fruit/val_list.txt", "w") as f:
26
        for filename in part2_files:
27
            abs_path = os.path.join(folder_path, filename)
28
            f.write(abs_path + "\n")
29

30
    print("Splitting files completed.")
31

32
if __name__ == "__main__":
33
    folder_path = r"W:\MingYi\programming project\Python\yolov7-main\data\fruit\images\potato1"  # 替换为实际的文件夹路径
34
    ratio = 0.9  # 指定训练部分的文件数量比例
35

36
    split_files(folder_path, ratio)

1
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
2
train: ./data/fruit/train_list.txt
3
val: ./data/fruit/val_list.txt
4
# number of classes
5
nc: 1
6

7
# class names
8
names: ['potato']

1
nc---有几种类别
2
names---类名

1
import cv2
2
import tensorrt as trt
3
import torch
4
import numpy as np
5
from collections import OrderedDict, namedtuple
6

7

8
class TRT_engine():
9
    def __init__(self, weight) -> None:
10
        self.imgsz = [640, 640]
11
        self.weight = weight
12
        self.device = torch.device('cuda:0')
13
        self.init_engine()
14

15
    def init_engine(self):
16
        # Infer TensorRT Engine
17
        self.Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
18
        self.logger = trt.Logger(trt.Logger.INFO)
19
        trt.init_libnvinfer_plugins(self.logger, namespace="")
20
        with open(self.weight, 'rb') as self.f, trt.Runtime(self.logger) as self.runtime:
21
            self.model = self.runtime.deserialize_cuda_engine(self.f.read())
22
        self.bindings = OrderedDict()
23
        self.fp16 = False
24

25
        print("Available TensorRT Bindings:")
26
        for index in range(self.model.num_bindings):
27
            self.name = self.model.get_tensor_name(index)  # 使用新 API
28
            self.dtype = trt.nptype(self.model.get_binding_dtype(index))
29
            self.shape = tuple(self.model.get_binding_shape(index))
30
            self.data = torch.from_numpy(np.empty(self.shape, dtype=np.dtype(self.dtype))).to(self.device)
31
            self.bindings[self.name] = self.Binding(self.name, self.dtype, self.shape, self.data,
32
                                                    int(self.data.data_ptr()))
33
            if self.model.binding_is_input(index) and self.dtype == np.float16:
34
                self.fp16 = True
35
            print(f"  {index}: {self.name} -> shape={self.shape}, dtype={self.dtype}")
36

37
        self.binding_addrs = OrderedDict((n, d.ptr) for n, d in self.bindings.items())
38
        self.context = self.model.create_execution_context()
39

40
    def preprocess(self, image):
41
        # 检查图像是否为空
42
        if image is None:
43
            raise ValueError("图像加载失败，请检查路径是否正确。")
44

45
        img = cv2.resize(image, (640, 640))
46
        img = img.transpose((2, 0, 1))
47
        img = np.expand_dims(img, 0)
48
        img = np.ascontiguousarray(img)
49
        img = torch.from_numpy(img).to(self.device)
50
        img = img.float()
51
        return img
52

53
    def predict(self, img, threshold):
54
        img = self.preprocess(img)
55
        self.binding_addrs['images'] = int(img.data_ptr())
56
        self.context.execute_v2(list(self.binding_addrs.values()))
57

58
        # 获取输出结果
59
        output = self.bindings['output'].data.cpu().numpy()  # 转为 NumPy 数组
60

61
        # 解析 YOLOv7 输出
62
        batch, num_anchors, grid_h, grid_w, num_attrs = output.shape  # (1, 3, 80, 80, 6)
63
        detections = output.reshape(-1, 6)  # 转成 (N, 6) 形式
64

65
        new_bboxes = []
66
        for det in detections:
67
            x, y, w, h, conf, cls = det
68

69
            # 如果存在 NaN 或者置信度低于阈值，则跳过该框
70
            if np.any(np.isnan(det)) or conf < threshold:
71
                continue
72

73
            # 转换 bbox 形式 (cx, cy, w, h) -> (xmin, ymin, xmax, ymax)
74
            xmin = (x - w / 2)
75
            ymin = (y - h / 2)
76
            xmax = (x + w / 2)
77
            ymax = (y + h / 2)
78

79
            new_bboxes.append([cls, conf, xmin, ymin, xmax, ymax])
80

81
        return new_bboxes
82

83

84
def visualize(img, bbox_array):
85
    for temp in bbox_array:
86
        xmin, ymin, xmax, ymax = map(int, temp[2:6])
87

88
        # 检查坐标是否合法
89
        if xmin < 0 or ymin < 0 or xmax < 0 or ymax < 0:
90
            continue  # 跳过无效的框
91

92
        clas = int(temp[0])
93
        score = temp[1]
94
        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (105, 237, 249), 2)
95
        img = cv2.putText(img, f"class:{clas} {round(score, 2)}", (xmin, ymin - 5),
96
                          cv2.FONT_HERSHEY_SIMPLEX, 0.5, (105, 237, 249), 1)
97
    return img
98

99

100
def process_image(img_path, trt_engine):
101
    img = cv2.imread(img_path)
102
    if img is None:
103
        print(f"无法读取图像文件: {img_path}")
104
        return
105
    results = trt_engine.predict(img, threshold=0.5)
106
    img = visualize(img, results)
107
    cv2.imshow("img", img)
108
    cv2.waitKey(0)
109
    cv2.destroyAllWindows()
110

111

112
# 用户输入选择
113
print("选择模式：输入 1 预测单张图片，输入 2 进行摄像头实时检测: ")
114
mode = input()
115

116
trt_engine = TRT_engine("./best_fp16.engine")
117

118
if mode == "1":
119
    img_path = input("请输入图片路径: ")
120
    process_image(img_path, trt_engine)
121
elif mode == "2":
122
    cap = cv2.VideoCapture(0)
123
    while True:
124
        ret, frame = cap.read()
125
        if not ret:
126
            break
127

128
        results = trt_engine.predict(frame, threshold=0.5)
129
        frame = visualize(frame, results)
130

131
        cv2.imshow("Camera", frame)
132

133
        if cv2.waitKey(1) & 0xFF == ord('q'):
134
            break
135

136
    cap.release()
137
    cv2.destroyAllWindows()

1
# 关键参数示例（简化版）
2
parser.add_argument('--weights', type=str, default='weights/yolov7.pt', help='预训练权重路径')
3
parser.add_argument('--cfg', type=str, default='models/yolov7.yaml', help='模型配置文件')
4
parser.add_argument('--data', type=str, default='data/coco.yaml', help='数据集配置文件')
5
parser.add_argument('--hyp', type=str, default='data/hyp.scratch.p5.yaml', help='超参数文件')
6
parser.add_argument('--epochs', type=int, default=300, help='训练总轮次')
7
parser.add_argument('--batch-size', type=int, default=32, help='批量大小')
8
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[训练, 测试] 图像尺寸')
9
parser.add_argument('--device', default='0', help='GPU设备编号')

1
import torch
2
torch.cuda.is_available()