huangxin/zjsh_yolov11
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

增加obb

Browse Source
This commit is contained in:
琉璃月光
2025-08-18 15:13:54 +08:00
parent 254afbbc43
commit 6e553f6a20
700 changed files with 576 additions and 1390 deletions
Download Patch File Download Diff File Expand all files Collapse all files

94
yolo11_obb/d_data.py Normal file
View File

@ -0,0 +1,94 @@
import os
import random
import shutil
from tqdm import tqdm
def split_dataset(
images_dir: str, # 图像文件夹路径
labels_dir: str, # 标签文件夹路径
output_root: str, # 输出根目录
train_ratio: float = 0.7, # 训练集比例
val_ratio: float = 0.2, # 验证集比例
image_extensions: tuple = (".jpg", ".png", ".jpeg"), # 支持的图像文件格式
seed: int = 42, # 随机种子
):
"""
将数据集按照 7:2:1 的比例划分为训练集、验证集和测试集，并将图像和标签文件分别复制到相应的文件夹中。
参数:
images_dir (str): 图像文件夹路径。
labels_dir (str): 标签文件夹路径。
output_root (str): 输出根目录。
train_ratio (float): 训练集比例，默认为 0.7。
val_ratio (float): 验证集比例，默认为 0.2。
image_extensions (tuple): 支持的图像文件格式，默认为 (".jpg", ".png", ".jpeg")。
seed (int): 随机种子，默认为 42。
"""
# 设置随机种子
random.seed(seed)
# 创建输出文件夹
output_images_dir = os.path.join(output_root, "images")
output_labels_dir = os.path.join(output_root, "labels")
# 创建训练集、验证集和测试集的图像文件夹
for subset in ["train", "val", "test"]:
os.makedirs(os.path.join(output_images_dir, subset), exist_ok=True) # 图像文件夹
# 创建训练集和验证集的标签文件夹（指定路径）
os.makedirs(os.path.join(output_labels_dir, "train_original"), exist_ok=True) # 训练集标签文件夹
os.makedirs(os.path.join(output_labels_dir, "val_original"), exist_ok=True) # 验证集标签文件夹
os.makedirs(os.path.join(output_labels_dir, "test"), exist_ok=True) # 测试集标签文件夹
# 获取所有图像文件的文件名列表
image_files = [f for f in os.listdir(images_dir) if f.endswith(image_extensions)]
random.shuffle(image_files) # 随机打乱文件列表
# 计算训练集、验证集和测试集的大小
total_files = len(image_files)
train_size = int(total_files * train_ratio) # 训练集大小
val_size = int(total_files * val_ratio) # 验证集大小
test_size = total_files - train_size - val_size # 测试集大小
# 复制图像和标签文件到相应的子集文件夹中
for i, image_file in enumerate(tqdm(image_files, desc="分割数据集中")):
base_file_name = os.path.splitext(image_file)[0] # 获取文件名（不包括扩展名）
image_path = os.path.join(images_dir, image_file) # 图像文件路径
label_path = os.path.join(labels_dir, base_file_name + ".txt") # 标签文件路径
# 根据索引判断文件应复制到训练集、验证集还是测试集
if i < train_size:
# 复制到训练集
shutil.copy(image_path, os.path.join(output_images_dir, "train", image_file)) # 复制图像
shutil.copy(label_path, os.path.join(output_labels_dir, "train_original", base_file_name + ".txt")) # 复制标签
elif i < train_size + val_size:
# 复制到验证集
shutil.copy(image_path, os.path.join(output_images_dir, "val", image_file)) # 复制图像
shutil.copy(label_path, os.path.join(output_labels_dir, "val_original", base_file_name + ".txt")) # 复制标签
else:
# 复制到测试集
shutil.copy(image_path, os.path.join(output_images_dir, "test", image_file)) # 复制图像
shutil.copy(label_path, os.path.join(output_labels_dir, "test", base_file_name + ".txt")) # 复制标签
print(f"数据集分割完成！训练集: {train_size} 个样本，验证集: {val_size} 个样本，测试集: {test_size} 个样本。")
# 示例调用
if __name__ == "__main__":
# 数据集路径
images_dir = r"/home/hx/桌面/image/image" # 需要读取的所有图像文件夹路径
labels_dir = r"/home/hx/桌面/image/2" # 需要读取的所有图像与之对应的txt标签文件夹路径
# 输出路径
output_root = r"/home/hx/桌面/image" # 保存最终数据集的根目录
# 调用函数
split_dataset(
images_dir=images_dir,
labels_dir=labels_dir,
output_root=output_root,
train_ratio=0.7, # 训练集比例
val_ratio=0.2, # 验证集比例
image_extensions=(".jpg", ".png", ".jpeg"), # 支持的图像文件格式
seed=42, # 随机种子
)

