|
|
import os
|
|
|
import shutil
|
|
|
import random
|
|
|
|
|
|
|
|
|
def get_all_dirs_with_files(root_dir):
|
|
|
"""递归获取 root_dir 下所有包含文件的目录"""
|
|
|
result = []
|
|
|
for dirpath, dirnames, filenames in os.walk(root_dir):
|
|
|
if filenames:
|
|
|
result.append(dirpath)
|
|
|
return result
|
|
|
|
|
|
|
|
|
def process_directory(dir_path, seq_num, all_pairs):
|
|
|
"""处理单个目录：修改文本文件第一列为序列号，BMP 复制对应的 txt"""
|
|
|
all_files = os.listdir(dir_path)
|
|
|
|
|
|
bmp_files = [f for f in all_files if f.lower().endswith('.bmp')]
|
|
|
txt_files = [f for f in all_files if f.lower().endswith('.txt')]
|
|
|
other_files = [f for f in all_files
|
|
|
if not f.lower().endswith('.bmp') and not f.lower().endswith('.txt')]
|
|
|
|
|
|
text_candidates = txt_files + other_files
|
|
|
|
|
|
if not text_candidates:
|
|
|
print(f" [{os.path.basename(dir_path)}] 未找到文本文件，跳过")
|
|
|
return
|
|
|
|
|
|
text_filename = text_candidates[0]
|
|
|
text_path = os.path.join(dir_path, text_filename)
|
|
|
|
|
|
if not os.path.isfile(text_path):
|
|
|
print(f" [{os.path.basename(dir_path)}] 跳过非文件: {text_filename}")
|
|
|
return
|
|
|
|
|
|
short_name = os.path.basename(dir_path)
|
|
|
print(f" 处理: {short_name} (seq={seq_num}), 文本: {text_filename}")
|
|
|
|
|
|
# 读取文件
|
|
|
try:
|
|
|
with open(text_path, 'r', encoding='utf-8') as f:
|
|
|
lines = f.readlines()
|
|
|
except (UnicodeDecodeError, PermissionError, OSError):
|
|
|
try:
|
|
|
with open(text_path, 'r', encoding='gbk') as f:
|
|
|
lines = f.readlines()
|
|
|
except Exception as e:
|
|
|
print(f" 读取失败: {e}")
|
|
|
return
|
|
|
|
|
|
modified_lines = []
|
|
|
for line in lines:
|
|
|
line = line.rstrip('\n').rstrip('\r')
|
|
|
if not line.strip():
|
|
|
modified_lines.append(line)
|
|
|
continue
|
|
|
parts = line.split()
|
|
|
if parts:
|
|
|
parts[0] = str(seq_num)
|
|
|
modified_lines.append(' '.join(parts))
|
|
|
|
|
|
# 写回
|
|
|
with open(text_path, 'w', encoding='utf-8') as f:
|
|
|
f.write('\n'.join(modified_lines))
|
|
|
print(f" 已修改 {len([l for l in modified_lines if l.strip()])} 行")
|
|
|
|
|
|
# 为每个 BMP 复制/关联文档
|
|
|
if bmp_files:
|
|
|
for bmp_filename in bmp_files:
|
|
|
bmp_basename = os.path.splitext(bmp_filename)[0]
|
|
|
copy_filename = f"{bmp_basename}.txt"
|
|
|
copy_path = os.path.join(dir_path, copy_filename)
|
|
|
|
|
|
if os.path.normpath(text_path) == os.path.normpath(copy_path):
|
|
|
all_pairs.append((os.path.join(dir_path, bmp_filename), copy_path))
|
|
|
else:
|
|
|
shutil.copy2(text_path, copy_path)
|
|
|
all_pairs.append((os.path.join(dir_path, bmp_filename), copy_path))
|
|
|
print(f" BMP 数量: {len(bmp_files)}")
|
|
|
else:
|
|
|
print(f" 无 BMP 图片")
|
|
|
|
|
|
|
|
|
def main():
|
|
|
base_dir = r"D:\xwechat_files\wxid_ttkf0xgmyihv22_2388\msg\file\2026-07\抓图(1)\抓图\新建文件夹\1\1\新建文件夹"
|
|
|
|
|
|
output_dir = os.path.join(os.path.dirname(base_dir.rstrip('\\/')), "yolo_output")
|
|
|
train_img_dir = os.path.join(output_dir, "train", "images")
|
|
|
train_lbl_dir = os.path.join(output_dir, "train", "labels")
|
|
|
val_img_dir = os.path.join(output_dir, "val", "images")
|
|
|
val_lbl_dir = os.path.join(output_dir, "val", "labels")
|
|
|
|
|
|
for d in [train_img_dir, train_lbl_dir, val_img_dir, val_lbl_dir]:
|
|
|
os.makedirs(d, exist_ok=True)
|
|
|
|
|
|
if not os.path.isdir(base_dir):
|
|
|
print(f"错误: 路径不存在: {base_dir}")
|
|
|
return
|
|
|
|
|
|
top_dirs = [d for d in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, d))]
|
|
|
top_dirs.sort()
|
|
|
|
|
|
if not top_dirs:
|
|
|
print("没有找到子文件夹")
|
|
|
return
|
|
|
|
|
|
seq_start = 20
|
|
|
all_pairs = [] # (bmp_path, txt_path)
|
|
|
mapping = {} # 顶层文件夹名 -> 序列号
|
|
|
|
|
|
for idx, top_name in enumerate(top_dirs):
|
|
|
top_path = os.path.join(base_dir, top_name)
|
|
|
seq_num = seq_start + idx
|
|
|
mapping[top_name] = seq_num
|
|
|
|
|
|
print(f"\n{'='*50}")
|
|
|
print(f"顶层文件夹: {top_name} -> 序列号: {seq_num}")
|
|
|
|
|
|
# 递归获取该顶层文件夹下所有含文件的目录
|
|
|
leaf_dirs = get_all_dirs_with_files(top_path)
|
|
|
|
|
|
if not leaf_dirs:
|
|
|
print(f" 无文件，跳过")
|
|
|
continue
|
|
|
|
|
|
for leaf_dir in leaf_dirs:
|
|
|
process_directory(leaf_dir, seq_num, all_pairs)
|
|
|
|
|
|
# ---- 输出对应关系 ----
|
|
|
print("\n" + "=" * 50)
|
|
|
print("文件夹与序列号对应关系:")
|
|
|
print("-" * 40)
|
|
|
for name, num in mapping.items():
|
|
|
print(f" {num} : {name}")
|
|
|
print("=" * 50)
|
|
|
|
|
|
# ---- YOLO 格式 train/val 划分 ----
|
|
|
if not all_pairs:
|
|
|
print("\n没有 BMP 图片，跳过 YOLO 输出")
|
|
|
return
|
|
|
|
|
|
random.seed(42)
|
|
|
random.shuffle(all_pairs)
|
|
|
|
|
|
split_idx = int(len(all_pairs) * 0.8)
|
|
|
train_pairs = all_pairs[:split_idx]
|
|
|
val_pairs = all_pairs[split_idx:]
|
|
|
|
|
|
print(f"\n总图片数: {len(all_pairs)}")
|
|
|
print(f"train: {len(train_pairs)} | val: {len(val_pairs)}")
|
|
|
|
|
|
def copy_pairs(pairs, img_dir, lbl_dir):
|
|
|
for bmp_path, txt_path in pairs:
|
|
|
bmp_name = os.path.basename(bmp_path)
|
|
|
txt_name = os.path.basename(txt_path)
|
|
|
|
|
|
dst_bmp = bmp_name
|
|
|
dst_txt = txt_name
|
|
|
|
|
|
counter = 0
|
|
|
while os.path.exists(os.path.join(img_dir, dst_bmp)):
|
|
|
counter += 1
|
|
|
base, ext = os.path.splitext(bmp_name)
|
|
|
dst_bmp = f"{base}_{counter}{ext}"
|
|
|
dst_txt = f"{base}_{counter}.txt"
|
|
|
|
|
|
shutil.copy2(bmp_path, os.path.join(img_dir, dst_bmp))
|
|
|
shutil.copy2(txt_path, os.path.join(lbl_dir, dst_txt))
|
|
|
|
|
|
copy_pairs(train_pairs, train_img_dir, train_lbl_dir)
|
|
|
copy_pairs(val_pairs, val_img_dir, val_lbl_dir)
|
|
|
|
|
|
print(f"\nYOLO 格式输出: {output_dir}")
|
|
|
print(f" train/images/ ({len(train_pairs)} 张)")
|
|
|
print(f" train/labels/ ({len(train_pairs)} 个)")
|
|
|
print(f" val/images/ ({len(val_pairs)} 张)")
|
|
|
print(f" val/labels/ ({len(val_pairs)} 个)")
|
|
|
print("\n全部完成!")
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
main()
|