一. 环境安装
参考视频
- Pytorch环境安装细节
- pytorch安装:一个单独的环境中,能使用pip就尽量使用pip,实在有问题的情况,例如没有合适的编译好的系统版本的安装包,再使用conda进行安装,不要来回混淆
- CUDA是否要安装:如果只需要训练、简单推理,则无需单独安装cUDA,直接安装pytorch;如果有部署需求,例如导出TensorRT模型,则需要进行CUDA安装
- Pytorch安装注意事项:必须使用官网的命令进行安装,否则安装的是cpu的版本
- 正确使用windows终端:
- 使用cmd,而不是powershell(无法激活环境)。
- 在其他软件,如Pycharm,vscode中也要注意!
- 可能出现的问题
- Arial.ttf字体文件无法下载
- 手动下载,放到对应的位置,windows下的目录是:~/AppData/Roaming/Ultralytics
- 页面文件太小,无法完成操作
- 调整训练参数中的workers,设置为1
- 修改虚拟内存,将环境安装位置所在的盘,设置一个较大的参数
- ‘Upsample’ object has no attribute ‘recompute _scale_factor’
- pytorch版本过高导致,可以选择降版本,1.8.2目前是不会报错的版本
- 如不想降低版本,可以修改pytorch源码,打开报错的unsampling.py,删除
recompute_scale_factor这个参数
设置电脑虚拟内存:
二 .数据集制作
数据转换代码数据格式转化关键根据相应的文件格式解析处class以及bbox
- voc2yolo
def convert_xml_to_txt_format(xml_file_path: str, class_list: list) -> list:"""将单个xml文件转换为txt格式。参数:xml_file_path (str): XML文件的路径。class_list (list): 包含所有类别的列表。返回:list: 包含YOLO格式的标签列表。"""# 根据class_list构建class_to_id_mapclass_to_id_map = {class_name: index for index, class_name in enumerate(class_list)}tree = ET.parse(xml_file_path)root = tree.getroot()width = int(root.find('.//size/width').text)height = int(root.find('.//size/height').text)txt_format = []for obj in root.findall('.//object'):class_name = obj.find('name').textif class_name not in class_to_id_map:continueclass_id = class_to_id_map[class_name]bbox = obj.find('bndbox')xmin = float(bbox.find('xmin').text)ymin = float(bbox.find('ymin').text)xmax = float(bbox.find('xmax').text)ymax = float(bbox.find('ymax').text)x_center = (xmin + xmax) / 2 / widthy_center = (ymin + ymax) / 2 / heightbbox_width = (xmax - xmin) / widthbbox_height = (ymax - ymin) / heighttxt_format.append(f"{class_id} {x_center} {y_center} {bbox_width} {bbox_height}")return txt_formatdef write_yolo_format_labels_to_file(yolo_labels_dir: str, yolo_label_file: str, yolo_content: list):"""将YOLO格式的标签写入到文件。参数:yolo_labels_dir (str): YOLO标签目录的路径。yolo_label_file (str): YOLO标签文件的名称。yolo_content (list): 包含YOLO格式标签的列表。"""yolo_label_path = os.path.join(yolo_labels_dir, yolo_label_file)with open(yolo_label_path, 'w') as f:for line in yolo_content:f.write(line + '\n')def convert_voc2yolo(voc_annotations_dir: str, yolo_labels_dir: str, class_list: list):"""将VOC格式的XML注释转换为YOLO格式的注释。参数:voc_annotations_dir (str): VOC注释目录的路径。yolo_labels_dir (str): YOLO标签目录的路径。class_list (list): 类别名称的list。"""if not os.path.exists(yolo_labels_dir):os.makedirs(yolo_labels_dir)for xml_file in os.listdir(voc_annotations_dir):if xml_file.endswith('.xml'):yolo_label_file = xml_file.replace('.xml', '.txt')xml_file_path = os.path.join(voc_annotations_dir, xml_file)yolo_content = convert_xml_to_txt_format(xml_file_path,class_list)write_yolo_format_labels_to_file(yolo_labels_dir, yolo_label_file, yolo_content)with open(os.path.join(yolo_labels_dir,"classes.txt"),"w",encoding='utf-8') as f:for c in class_list:f.write(c+'\n')
if __name__ == "__main__":dataset_dir = "./dataset"labels_voc_dir = "labels_voc"labels_yolo_dir = "labels_yolo"class_names = ["橘子", "香蕉", "草莓"]convert_voc2yolo(labels_voc_dir,labels_yolo_dir,class_names)
- yolo2voc
def convert_txt_to_xml_format(txt_file_path: str, voc_file_path: str, class_names: list, images_dir: str):"""将单个txt文件转换为xml格式。参数:txt_file_path (str): YOLO标签文件txt的路径。voc_file_path (str): XML文件的路径。class_names (list): 包含类别名称的列表。images_dir (str): 包含图像文件的目录路径。"""with open(txt_file_path, 'r') as file:lines = file.readlines()# 获取txt文件的文件名,然后去掉后缀txt_filename = os.path.basename(txt_file_path)txt_filename_without_ext = os.path.splitext(txt_filename)[0]# 使用glob库在images_dir下搜索与txt文件同名的图片文件pattern = os.path.join(images_dir, f"{txt_filename_without_ext}.*")matching_files = glob.glob(pattern)if not matching_files:raise FileNotFoundError(f"No matching image found for {txt_filename_without_ext} in {images_dir}")img_file = matching_files[0] # 假设找到的第一个匹配文件就是正确的图片# 从与YOLO标签同名的图像文件中读取宽度和高度img = Image.open(img_file)width, height = img.sizeroot = ET.Element("annotation")folder = ET.SubElement(root, "folder")folder.text = os.path.basename(images_dir)filename = ET.SubElement(root, "filename")filename.text = os.path.basename(img_file)path = ET.SubElement(root, "path")path.text = os.path.abspath(img_file)source = ET.SubElement(root, "source")database = ET.SubElement(source, "database")database.text = "Unknown"size = ET.SubElement(root, "size")ET.SubElement(size, "width").text = str(width)ET.SubElement(size, "height").text = str(height)ET.SubElement(size, "depth").text = "3"segmented = ET.SubElement(root, "segmented")segmented.text = "0"for line in lines:values = line.strip().split()if len(values) < 5:continueclass_id, x_center, y_center, bbox_width, bbox_height = map(float, values)object = ET.SubElement(root, "object")name = ET.SubElement(object, "name")name.text = class_names[int(class_id)]pose = ET.SubElement(object, "pose")pose.text = "Unspecified"truncated = ET.SubElement(object, "truncated")truncated.text = "0"difficult = ET.SubElement(object, "difficult")difficult.text = "0"bndbox = ET.SubElement(object, "bndbox")xmin = int((x_center - bbox_width / 2) * width)ymin = int((y_center - bbox_height / 2) * height)xmax = int((x_center + bbox_width / 2) * width)ymax = int((y_center + bbox_height / 2) * height)ET.SubElement(bndbox, "xmin").text = str(xmin)ET.SubElement(bndbox, "ymin").text = str(ymin)ET.SubElement(bndbox, "xmax").text = str(xmax)ET.SubElement(bndbox, "ymax").text = str(ymax)# 使用prettify函数来格式化输出pretty_xml_as_string = prettify(root)with open(voc_file_path, 'w', encoding='utf-8') as f:f.write(pretty_xml_as_string)def prettify(elem):from xml.dom import minidom"""Return a pretty-printed XML string for the Element."""rough_string = ET.tostring(elem, 'utf-8')reparsed = minidom.parseString(rough_string)return reparsed.toprettyxml(indent=" ")def convert_yolo2voc(yolo_labels_dir: str, voc_labels_dir: str, images_dir: str, class_list: list):"""将YOLO格式的标签文件夹转换为VOC格式的XML注释文件夹。参数:yolo_labels_dir (str): YOLO标签文件夹的路径。voc_labels_dir (str): 转换后的VOC XML文件夹的路径。images_dir (str): 包含图像文件的目录路径。class_names (list): 包含类别名称的列表。"""# 确保voc_labels_dir存在if not os.path.exists(voc_labels_dir):os.makedirs(voc_labels_dir)# 遍历YOLO标签文件夹for label_file in os.listdir(yolo_labels_dir):if label_file.endswith(".txt"):yolo_file_path = os.path.join(yolo_labels_dir, label_file)voc_file_path = os.path.join(voc_labels_dir, label_file.replace(".txt", ".xml"))convert_txt_to_xml_format(yolo_file_path, voc_file_path, class_list, images_dir)
if __name__ == "__main__":dataset_dir = "./dataset"labels_voc_dir = "labels_voc"labels_yolo_dir = "labels_yolo"class_names = ["橘子", "香蕉", "草莓"]convert_yolo2voc(labels_yolo_dir,labels_voc_dir,"imgs",class_names)
- Json转Yolo
提示可以采用 data_df = pd.read_json(json_path),然后利用pandas的语法进行提取bbox和class_id
- 划分数据集
def split_dataset(images_origin_dir: str, labels_origin_dir: str, dataset_dir: str, train_ratio: float = 0.8):"""将原始数据集分割为训练集和验证集,并将相应的文件复制到新的文件夹中。参数:images_origin (str): 存放所有图片文件夹路径labels_origin (str): 存放所有标签文件夹路径。dataset_dir (str): 新的数据集根目录。train_ratio (float, optional): 训练集的比例,默认为 0.8。"""images_dir = os.path.join(dataset_dir, "images")labels_dir = os.path.join(dataset_dir, "labels")train_dir = os.path.join(images_dir, "train")val_dir = os.path.join(images_dir, "val")# 创建训练集和验证集的文件夹os.makedirs(train_dir, exist_ok=True)os.makedirs(val_dir, exist_ok=True)os.makedirs(os.path.join(labels_dir, "train"), exist_ok=True)os.makedirs(os.path.join(labels_dir, "val"), exist_ok=True)image_files=[]for file in os.listdir(images_origin_dir):if file.endswith(".jpg") or file.endswith(".png") or file.endswith(".jpeg"):image_files.append(os.path.join(images_origin_dir, file))# 使用 random.sample 函数进行随机划分random.seed(42) # 设置随机种子train_images = random.sample(image_files, k=int(len(image_files) * train_ratio))val_images = [f for f in image_files if f not in train_images]# 复制文件到训练集和验证集文件夹for subset, images in [('train', train_images), ('val', val_images)]:for image_file in images:# 复制图片文件target_image_path = os.path.join(images_dir, subset, os.path.basename(image_file))shutil.copyfile(image_file, target_image_path)# 复制标签文件label_file = os.path.splitext(os.path.basename(image_file))[0] + '.txt'label_path = os.path.join(labels_origin_dir, label_file)if os.path.exists(label_path):target_label_path = os.path.join(labels_dir, subset, label_file)shutil.copyfile(label_path, target_label_path)print(f"训练集大小: {len(train_images)}")print(f"验证集大小: {len(val_images)}")if __name__ == "__main__":dataset_dir = "./dataset"labels_voc_dir = "labels_voc"labels_yolo_dir = "labels_yolo"class_names = ["橘子", "香蕉", "草莓"]split_dataset("imgs",labels_yolo_dir,dataset_dir)
三.模型使用
- YOLOV5 模型调用(本地)
# Model
model = torch.hub.load('.', 'custom', path=r'best.pt',source='local')
.
: 本地的yolov5的根目录
source=“local”
:从本地调用该模型
- YOLOV8 模型调用
# Load the YOLOv8 model
model = YOLO("yolov8n.pt") # Make sure the model file is in the correct path
如果本地不存在yolov8n.pt,则从网上下载
四. 结果处理
.xyxy()
:即为[x_min,y_min,x_max,y_max]
- YOLOV5 参考博客
import torch
model = torch.hub.load('.', 'yolov5s', source='local')
im = r'data\images\bus.jpg' # file, Path, PIL.Image, OpenCV, nparray, list`
results = model(im) # inference
# results.crop() # or .show(), .save(), .crop(), .pandas(), etc.
results是一个Detections对象,主要有以下方法:
# results.crop() # or .show(), .save(), .crop(), .pandas(), etc.
results.xyxy[0] #Tensor类型
results.pandas().xyxy[0] #转为pandas格式
results.print() #控制台显示
frame = results.render()[0] # 将模型的输出绘制回帧上
另外
:想要个性化输出,限定类别
# 将检测结果转换为Pandas DataFrame格式
detections_df = results.pandas().xyxy[0]# 筛选出标签名为'person'的所有检测结果
person_detections = detections_df[detections_df['name'] == 'person'].to_numpy()# 遍历检测到的人
for detection in person_detections:label_name = detection[6] # 获取标签名bbox = detection[:4].astype('int') # 获取边界框坐标并转换为整数# 在图像上绘制边界框cv.rectangle(image_temp, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 0, 255), 3)# 在边界框左上角添加标签名cv.putText(image_temp, label_name, (bbox[0] - 10, bbox[1] - 10), cv.FONT_ITALIC, 1, (0, 255, 0), 2)
- YOLOV8
建议
:ultalytics源码安装,在ultralytics的根目录下pip install -e . -e
指的是本地可编辑,下载该库时不会直接下载到site-packages下,而是以链接的方式链接到源码的根目录下
from ultralytics import YOLO
model=YOLO("yolov8n-seg.pt")
results=model(r"images\bus.jpg")
这里的results是list类型。
for idx, result in enumerate(results):boxes = result.boxes # Boxes object for bounding box outputsmasks = result.masks # Masks object for segmentation masks outputskeypoints = result.keypoints # Keypoints object for pose outputsprobs = result.probs # Probs object for classification outputsobb = result.obb # Oriented boxes object for OBB outputsif len(boxes.cls) == 0:continuexyxy = boxes.xyxy.data.cpu().numpy().round()cls = boxes.cls.data.cpu().numpy().round()conf = boxes.conf.data.cpu().numpy()
另外
:如果想要个性化输出,可以使用 patched_yolo_infer
库(我只验证了YOLOV8,其它的YOLO系列不知道)
import cv2
from ultralytics import YOLO
from patched_yolo_infer import visualize_results_usual_yolo_inference# Load the image
img_path = 'images/bus.jpg'
img = cv2.imread(img_path)img=visualize_results_usual_yolo_inference(img,model,conf=0.4,iou=0.7,show_classes_list=[0], #Whether to perform instance segmentation. Default is False.segment=True, #是否分割thickness=5,show_boxes=False,fill_mask=False, alpha=0.7, #The transparency of filled masks. Default is 0.3.show_class=False,delta_colors=25, #The random seed offset for color variation. Default is 0.inference_extra_args={'retina_masks':True}, #increase the accuracy of the contoursreturn_image_array=True)
附
:Pyme组件
import cv2 as cv
from PIL import ImageTk, Image
class VideoPlayer():def __init__(self,elementName,video_source=0):super().__init__()self.video_source = video_sourceself.vid = cv.VideoCapture(self.video_source)self.label= Fun.GetElement(uiName,elementName) #如果想要替换为tk,这里直接注入label对象self.delay = 15self.paused = False # 暂停状态标志self.update()def update(self):# 检查是否成功读取帧if self.vid.isOpened():if not self.paused:ret, frame = self.vid.read()if ret:frame = self.process_frame(frame)# 转换为Imageself.photo = ImageTk.PhotoImage(image=Image.fromarray(frame))self.label.config(image=self.photo)else:self.vid.release()returnself.label.after(self.delay, self.update)def process_frame(self, frame):#convert BGR to RGBframe = cv.cvtColor(frame, cv.COLOR_BGR2RGB)return framedef pause(self):self.paused = not self.pausedreturn self.pauseddef stop(self):self.vid.release()