• YOLOv7训练自己的数据集

    目录

    1、制作YOLO格式数据集

    1.1、数据集

    1.2、如何转换为YOLOv7所需的格式?

    1.3、如何批量化生成YOLO格式的txt标注

    1.4、如何划分YOLO的train、val和test

    2、使用YOLOv７训练自己的模型

    2.1、测试预训练的yolov7.pt

    （1）测试图片

    （2）测试本地摄像头

    （3）测试视频流效果

    2.2、训练自己数据的YOLOv7模型

    2.3、测试自己训练的模型

    2.4、测试关键点检测

    YOLOv7下载地址：YOLOv7: Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors

    1、制作YOLO格式数据集

    1.1、数据集

    本文采用的是EDS数据集：包含了来自 3 台不同 X 光机器的 14219 张图片, 其中 10 类物品, 共计 31655 个目标实例，均由专业标注人员进行标注。

    每一台机器对应一个数据集，分别对应domain1、domain2和domain3，下图对应每个数据集中的类别分布且相对均匀。

    代码显示部分图像：

    1. import matplotlib.pyplot as plt
    2. import glob
    3. import cv2
    4.  
    5. def show_multi_img(imgpath,num):
    6.     """
    7.     :param imgpath: 图像地址
    8.     :param num: 输出图像的数量：eg:6*6,一幅图展示36张
    9.     :return:
    10.     """
    11.     img_path = glob.glob(imgpath+"/*")
    12.     plt.figure()
    13.     for i in range(1,num*num+1):
    14.         img = cv2.imread(img_path[i])
    15.         title = img_path[i].split("\\")[1]
    16.         plt.subplot(num,num,i)
    17.         plt.imshow(img)
    18.         plt.title(title,fontsize=6)
    19.         plt.xticks([])
    20.         plt.yticks([])
    21.         plt.axis("on")
    22.         plt.savefig("final.png")
    23.     plt.show()
    24.  
    25. if __name__ == "__main__":
    26.     image_dir = "./domain2/image"
    27.     show_multi_img(image_dir,6)

    每个domain分别由image和txt组成：

     1.2、如何转换为YOLOv7所需的格式?

    首先来看一下yolo数据的标注：

     EDS数据集格式：

     假设图像的高度和宽度分别为H和W，bbox的左上角坐标为（xmin,ymin)，右下角坐标为（xmax,ymax），则中心点（x_center,y_center），即

    x_center = xmin + (xmax - xmin)/2

    y_center = ymin + (ymax - ymin)/2

    W = xmax - xmin

    H = ymax - ymin

    则YOLO数据格式为：label, x_, y_, w_, h_，则有对应关系：

    x_ = x_center / img_width

    y_ = y_center / img_height

    w_ = W / img_width

    h_ = H / img_height

    其中label对应的是数字，需要将EDS中的类名转换为数字表示，img_width和img_height为图像的原始的宽度和高度，可以通过cv2.imread()读取，然后shape获取宽度和高度

    1. img= cv2.imread("./domain/image/00001.jpg")
    2. img_height,img_width,_ = img.shape

    显示一幅图像并将bbox绘制在原图中：

    1. import cv2
    2. f = open("./domain1/txt/00004.txt",encoding="utf-8")
    3. img = cv2.imread('./domain1/image/00004.jpg')
    4. img_height,img_width,_ = img.shape
    5. for line in f.readlines():
    6.     text = str(line.split(" ")[1])
    7.     xmin = float(line.split(" ")[2])
    8.     ymin = float(line.split(" ")[3])
    9.     xmax = float(line.split(" ")[4])
    10.     ymax = float(line.split(" ")[5])
    11.     print("xmin:{},xmax:{},ymin:{},ymax:{}".format(xmin,xmax,ymin,ymax))
    12.     x_center = xmin + (xmax - xmin) / 2
    13.     y_center = ymin + (ymax - ymin) / 2
    14.     w = xmax - xmin
    15.     h = ymax - ymin
    16.     # 保留6位小数
    17.     x_center = round(x_center / img_width, 6)
    18.     y_center = round(y_center / img_height, 6)
    19.     w = round(w / img_width, 6)
    20.     h = round(h / img_height, 6)
    21.     # print(x_center,y_center,w,h)
    22.     # 将yolo格式转换原始的格式进行验证
    23.     x1 = int((float(x_center)-float(w)/2)*img_width)
    24.     y1 = int((float(y_center) - float(h) / 2) * img_height)
    25.     x2 = int((float(x_center) + float(w) / 2) * img_width)
    26.     y2 = int((float(y_center) + float(h) / 2) * img_height)
    27.     print(x1,y1,x2,y2)
    28.     cv2.rectangle(img,(x1,y1),(x2,y2),(0,255,255),3)
    29.     cv2.putText(img,text,(int(xmin),int(ymin)-5),cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
    30.     cv2.imshow("show",img)
    31.     cv2.waitKey(0)
    32. cv2.imwrite("bbox.png",img)

    转换前：xmin:84.0,ymin:369.0,xmax:342.0,ymax:554.0
    转换后：xmin:83,ymin:368,xmax:341,ymax:553
    转换前：xmin:210.0,ymin:409.0,xmax:591.0,ymax:691.0
    转换后：xmin:210,ymin:409,xmax:591,ymax:691
    转换前：xmin:182.0,ymin:457.0,xmax:364.0,ymax:550.0
    转换后：xmin:181,ymin:456,xmax:364,ymax:549

    -------------------------------------------------------------------------------------

    这里还是存在一些转换的误差，不过影响没那么大。

    注意：如果没有的标注数据可以用，可以下载LabelImg，进行YOLO格式的数据集标注，直接生存对应的yolo格式的数据集。

    1.3、如何批量化生成YOLO格式的txt标注

    1. import glob
    2. import os
    3. import cv2
    4.  
    5. txt_file = r".\domain1\txt"
    6. name = glob.glob(os.path.join(txt_file,"*.txt"))
    7. list_1 = []
    8. for i in name:
    9.     f = open(i,encoding="utf-8")
    10.     byt = f.readlines()
    11.     for line in byt:
    12.         list_1.append(line.split(" ")[1])
    13.         x = line.split(" ")[2]
    14.         y = line.split(" ")[3]
    15.         w = line.split(" ")[4]
    16.         h = line.split(" ")[5]
    17.         # print(x,y,w,h)
    18. # 读取所有txt中的目标，并去重
    19. list2 = list(set(list_1))
    20. # print(list2)
    21. l = {} # EDS数据类名对应的数字
    22. j = 0
    23. for i in list2:
    24.     l[i] = j
    25.     j += 1
    26. print(l) #对应的字典形式
    27.  
    28. # yolov7的第一列是cls_id x y w h 其中坐标(x,y)是中心点坐标，并且是相对于图片宽高的比例值 ，并非绝对坐标
    29. img_path = "./domain1/image"
    30. out_path = "./out"
    31. list_1 = []
    32. name = glob.glob(os.path.join(txt_file,"*.txt"))
    33. for i in name:
    34.     if not os.path.exists(out_path):
    35.         os.mkdir(out_path)
    36.     with open(os.path.join(out_path, i.split("\\")[3].split(".")[0] + ".txt"), "w") as f_1:
    37.         img_name = i.split("\\")[3].split(".")[0] + ".jpg"
    38.         img = os.path.join(img_path,img_name)
    39.         img_ = cv2.imread(img)
    40.         img_height, img_width, _ = img_.shape
    41.         f = open(i,encoding="utf-8")
    42.         byt = f.readlines()
    43.         for line in byt:
    44.             class_num = l[line.split(" ")[1]]
    45.             xmin = float(line.split(" ")[2])
    46.             ymin = float(line.split(" ")[3])
    47.             xmax = float(line.split(" ")[4])
    48.             ymax = float(line.split(" ")[5])
    49.             x_center = xmin + (xmax - xmin) / 2
    50.             y_center = ymin + (ymax - ymin) / 2
    51.             w = xmax - xmin
    52.             h = ymax - ymin
    53.             x_center = round(x_center / img_width, 6)
    54.             y_center = round(y_center / img_height, 6)
    55.             w = round(w / img_width, 6)
    56.             h = round(h / img_height, 6)
    57.             info = [str(i) for i in [class_num, x_center, y_center, w, h]]
    58.             print(info)
    59.             f_1.write(" ".join(info)+"\n")

    1.4、如何划分YOLO的train、val和test

    本文制作好的数据集：YOLO格式的EDS数据集，免费欢迎下载！感谢支持！

    1. # 将图片和标注数据按比例切分为 训练集和测试集
    2. import shutil
    3. import random
    4. import os
    5.  
    6. # 原始路径，需要修改
    7. image_original_path = './domain1/image/'
    8. label_original_path = './out/'
    9. # 训练集路径，需要修改
    10. train_image_path = 'E:\yolov7\data\images\\train\\'
    11. train_label_path = 'E:\yolov7\data\labels\\train\\'
    12. # 验证集路径，需要修改
    13. val_image_path = 'E:\yolov7\data\images\\val\\'
    14. val_label_path = 'E:\yolov7\data\labels\\val\\'
    15. # 测试集路径，需要修改
    16. test_image_path = 'E:\yolov7\data\images\\test\\'
    17. test_label_path = 'E:\yolov7\data\labels\\test\\'
    18.  
    19. # 数据集划分比例，训练集75%，验证集15%，测试集15%，按需修改
    20. train_percent = 0.7
    21. val_percent = 0.15
    22. test_percent = 0.1
    23.  
    24. # 检查文件夹是否存在
    25. def mkdir():
    26.     if not os.path.exists(train_image_path):
    27.         os.makedirs(train_image_path)
    28.     if not os.path.exists(train_label_path):
    29.         os.makedirs(train_label_path)
    30.  
    31.     if not os.path.exists(val_image_path):
    32.         os.makedirs(val_image_path)
    33.     if not os.path.exists(val_label_path):
    34.         os.makedirs(val_label_path)
    35.  
    36.     if not os.path.exists(test_image_path):
    37.         os.makedirs(test_image_path)
    38.     if not os.path.exists(test_label_path):
    39.         os.makedirs(test_label_path)
    40.  
    41. def main():
    42.     mkdir()
    43.     total_txt = os.listdir(label_original_path)
    44.     num_txt = len(total_txt)
    45.     list_all_txt = range(num_txt)  # 范围 range(0, num)
    46.  
    47.     num_train = int(num_txt * train_percent)
    48.     num_val = int(num_txt * val_percent)
    49.     num_test = num_txt - num_train - num_val
    50.  
    51.     train = random.sample(list_all_txt, num_train)
    52.     # train从list_all_txt取出num_train个元素
    53.     # 所以list_all_txt列表只剩下了这些元素：val_test
    54.     val_test = [i for i in list_all_txt if not i in train]
    55.     # 再从val_test取出num_val个元素，val_test剩下的元素就是test
    56.     val = random.sample(val_test, num_val)
    57.  
    58.     print("训练集数目：{}, 验证集数目：{},测试集数目：{}".format(len(train), len(val), len(val_test) - len(val)))
    59.     for i in list_all_txt:
    60.         name = total_txt[i][:-4]
    61.  
    62.         srcImage = image_original_path + name + '.jpg'
    63.         srcLabel = label_original_path + name + '.txt'
    64.  
    65.         if i in train:
    66.             dst_train_Image = train_image_path + name + '.jpg'
    67.             dst_train_Label = train_label_path + name + '.txt'
    68.             shutil.copyfile(srcImage, dst_train_Image)
    69.             shutil.copyfile(srcLabel, dst_train_Label)
    70.         elif i in val:
    71.             dst_val_Image = val_image_path + name + '.jpg'
    72.             dst_val_Label = val_label_path + name + '.txt'
    73.             shutil.copyfile(srcImage, dst_val_Image)
    74.             shutil.copyfile(srcLabel, dst_val_Label)
    75.         else:
    76.             dst_test_Image = test_image_path + name + '.jpg'
    77.             dst_test_Label = test_label_path + name + '.txt'
    78.             shutil.copyfile(srcImage, dst_test_Image)
    79.             shutil.copyfile(srcLabel, dst_test_Label)
    80.  
    81.  
    82. if __name__ == '__main__':
    83.     main()
    84.

    2、使用YOLOv７训练自己的模型

    官方地址：https://github.com/wongkinyiu/yolov7

    采用git拉取：

    git clone https://github.com/wongkinyiu/yolov7

    2.1、测试预训练的yolov7.pt

    官网提供了下载链接，可以直接下载，或者直接从csdn里下载：YOLOv7预训练权重

    预训练权重下载完成后，打开detect.py

    直接运行即可，其他都选择默认的参数！

    （1）测试图片

    或者修改--source为自己的图像路径，同样也可以修改--weights=your_weight_path，测试自己训练的模型

    测试yolov7.pt的识别效果 

     

    看下官方提供处理图像的代码：utils.datasets

    1. class LoadImages:  # for inference
    2.     def __init__(self, path, img_size=640, stride=32):
    3.         """
    4.             path:图像路径
    5.             img_size:最终要测试的图像尺寸
    6.             stride:这个主要用于pad一些小的图像以满足实际测试图像的尺寸
    7.             return：
    8.             path：图像的路径
    9.             img：resize后的图像
    10.             img0：原始图像
    11.             self.cap
    12.         """
    13.         # 遍历输入的测试图像路径，files保存测试的地址
    14.         p = str(Path(path).absolute())  # os-agnostic absolute path
    15.         if '*' in p:
    16.             files = sorted(glob.glob(p, recursive=True))  # glob
    17.         elif os.path.isdir(p):
    18.             files = sorted(glob.glob(os.path.join(p, '*.*')))  # dir
    19.         elif os.path.isfile(p):
    20.             files = [p]  # files
    21.         else:
    22.             raise Exception(f'ERROR: {p} does not exist')
    23.  
    24.         # 通过图像地址的后缀判断是图像还是视频，然后用list格式保存
    25.         images = [x for x in files if x.split('.')[-1].lower() in img_formats]
    26.         videos = [x for x in files if x.split('.')[-1].lower() in vid_formats]
    27.         # 文件里总共有多少和多少视频
    28.         ni, nv = len(images), len(videos)
    29.  
    30.         self.img_size = img_size
    31.         self.stride = stride
    32.         self.files = images + videos # list格式
    33.         self.nf = ni + nv  # number of files
    34.         self.video_flag = [False] * ni + [True] * nv # 用于判断是不是视频
    35.         self.mode = 'image'
    36.         if any(videos): # 判断videos是否存在
    37.             self.new_video(videos[0])  # new video
    38.         else:
    39.             self.cap = None
    40.         assert self.nf > 0, f'No images or videos found in {p}. ' \
    41.                             f'Supported formats are:\nimages: {img_formats}\nvideos: {vid_formats}'
    42.     
    43.     # __iter__迭代器，系统定义的名字
    44.     def __iter__(self):
    45.         self.count = 0
    46.         return self
    47.  
    48.     def __next__(self):
    49.         if self.count == self.nf:
    50.             raise StopIteration
    51.         path = self.files[self.count]
    52.  
    53.         if self.video_flag[self.count]:
    54.             # Read video
    55.             self.mode = 'video'
    56.             ret_val, img0 = self.cap.read()
    57.             if not ret_val:
    58.                 self.count += 1
    59.                 self.cap.release()
    60.                 if self.count == self.nf:  # last video
    61.                     raise StopIteration
    62.                 else:
    63.                     path = self.files[self.count]
    64.                     self.new_video(path)
    65.                     ret_val, img0 = self.cap.read()
    66.  
    67.             self.frame += 1
    68.             print(f'video {self.count + 1}/{self.nf} ({self.frame}/{self.nframes}) {path}: ', end='')
    69.  
    70.         else:
    71.             # Read image
    72.             self.count += 1
    73.             img0 = cv2.imread(path)  # BGR
    74.             assert img0 is not None, 'Image Not Found ' + path
    75.             #print(f'image {self.count}/{self.nf} {path}: ', end='')
    76.  
    77.         # Padded resize
    78.         img = letterbox(img0, self.img_size, stride=self.stride)[0]
    79.  
    80.         # Convert
    81.         img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
    82.         img = np.ascontiguousarray(img)
    83.  
    84.         return path, img, img0, self.cap
    85.  
    86.     def new_video(self, path):
    87.         self.frame = 0
    88.         self.cap = cv2.VideoCapture(path)
    89.         self.nframes = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
    90.  
    91.     def __len__(self):
    92.         return self.nf  # number of files

    （2）测试本地摄像头

    简单的一个获取本地摄像头的代码

    1. import cv2
    2.  
    3. def access_camera(url,output_path):
    4.     if url == int(0):
    5.         cap = cv2.VideoCapture(url)
    6.     else:
    7.         cap = cv2.VideoCapture(url)
    8.     while(cap.isOpened()):
    9.         # Capture frame-by-frame
    10.         ret, frame = cap.read()
    11.         # Display the resulting frame
    12.         cv2.imshow('frame',frame)
    13.         cv2.imwrite(output_path,frame)
    14.         print("图像保存成功！")
    15.         if cv2.waitKey(1) & 0xFF == ord('q'):
    16.             break
    17.     # When everything done, release the capture
    18.     cap.release()
    19.     cv2.destroyAllWindows()
    20.  
    21. if __name__ == "__main__":
    22.     url = 'http://admin:admin@192.168.1.3:8081/video'  # 调用IP摄像机
    23.     output_path = "./runs/detect/img.png"
    24.     # url = 0 # 调用笔记本摄像头
    25.     access_camera(url,output_path)

    yolov7提供的代码，其实思路是一样的

    1. class LoadWebcam:  # for inference
    2.     def __init__(self, pipe='0', img_size=640, stride=32):
    3.         """
    4.             pipe:0表示使用本地摄像头
    5.             img_size:图像大小
    6.             stride:
    7.         """
    8.         self.img_size = img_size
    9.         self.stride = stride
    10.  
    11.         if pipe.isnumeric():
    12.             pipe = eval(pipe)  # local camera
    13.         # pipe = 'rtsp://192.168.1.64/1'  # IP camera
    14.         # pipe = 'rtsp://username:password@192.168.1.64/1'  # IP camera with login
    15.         # pipe = 'http://wmccpinetop.axiscam.net/mjpg/video.mjpg'  # IP golf camera
    16.  
    17.         self.pipe = pipe
    18.         self.cap = cv2.VideoCapture(pipe)  # video capture object
    19.         self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 3)  # set buffer size
    20.  
    21.     def __iter__(self):
    22.         self.count = -1
    23.         return self
    24.  
    25.     def __next__(self):
    26.         self.count += 1
    27.         if cv2.waitKey(1) == ord('q'):  # q to quit
    28.             self.cap.release()
    29.             cv2.destroyAllWindows()
    30.             raise StopIteration
    31.  
    32.         # Read frame
    33.         if self.pipe == 0:  # local camera
    34.             ret_val, img0 = self.cap.read()
    35.             img0 = cv2.flip(img0, 1)  # flip left-right
    36.         else:  # IP camera
    37.             n = 0
    38.             while True:
    39.                 n += 1
    40.                 self.cap.grab()
    41.                 if n % 30 == 0:  # skip frames
    42.                     ret_val, img0 = self.cap.retrieve()
    43.                     if ret_val:
    44.                         break
    45.  
    46.         # Print
    47.         assert ret_val, f'Camera Error {self.pipe}'
    48.         img_path = 'webcam.jpg'
    49.         print(f'webcam {self.count}: ', end='')
    50.  
    51.         # Padded resize
    52.         img = letterbox(img0, self.img_size, stride=self.stride)[0]
    53.  
    54.         # Convert
    55.         img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
    56.         img = np.ascontiguousarray(img)
    57.  
    58.         return img_path, img, img0, None
    59.  
    60.     def __len__(self):
    61.         return 0

    （3）测试视频流效果

    1. class LoadStreams:  # multiple IP or RTSP cameras
    2.     def __init__(self, sources='streams.txt', img_size=640, stride=32):
    3.         self.mode = 'stream'
    4.         self.img_size = img_size
    5.         self.stride = stride
    6.  
    7.         if os.path.isfile(sources):
    8.             with open(sources, 'r') as f:
    9.                 sources = [x.strip() for x in f.read().strip().splitlines() if len(x.strip())]
    10.         else:
    11.             sources = [sources]
    12.  
    13.         n = len(sources)
    14.         self.imgs = [None] * n
    15.         self.sources = [clean_str(x) for x in sources]  # clean source names for later
    16.         for i, s in enumerate(sources):
    17.             # Start the thread to read frames from the video stream
    18.             print(f'{i + 1}/{n}: {s}... ', end='')
    19.             url = eval(s) if s.isnumeric() else s
    20.             if 'youtube.com/' in str(url) or 'youtu.be/' in str(url):  # if source is YouTube video
    21.                 check_requirements(('pafy', 'youtube_dl'))
    22.                 import pafy
    23.                 url = pafy.new(url).getbest(preftype="mp4").url
    24.             cap = cv2.VideoCapture(url)
    25.             assert cap.isOpened(), f'Failed to open {s}'
    26.             w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    27.             h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    28.             self.fps = cap.get(cv2.CAP_PROP_FPS) % 100
    29.  
    30.             _, self.imgs[i] = cap.read()  # guarantee first frame
    31.             thread = Thread(target=self.update, args=([i, cap]), daemon=True)
    32.             print(f' success ({w}x{h} at {self.fps:.2f} FPS).')
    33.             thread.start()
    34.         print('')  # newline
    35.  
    36.         # check for common shapes
    37.         s = np.stack([letterbox(x, self.img_size, stride=self.stride)[0].shape for x in self.imgs], 0)  # shapes
    38.         self.rect = np.unique(s, axis=0).shape[0] == 1  # rect inference if all shapes equal
    39.         if not self.rect:
    40.             print('WARNING: Different stream shapes detected. For optimal performance supply similarly-shaped streams.')
    41.  
    42.     def update(self, index, cap):
    43.         # Read next stream frame in a daemon thread
    44.         n = 0
    45.         while cap.isOpened():
    46.             n += 1
    47.             # _, self.imgs[index] = cap.read()
    48.             cap.grab()
    49.             if n == 4:  # read every 4th frame
    50.                 success, im = cap.retrieve()
    51.                 self.imgs[index] = im if success else self.imgs[index] * 0
    52.                 n = 0
    53.             time.sleep(1 / self.fps)  # wait time
    54.  
    55.     def __iter__(self):
    56.         self.count = -1
    57.         return self
    58.  
    59.     def __next__(self):
    60.         self.count += 1
    61.         img0 = self.imgs.copy()
    62.         if cv2.waitKey(1) == ord('q'):  # q to quit
    63.             cv2.destroyAllWindows()
    64.             raise StopIteration
    65.  
    66.         # Letterbox
    67.         img = [letterbox(x, self.img_size, auto=self.rect, stride=self.stride)[0] for x in img0]
    68.  
    69.         # Stack
    70.         img = np.stack(img, 0)
    71.  
    72.         # Convert
    73.         img = img[:, :, :, ::-1].transpose(0, 3, 1, 2)  # BGR to RGB, to bsx3x416x416
    74.         img = np.ascontiguousarray(img)
    75.  
    76.         return self.sources, img, img0, None
    77.  
    78.     def __len__(self):
    79.         return 0  # 1E12 frames = 32 streams at 30 FPS for 30 years

    获取到所有的图像或者视频流，然后将获取的图像输入对应的model中，查看一下官方提供的detect.py代码

    1. def detect(save_img=False):
    2.     source, weights, view_img, save_txt, imgsz, trace = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size, not opt.no_trace
    3.     save_img = not opt.nosave and not source.endswith('.txt')  # save inference images
    4.     webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
    5.         ('rtsp://', 'rtmp://', 'http://', 'https://'))
    6.  
    7.     # Directories
    8.     save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))  # increment run
    9.     (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
    10.  
    11.     # Initialize
    12.     set_logging()
    13.     device = select_device(opt.device)
    14.     half = device.type != 'cpu'  # half precision only supported on CUDA
    15.  
    16.     # Load model
    17.     model = attempt_load(weights, map_location=device)  # load FP32 model
    18.     stride = int(model.stride.max())  # model stride
    19.     imgsz = check_img_size(imgsz, s=stride)  # check img_size
    20.  
    21.     if trace:
    22.         model = TracedModel(model, device, opt.img_size)
    23.  
    24.     if half:
    25.         model.half()  # to FP16
    26.  
    27.     # Second-stage classifier
    28.     classify = False
    29.     if classify:
    30.         modelc = load_classifier(name='resnet101', n=2)  # initialize
    31.         modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()
    32.  
    33.     # Set Dataloader
    34.     vid_path, vid_writer = None, None
    35.     if webcam:
    36.         view_img = check_imshow()
    37.         cudnn.benchmark = True  # set True to speed up constant image size inference
    38.         dataset = LoadStreams(source, img_size=imgsz, stride=stride)
    39.     else:
    40.         dataset = LoadImages(source, img_size=imgsz, stride=stride)
    41.  
    42.     # Get names and colors
    43.     names = model.module.names if hasattr(model, 'module') else model.names
    44.     colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
    45.  
    46.     # Run inference
    47.     if device.type != 'cpu':
    48.         model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
    49.     t0 = time.time()
    50.     for path, img, im0s, vid_cap in dataset:
    51.         img = torch.from_numpy(img).to(device)
    52.         img = img.half() if half else img.float()  # uint8 to fp16/32
    53.         img /= 255.0  # 0 - 255 to 0.0 - 1.0
    54.         if img.ndimension() == 3:
    55.             img = img.unsqueeze(0)
    56.  
    57.         # Inference
    58.         t1 = time_synchronized()
    59.         pred = model(img, augment=opt.augment)[0]
    60.  
    61.         # Apply NMS
    62.         pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
    63.         t2 = time_synchronized()
    64.  
    65.         # Apply Classifier
    66.         if classify:
    67.             pred = apply_classifier(pred, modelc, img, im0s)
    68.  
    69.         # Process detections
    70.         for i, det in enumerate(pred):  # detections per image
    71.             if webcam:  # batch_size >= 1
    72.                 p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
    73.             else:
    74.                 p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
    75.  
    76.             p = Path(p)  # to Path
    77.             save_path = str(save_dir / p.name)  # img.jpg
    78.             txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # img.txt
    79.             s += '%gx%g ' % img.shape[2:]  # print string
    80.             gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
    81.             if len(det):
    82.                 # Rescale boxes from img_size to im0 size
    83.                 det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
    84.  
    85.                 # Print results
    86.                 for c in det[:, -1].unique():
    87.                     n = (det[:, -1] == c).sum()  # detections per class
    88.                     s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
    89.  
    90.                 # Write results
    91.                 for *xyxy, conf, cls in reversed(det):
    92.                     if save_txt:  # Write to file
    93.                         xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
    94.                         line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh)  # label format
    95.                         with open(txt_path + '.txt', 'a') as f:
    96.                             f.write(('%g ' * len(line)).rstrip() % line + '\n')
    97.  
    98.                     if save_img or view_img:  # Add bbox to image
    99.                         label = f'{names[int(cls)]} {conf:.2f}'
    100.                         plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
    101.  
    102.             # Print time (inference + NMS)
    103.             #print(f'{s}Done. ({t2 - t1:.3f}s)')
    104.  
    105.             # Stream results
    106.             if view_img:
    107.                 cv2.imshow(str(p), im0)
    108.                 cv2.waitKey(1)  # 1 millisecond
    109.  
    110.             # Save results (image with detections)
    111.             if save_img:
    112.                 if dataset.mode == 'image':
    113.                     cv2.imwrite(save_path, im0)
    114.                     print(f" The image with the result is saved in: {save_path}")
    115.                 else:  # 'video' or 'stream'
    116.                     if vid_path != save_path:  # new video
    117.                         vid_path = save_path
    118.                         if isinstance(vid_writer, cv2.VideoWriter):
    119.                             vid_writer.release()  # release previous video writer
    120.                         if vid_cap:  # video
    121.                             fps = vid_cap.get(cv2.CAP_PROP_FPS)
    122.                             w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    123.                             h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    124.                         else:  # stream
    125.                             fps, w, h = 30, im0.shape[1], im0.shape[0]
    126.                             save_path += '.mp4'
    127.                         vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
    128.                     vid_writer.write(im0)
    129.  
    130.     if save_txt or save_img:
    131.         s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
    132.         #print(f"Results saved to {save_dir}{s}")
    133.  
    134.     print(f'Done. ({time.time() - t0:.3f}s)')

    2.2、训练自己数据的YOLOv7模型

    根据前面的方法制作自己数据集，放在yolov7/data目录下

     在yolov7/data目录下创建一个dataset.yaml文件，可以参考官方提供的coco.yaml

     自己的配置文件dataset.yaml

    1. train: E:/yolov7/data/images/train # train images
    2. val: E:/yolov7/data/images/val # val images
    3. test: E:/yolov7/data/images/test # test images (optional)
    4.  
    5. # Classes
    6. nc: 10  # number of classes
    7. names: ['laptop','pressure','device','plasticbottle','scissor','knife','lighter','powerbank','glassbottle','umbrella']  # class names

    开始训练。。。漫长的等待了，最终所有的训练信息都保存在yolov7/runs/train/exp目录下

    2.3、测试自己训练的模型

    修改detect.py中的weights地址 ，这个模型我只训练了10次，效果也还行。

    2.4、测试关键点检测

    首先下载官方提供的预训练模型yolov7-w6-pose.pt

    1. import matplotlib
    2. """
    3.     未使用matplotlib.use('TkAgg')
    4.     出现问题：UserWarning: Matplotlib is currently using agg, which is a non-GUI backend
    5. """
    6. matplotlib.use('TkAgg')
    7. import matplotlib.pyplot as plt
    8. print(matplotlib.get_backend())
    9. import torch
    10. import cv2
    11. from torchvision import transforms
    12. import numpy as np
    13. from utils.datasets import letterbox
    14. # 查看GUI backbend环境，主要是查看每个调用的代码环境下是否都是相同的环境
    15. print(matplotlib.get_backend())
    16. from utils.general import non_max_suppression_kpt
    17. print(matplotlib.get_backend())
    18. from utils.plots import output_to_keypoint, plot_skeleton_kpts
    19. # plots中matplotlib.get_backend()设置不同，记得修改
    20. print(matplotlib.get_backend())
    21.  
    22. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    23. weigths = torch.load('../yolov7-w6-pose.pt')
    24. model = weigths['model']
    25. model = model.half().to(device)
    26. _ = model.eval()
    27.  
    28. image = cv2.imread('../person.jpeg')
    29. image = letterbox(image, 960, stride=64, auto=True)[0]
    30. image_ = image.copy()
    31. image = transforms.ToTensor()(image)
    32. image = torch.tensor(np.array([image.numpy()]))
    33. image = image.to(device)
    34. image = image.half()
    35.  
    36. output, _ = model(image)
    37.  
    38. output = non_max_suppression_kpt(output, 0.25, 0.65, nc=model.yaml['nc'], nkpt=model.yaml['nkpt'], kpt_label=True)
    39. output = output_to_keypoint(output)
    40. nimg = image[0].permute(1, 2, 0) * 255
    41. nimg = nimg.cpu().numpy().astype(np.uint8)
    42. nimg = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)
    43. for idx in range(output.shape[0]):
    44.     plot_skeleton_kpts(nimg, output[idx, 7:].T, 3)
    45.  
    46. plt.figure(figsize=(8,8))
    47. plt.axis('off')
    48. plt.imshow(nimg)
    49. plt.savefig("person_detection.png")
    50. plt.show()
    51.

     报错信息：

     在utils.plots.py中的442-443行中的增加detach()不在进行反向传播即可

     未完待续。。。

  • 相关阅读:
    java计算机毕业设计基于springboo高校学报论文在线投稿系统
    微服务-Eureka
    「SpringCloud」10 Stream消息驱动
    【C++】多态
    面试五 -bind 和 function
    动态规划：11分割等和子集
    [Linux打怪升级之路]-yun安装和gcc的使用
    神经网络算法有哪些模型,神经网络模型应用实例
    我的创作纪念日——创作者2年
    通过 filesystem 的 inode 设计，理解数组与链表
  • 原文地址：https://blog.csdn.net/weixin_43687366/article/details/126166982