73 lines
2.8 KiB
Python
73 lines
2.8 KiB
Python
|
# 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")
|