一. 环境安装
参考视频
- 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_map
class_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').text
if class_name not in class_to_id_map:
continue
class_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 / width
y_center = (ymin + ymax) / 2 / height
bbox_width = (xmax - xmin) / width
bbox_height = (ymax - ymin) / height
txt_format.append(f"{class_id} {x_center} {y_center} {bbox_width} {bbox_height}")
return txt_format
def 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.size
root = 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:
continue
class_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 outputs
masks = result.masks # Masks object for segmentation masks outputs
keypoints = result.keypoints # Keypoints object for pose outputs
probs = result.probs # Probs object for classification outputs
obb = result.obb # Oriented boxes object for OBB outputs
if len(boxes.cls) == 0:
continue
xyxy = 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 contours
return_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_source
self.vid = cv.VideoCapture(self.video_source)
self.label= Fun.GetElement(uiName,elementName) #如果想要替换为tk,这里直接注入label对象
self.delay = 15
self.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)
# 转换为Image
self.photo = ImageTk.PhotoImage(image=Image.fromarray(frame))
self.label.config(image=self.photo)
else:
self.vid.release()
return
self.label.after(self.delay, self.update)
def process_frame(self, frame):
#convert BGR to RGB
frame = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
return frame
def pause(self):
self.paused = not self.paused
return self.paused
def stop(self):
self.vid.release()