4
yolo11_obb/d_data_converter .py Normal file
View File

@ -0,0 +1,4 @@
from ultralytics.data.converter import convert_dota_to_yolo_obb
# 输入需要转换TXT标签文件的数据集文件夹名字的路径
convert_dota_to_yolo_obb(r"/home/hx/桌面/image")

225
yolo11_obb/data_trans_dota.py Normal file
View File

@ -0,0 +1,225 @@
import os
import xml.etree.ElementTree as ET
import math
"""
obb转换代码使用：
只用修改if __name__ == '__main__': 下面的三条路径即可：
roxml_path = r"path/已经存放着原始XML文件的文件夹名字"
dotaxml_path = r"path/准备要存放DOTA格式的XML文件的文件夹名字"
# （小建议：先手动创建该文件夹，然后路径放在这里）
out_path = r"path/准备要存放DOTA格式的TXT文件的文件夹名字"
# （小建议：先手动创建该文件夹，然后路径放在这里）
"""
def edit_xml(xml_file, dotaxml_file):
"""
修改 XML 文件，将 bndbox 或 robndbox 转换为四点坐标格式。
:param xml_file: 原始 XML 文件路径
:param dotaxml_file: 转换后的 XML 文件保存路径
"""
tree = ET.parse(xml_file)
objs = tree.findall('object')
for ix, obj in enumerate(objs):
x0 = ET.Element("x0") # 创建节点
y0 = ET.Element("y0")
x1 = ET.Element("x1")
y1 = ET.Element("y1")
x2 = ET.Element("x2")
y2 = ET.Element("y2")
x3 = ET.Element("x3")
y3 = ET.Element("y3")
if obj.find('robndbox') is None:
# 处理 bndbox
obj_bnd = obj.find('bndbox')
obj_xmin = obj_bnd.find('xmin')
obj_ymin = obj_bnd.find('ymin')
obj_xmax = obj_bnd.find('xmax')
obj_ymax = obj_bnd.find('ymax')
xmin = float(obj_xmin.text)
ymin = float(obj_ymin.text)
xmax = float(obj_xmax.text)
ymax = float(obj_ymax.text)
obj_bnd.remove(obj_xmin) # 删除节点
obj_bnd.remove(obj_ymin)
obj_bnd.remove(obj_xmax)
obj_bnd.remove(obj_ymax)
x0.text = str(xmin)
y0.text = str(ymax)
x1.text = str(xmax)
y1.text = str(ymax)
x2.text = str(xmax)
y2.text = str(ymin)
x3.text = str(xmin)
y3.text = str(ymin)
else:
# 处理 robndbox
obj_bnd = obj.find('robndbox')
obj_bnd.tag = 'bndbox' # 修改节点名
obj_cx = obj_bnd.find('cx')
obj_cy = obj_bnd.find('cy')
obj_w = obj_bnd.find('w')
obj_h = obj_bnd.find('h')
obj_angle = obj_bnd.find('angle')
cx = float(obj_cx.text)
cy = float(obj_cy.text)
w = float(obj_w.text)
h = float(obj_h.text)
angle = float(obj_angle.text)
obj_bnd.remove(obj_cx) # 删除节点
obj_bnd.remove(obj_cy)
obj_bnd.remove(obj_w)
obj_bnd.remove(obj_h)
obj_bnd.remove(obj_angle)
x0.text, y0.text = rotatePoint(cx, cy, cx - w / 2, cy - h / 2, -angle)
x1.text, y1.text = rotatePoint(cx, cy, cx + w / 2, cy - h / 2, -angle)
x2.text, y2.text = rotatePoint(cx, cy, cx + w / 2, cy + h / 2, -angle)
x3.text, y3.text = rotatePoint(cx, cy, cx - w / 2, cy + h / 2, -angle)
obj_bnd.append(x0) # 新增节点
obj_bnd.append(y0)
obj_bnd.append(x1)
obj_bnd.append(y1)
obj_bnd.append(x2)
obj_bnd.append(y2)
obj_bnd.append(x3)
obj_bnd.append(y3)
tree.write(dotaxml_file, method='xml', encoding='utf-8') # 更新 XML 文件
def rotatePoint(xc, yc, xp, yp, theta):
"""
计算旋转后的点坐标。
:param xc: 旋转中心 x 坐标
:param yc: 旋转中心 y 坐标
:param xp: 点 x 坐标
:param yp: 点 y 坐标
:param theta: 旋转角度（弧度）
:return: 旋转后的点坐标 (x, y)
"""
xoff = xp - xc
yoff = yp - yc
cosTheta = math.cos(theta)
sinTheta = math.sin(theta)
pResx = cosTheta * xoff + sinTheta * yoff
pResy = -sinTheta * xoff + cosTheta * yoff
return str(int(xc + pResx)), str(int(yc + pResy))
def get_unique_classes(xml_path):
"""
从 XML 文件中提取所有唯一的类别名称。
:param xml_path: XML 文件所在的目录
:return: 包含所有唯一类别名称的集合
"""
unique_classes = set() # 使用集合存储唯一的类别名称
files = os.listdir(xml_path)
for file in files:
if not file.endswith('.xml'):
continue
tree = ET.parse(os.path.join(xml_path, file))
root = tree.getroot()
for obj in root.findall('object'):
cls = obj.find('name').text
unique_classes.add(cls)
return unique_classes
def generate_class_mapping(unique_classes):
"""
根据唯一的类别名称生成类别映射字典。
:param unique_classes: 包含所有唯一类别名称的集合
:return: 类别名称到编号的映射字典
"""
class_mapping = {}
for index, cls in enumerate(sorted(unique_classes)): # 按字母顺序排序
class_mapping[cls] = index
return class_mapping
def totxt(xml_path, out_path, class_mapping):
"""
将 XML 文件转换为 TXT 文件，并根据类别名称动态生成编号。
:param xml_path: XML 文件所在的目录
:param out_path: 保存 TXT 文件的目录
:param class_mapping: 类别名称到编号的映射字典
"""
# 确保输出目录存在
if not os.path.exists(out_path):
os.makedirs(out_path)
# 遍历 xml_path 下的所有文件
files = os.listdir(xml_path)
for file in files:
# 只处理 .xml 文件
if not file.endswith('.xml'):
continue
# 解析 XML 文件
tree = ET.parse(os.path.join(xml_path, file))
root = tree.getroot()
# 获取文件名（不带扩展名）
name = os.path.splitext(file)[0]
# 设置输出文件路径
output = os.path.join(out_path, name + '.txt') # 使用 os.path.join 确保路径正确
with open(output, 'w') as f: # 使用 with 打开文件，确保文件正确关闭
objs = tree.findall('object')
for obj in objs:
cls = obj.find('name').text
box = obj.find('bndbox')
x0 = int(float(box.find('x0').text))
y0 = int(float(box.find('y0').text))
x1 = int(float(box.find('x1').text))
y1 = int(float(box.find('y1').text))
x2 = int(float(box.find('x2').text))
y2 = int(float(box.find('y2').text))
x3 = int(float(box.find('x3').text))
y3 = int(float(box.find('y3').text))
# 根据类别名称获取对应的编号
cls_index = class_mapping.get(cls, -1) # 如果类别不存在，默认返回 -1
if cls_index == -1:
print(f"Warning: Class '{cls}' not found in class_mapping. Skipping this object.")
continue
# 写入文件
f.write("{} {} {} {} {} {} {} {} {} {}\n".format(x0, y0, x1, y1, x2, y2, x3, y3, cls, cls_index))
print(f"Generated: {output}")
if __name__ == '__main__':
# 设置路径
roxml_path = r"/home/hx/桌面/image/1" # 存放原始 XML 文件夹名字的路径
dotaxml_path = r"/home/hx/桌面/image/2" # 存放 DOTA格式的 XML文件夹名字（小建议：先手动创建该文件夹，然后路径放在这里）
out_path = r"/home/hx/桌面/image/2" # 存放 DOTA格式的 TXT 文件夹名字（小建议：先手动创建该文件夹，然后路径放在这里）
# 第一步：将 XML 文件统一转换成旋转框的 XML 文件
filelist = os.listdir(roxml_path)
for file in filelist:
print(f"Processing: {os.path.join(roxml_path, file)}")
edit_xml(os.path.join(roxml_path, file), os.path.join(dotaxml_path, file))
# 第二步：从 XML 文件中提取所有唯一的类别名称
unique_classes = get_unique_classes(dotaxml_path)
print(f"Unique classes found: {unique_classes}")
# 第三步：生成类别映射字典
class_mapping = generate_class_mapping(unique_classes)
print(f"Generated class mapping: {class_mapping}")
# 第四步：将旋转框 XML 文件转换成 TXT 格式
totxt(dotaxml_path, out_path, class_mapping)

