xml to yolo
import copy
from lxml.etree import Element, SubElement, tostring, ElementTree
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import cv2
import torch
# 检测是否有可用的cuda
if torch.cuda.is_available():
device = torch.device('cuda')
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
device = torch.device('cpu')
classes = ['AUV','UWSN node','BUOY','SUB_BUOY','holothurian','echinus','scallop','starfish'] # 修改类别
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
imgpath = 'C:/Users/32566/Desktop/dataset/images' #修改路径
def convert(size, box):
dw = 1. / size[0]
dh = 1. / size[1]
x = (box[0] + box[1]) / 2.0
y = (box[2] + box[3]) / 2.0
w = box[1] - box[0]
h = box[3] - box[2]
# 根据错误信息,错误可能来自于这里的tensor操作。让我们一项一项地做一些更改。
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return (x.item(), y.item(), w.item(), h.item())
def convert_annotation(image_id):
in_file = open('./annotations\%s.xml' % (image_id), encoding='UTF-8') #修改xml路径
out_file = open('./labels2\%s.txt' % (image_id), 'w') # 生成txt格式文件路径
tree = ET.parse(in_file)
root = tree.getroot()
# size = root.find('size')
# w = int(size.find('width').text)
# h = int(size.find('height').text)
# 读取图像,创建相应的torch tensor,这里没有使用unsqueeze,因为针对标量,unsqueeze可能会导致大小为1的额外维度
img = cv2.imread(imgpath + "/" + image_id + '.jpg')
w = torch.tensor(int(img.shape[1]), device=device)
h = torch.tensor(int(img.shape[0]), device=device)
for obj in root.iter('object'):
cls = obj.find('name').text
# print(cls)
if cls not in classes:
continue
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
float(xmlbox.find('ymax').text))
b = torch.tensor(list(b), device=device)
bb = convert((w, h), b)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
xml_path = os.path.join(CURRENT_DIR, './annotations/') #修改xml路径
# xml list
img_xmls = os.listdir(xml_path)
for img_xml in img_xmls:
label_name = img_xml.split('.')[0]
print(label_name)
convert_annotation(label_name)
划分数据集
# by CSDN 迪菲赫尔曼
import os
import random
import shutil
# 设置随机数种子
random.seed(123)
# 定义文件夹路径
root_dir = 'C:/Users/32566/Desktop/dataset' #更改路径
image_dir = os.path.join(root_dir, 'images')
label_dir = os.path.join(root_dir, 'labels')
output_dir = 'C:/Users/32566/Desktop/dataset/JOURNAL_dataset' #更改名字
# 定义训练集、验证集和测试集比例
train_ratio = 0.7
valid_ratio = 0.2
test_ratio = 0.1
# 获取所有图像文件和标签文件的文件名(不包括文件扩展名)
image_filenames = [os.path.splitext(f)[0] for f in os.listdir(image_dir)]
label_filenames = [os.path.splitext(f)[0] for f in os.listdir(label_dir)]
# 随机打乱文件名列表
random.shuffle(image_filenames)
# 计算训练集、验证集和测试集的数量
total_count = len(image_filenames)
train_count = int(total_count * train_ratio)
valid_count = int(total_count * valid_ratio)
test_count = total_count - train_count - valid_count
# 定义输出文件夹路径
train_image_dir = os.path.join(output_dir, 'train', 'images')
train_label_dir = os.path.join(output_dir, 'train', 'labels')
valid_image_dir = os.path.join(output_dir, 'valid', 'images')
valid_label_dir = os.path.join(output_dir, 'valid', 'labels')
test_image_dir = os.path.join(output_dir, 'test', 'images')
test_label_dir = os.path.join(output_dir, 'test', 'labels')
# 创建输出文件夹
os.makedirs(train_image_dir, exist_ok=True)
os.makedirs(train_label_dir, exist_ok=True)
os.makedirs(valid_image_dir, exist_ok=True)
os.makedirs(valid_label_dir, exist_ok=True)
os.makedirs(test_image_dir, exist_ok=True)
os.makedirs(test_label_dir, exist_ok=True)
# 将图像和标签文件划分到不同的数据集中
for i, filename in enumerate(image_filenames):
if i < train_count:
output_image_dir = train_image_dir
output_label_dir = train_label_dir
elif i < train_count + valid_count:
output_image_dir = valid_image_dir
output_label_dir = valid_label_dir
else:
output_image_dir = test_image_dir
output_label_dir = test_label_dir
# 复制图像文件
src_image_path = os.path.join(image_dir, filename + '.jpg')
dst_image_path = os.path.join(output_image_dir, filename + '.jpg')
shutil.copy(src_image_path, dst_image_path)
# 复制标签文件
src_label_path = os.path.join(label_dir, filename + '.txt')
dst_label_path = os.path.join(output_label_dir, filename + '.txt')
shutil.copy(src_label_path, dst_label_path)
在ultralytics/dataset中创建新的数据yaml文件
# moncake
train: D:\Experiment\improve3\ultralytics-main\datasets\JOURNAL_dataset\train # train images (relative to 'path') 128 images
val: D:\Experiment\improve3\ultralytics-main\datasets\JOURNAL_dataset\valid # val images (relative to 'path') 128 images
test: D:\Experiment\improve3\ultralytics-main\datasets\JOURNAL_dataset\test # test images (optional)
# Classes
names:
- UAV
- SENSOR
- BUOY
- SUB_BUOY
- holothurian
- echinus
- scallop
- starfish
建立一个train文件,yolov8.yaml在ultralytics/models/v8修改
from ultralytics import YOLO
if __name__ == '__main__':
# Load a model
model = YOLO(r'yolov8.yaml') # 不使用预训练权重训练
# model = YOLO(r'yolov8.yaml').load("yolov8n.pt") # 使用预训练权重训练
# Trainparameters ----------------------------------------------------------------------------------------------
model.train(
data=r'coco128.yaml', #这里写我们创建的数据集yaml
epochs= 300 , # (int) number of epochs to train for
patience= 50 , # (int) epochs to wait for no observable improvement for early stopping of training
batch= 32 , # (int) number of images per batch (-1 for AutoBatch)
imgsz= 640 , # (int) size of input images as integer or w,h
save= True , # (bool) save train checkpoints and predict results
save_period= -1, # (int) Save checkpoint every x epochs (disabled if < 1)
cache= False , # (bool) True/ram, disk or False. Use cache for data loading
device='' , # (int | str | list, optional) device to run on, i.e. cuda device=0 or device=0,1,2,3 or device=cpu
workers= 8 , # (int) number of worker threads for data loading (per RANK if DDP)
project= 'runs/train', # (str, optional) project name
name= 'exp' ,# (str, optional) experiment name, results saved to 'project/name' directory
exist_ok= False , # (bool) whether to overwrite existing experiment
pretrained= True , # (bool | str) whether to use a pretrained model (bool) or a model to load weights from (str)
optimizer= 'SGD', # (str) optimizer to use, choices=[SGD, Adam, Adamax, AdamW, NAdam, RAdam, RMSProp, auto]
verbose= True ,# (bool) whether to print verbose output
seed= 0 , # (int) random seed for reproducibility
deterministic= True , # (bool) whether to enable deterministic mode
single_cls= False , # (bool) train multi-class data as single-class
rect= False ,# (bool) rectangular training if mode='train' or rectangular validation if mode='val'
cos_lr= False , # (bool) use cosine learning rate scheduler
close_mosaic= 0, # (int) disable mosaic augmentation for final epochs
resume= False , # (bool) resume training from last checkpoint
amp= True, # (bool) Automatic Mixed Precision (AMP) training, choices=[True, False], True runs AMP check
fraction= 1.0 , # (float) dataset fraction to train on (default is 1.0, all images in train set)
profile= False, # (bool) profile ONNX and TensorRT speeds during training for loggers
# Segmentation
overlap_mask= True , # (bool) masks should overlap during training (segment train only)
mask_ratio= 4, # (int) mask downsample ratio (segment train only)
# Classification
dropout= 0.0, # (float) use dropout regularization (classify train only)
# Hyperparameters ----------------------------------------------------------------------------------------------
lr0=0.01, # (float) initial learning rate (i.e. SGD=1E-2, Adam=1E-3)
lrf=0.01, # (float) final learning rate (lr0 * lrf)
momentum=0.937, # (float) SGD momentum/Adam beta1
weight_decay=0.0005, # (float) optimizer weight decay 5e-4
warmup_epochs=3.0, # (float) warmup epochs (fractions ok)
warmup_momentum=0.8, # (float) warmup initial momentum
warmup_bias_lr=0.1, # (float) warmup initial bias lr
box=7.5, # (float) box loss gain
cls=0.5, # (float) cls loss gain (scale with pixels)
dfl=1.5, # (float) dfl loss gain
pose=12.0, # (float) pose loss gain
kobj=1.0, # (float) keypoint obj loss gain
label_smoothing=0.0, # (float) label smoothing (fraction)
nbs=64, # (int) nominal batch size
hsv_h=0.015, # (float) image HSV-Hue augmentation (fraction)
hsv_s=0.7, # (float) image HSV-Saturation augmentation (fraction)
hsv_v=0.4, # (float) image HSV-Value augmentation (fraction)
degrees=0.0, # (float) image rotation (+/- deg)
translate=0.1, # (float) image translation (+/- fraction)
scale=0.5, # (float) image scale (+/- gain)
shear=0.0, # (float) image shear (+/- deg)
perspective=0.0, # (float) image perspective (+/- fraction), range 0-0.001
flipud=0.0, # (float) image flip up-down (probability)
fliplr=0.5, # (float) image flip left-right (probability)
mosaic=1.0, # (float) image mosaic (probability)
mixup=0.0, # (float) image mixup (probability)
copy_paste=0.0, # (float) segment copy-paste (probability)
)
建立一个detect文件
from ultralytics import YOLO
if __name__ == '__main__':
# Load a model
model = YOLO(r'ultralytics\yolov8n.pt') # pretrained YOLOv8n model 权重文件地址
model.predict(
source=r'ultralytics\assets\bus.jpg', #资源地址
save=True, # save predict results
imgsz=640, # (int) size of input images as integer or w,h
conf=0.25, # object confidence threshold for detection (default 0.25 predict, 0.001 val)
iou=0.45, # # intersection over union (IoU) threshold for NMS
show=False, # show results if possible
project='runs/predict', # (str, optional) project name
name='exp', # (str, optional) experiment name, results saved to 'project/name' directory
save_txt=False, # save results as .txt file
save_conf=True, # save results with confidence scores
save_crop=False, # save cropped images with results
show_labels=True, # show object labels in plots
show_conf=True, # show object confidence scores in plots
vid_stride=1, # video frame-rate stride
line_width=3, # bounding box thickness (pixels)
visualize=False, # visualize model features
augment=False, # apply image augmentation to prediction sources
agnostic_nms=False, # class-agnostic NMS
retina_masks=False, # use high-resolution segmentation masks
boxes=True, # Show boxes in segmentation predictions
)
建立一个val文件,计算mAP,FPS的
from ultralytics import YOLO
if __name__ == '__main__':
# Load a model
model = YOLO(r'ultralytics\yolov8n.pt') # build a new model from YAML 加载权重
# Validate the model
model.val(
val=True, # (bool) validate/test during training
data=r'coco128.yaml', #数据集yaml
split='val', # (str) dataset split to use for validation, i.e. 'val', 'test' or 'train'
batch=1, # (int) number of images per batch (-1 for AutoBatch)
imgsz=640, # (int) size of input images as integer or w,h
device='', # (int | str | list, optional) device to run on, i.e. cuda device=0 or device=0,1,2,3 or device=cpu
workers=8, # (int) number of worker threads for data loading (per RANK if DDP)
save_json=False, # (bool) save results to JSON file
save_hybrid=False, # (bool) save hybrid version of labels (labels + additional predictions)
conf=0.001, # (float, optional) object confidence threshold for detection (default 0.25 predict, 0.001 val)
iou=0.6, # (float) intersection over union (IoU) threshold for NMS
project='runs/val', # (str, optional) project name
name='exp', # (str, optional) experiment name, results saved to 'project/name' directory
max_det=300, # (int) maximum number of detections per image
half=False, # (bool) use half precision (FP16)
dnn=False, # (bool) use OpenCV DNN for ONNX inference
plots=True, # (bool) save plots during train/val
)
查看网络参数量
from ultralytics import YOLO
if __name__ == '__main__':
# Load a model
model = YOLO(r'path/to/xxx.yaml') # build a new model from YAML yolov8的yaml文件,不是数据集的yaml
model.info()