73
yolo11_obb/final_trans_obb.py Normal file
View File

@ -0,0 +1,73 @@
# convert_dota_to_yolo_obb_manual.py
from pathlib import Path
def convert_dota_to_yolo_obb_manual(dota_root: str):
dota_root = Path(dota_root)
class_mapping = {"clamp": 0} # 你可以加更多类别
def convert_file(txt_path: Path, image_width: int, image_height: int, save_path: Path):
with open(txt_path, 'r') as f:
lines = f.readlines()
with open(save_path, 'w') as g:
for line in lines:
parts = line.strip().split()
if len(parts) < 9:
continue
x1, y1, x2, y2, x3, y3, x4, y4 = map(float, parts[:8])
class_name = parts[8]
# difficult = parts[9] # 可选，一般不用
class_id = class_mapping.get(class_name.lower(), -1)
if class_id == -1:
print(f"⚠️ 未知类别: {class_name}, 文件: {txt_path}")
continue
# 归一化坐标
coords = [x1, y1, x2, y2, x3, y3, x4, y4]
normalized = [
coords[i] / (image_width if i % 2 == 0 else image_height)
for i in range(8)
]
# 写入 YOLO-OBB 格式：class_id x1 y1 x2 y2 x3 y3 x4 y4
line_str = f"{class_id} " + " ".join(f"{v:.6g}" for v in normalized)
g.write(line_str + "\n")
for phase in ["train", "val"]:
image_dir = dota_root / "images" / phase
label_dir = dota_root / "labels" / phase
save_dir = dota_root / "labels_obb" / phase
save_dir.mkdir(parents=True, exist_ok=True)
for img_file in image_dir.iterdir():
if img_file.suffix.lower() not in [".jpg", ".jpeg", ".png"]:
continue
# 获取图像宽高
from PIL import Image
img = Image.open(img_file)
w, h = img.size
# 找对应标签文件
txt_file = label_dir / f"{img_file.stem}.txt"
if not txt_file.exists():
print(f"⚠️ 找不到标签文件: {txt_file}")
continue
save_file = save_dir / f"{img_file.stem}.txt"
convert_file(txt_file, w, h, save_file)
# 自动生成 dataset.yaml
dataset_yaml = dota_root / "dataset.yaml"
with open(dataset_yaml, 'w') as f:
f.write(f"""train: {dota_root / 'images' / 'train'}\n""")
f.write(f"""val: {dota_root / 'images' / 'val'}\n\n""")
f.write(f"nc: {len(class_mapping)}\n")
f.write(f"names: {list(class_mapping.keys())}\n")
print(f"✅ 转换完成！标签已保存到: {dota_root / 'labels_obb'}")
print(f"✅ dataset.yaml 已生成: {dataset_yaml}")
# 开始转换
convert_dota_to_yolo_obb_manual("/home/hx/桌面/image")