commit e7c35a017be5f8def7f7a0db544c09b358a010ec Author: ninghongbin <2409766686@qq.com> Date: Fri Sep 26 10:23:45 2025 +0800 数据预处理 diff --git a/README.md b/README.md new file mode 100644 index 0000000..f404791 --- /dev/null +++ b/README.md @@ -0,0 +1,9 @@ +### 数据集预处理说明 + + + +- 读取一段原始视频、切片为帧 + +- 然后从原始数据集贴图到原始视频帧、模拟显示识别的复杂场景 +- 然后生成最新的数据文件和标注文件 +- 即可使用最新生成的文件进行训练 \ No newline at end of file diff --git a/数据增强底图版.py b/数据增强底图版.py new file mode 100644 index 0000000..92874df --- /dev/null +++ b/数据增强底图版.py @@ -0,0 +1,272 @@ +import cv2 +import numpy as np +import os +import random +import albumentations as A +from tqdm import tqdm + +# --- 1. 用户配置（重点修改！！！）--- +# 请根据你的实际路径修改，三个核心目录需区分清楚： +# 1. 特征素材来源目录：存放有「待粘贴目标（如no_helmet）」的图片和标签（用于提取可粘贴的目标） +SOURCE_FEATURE_IMAGE_DIR = r"E:\NSFW-Detection-YOLO\data\images\val\images" # 有目标的原图 +SOURCE_FEATURE_LABEL_DIR = r"E:\NSFW-Detection-YOLO\data\images\val\labels" # 对应原图的标签 +# 2. 独立底图目录：存放你要粘贴目标的「空白/背景底图」（底图无需标签） +BASE_IMAGE_DIR = r"D:\DataPreHandler\images\valid" # 你的底图文件夹 +# 3. 输出目录：保存最终增强后的图片和标签 +OUTPUT_IMAGE_DIR = r"D:\DataPreHandler\data\valid\images" +OUTPUT_LABEL_DIR = r"D:\DataPreHandler\data\valid\labels" + +# 数据增强参数 +AUGMENTATION_FACTOR = 1 # 每张底图生成的增强图数量（如40张） + +# --- Copy-Paste 核心配置 --- +SMALL_OBJECT_CLASSES_TO_PASTE = [0,1,2,3,4,5,6] # 要粘贴的目标类别ID（如no_helmet是2） +PASTE_COUNT_RANGE = (5, 10) # 每张增强图上粘贴的目标数量（随机5-10个） + +# --- 2. 常规增强流水线（修复Albumentations参数）--- +transform_geometric = A.Compose([ + A.HorizontalFlip(p=0.5), + # 修改1：A.Affine参数：rotate_limit→rotate，cval→pad_val，新增border_mode + A.Affine(scale=(0.8, 1.2), shear=(-10, 10), translate_percent=0.1, + rotate=30, border_mode=cv2.BORDER_CONSTANT, pad_val=0, p=0.8), + A.Perspective(scale=(0.02, 0.05), p=0.4), +], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels'], min_visibility=0.25)) + +transform_quality = A.Compose([ + A.RandomBrightnessContrast(brightness_limit=0.25, contrast_limit=0.25, p=0.8), + A.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.7), + # 修改2：A.GaussNoise参数：var_limit→std_limit（方差转标准差，数值取平方根近似） + A.OneOf([A.GaussNoise(std_limit=(3.0, 8.0), p=1.0), A.ISONoise(p=1.0)], p=0.6), + A.OneOf([A.Blur(blur_limit=(3, 7), p=1.0), A.MotionBlur(blur_limit=(3, 7), p=1.0)], p=0.5), + # 修改3：A.ImageCompression参数：quality_lower/upper→quality_range（合并为元组） + A.ImageCompression(quality_range=(70, 95), p=0.3), +], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels'], min_visibility=0.25)) + +transform_mixed = A.Compose([ + A.HorizontalFlip(p=0.5), + # 修改4：A.Rotate参数：value→pad_val + A.Rotate(limit=15, p=0.5, border_mode=cv2.BORDER_CONSTANT, pad_val=0), + A.RandomBrightnessContrast(p=0.6), + A.GaussNoise(std_limit=(2.0, 6.0), p=0.4), # 同步修改GaussNoise参数 + A.Blur(blur_limit=3, p=0.3), +], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels'], min_visibility=0.25)) + +base_transforms = [transform_geometric, transform_quality, transform_mixed] # 随机选择增强策略 + + +# --- 3. 核心工具函数 --- +def harvest_objects_for_pasting(feature_image_dir, feature_label_dir, target_classes): + """ + 从「特征素材来源目录」提取目标，创建可粘贴的素材库 + :param feature_image_dir: 有目标的图片目录（如含no_helmet的原图） + :param feature_label_dir: 对应图片的标签目录 + :param target_classes: 要提取的目标类别（如[2]） + :return: 素材库 {类别ID: [目标图像1, 目标图像2, ...]} + """ + print(f"正在从 {feature_image_dir} 提取目标类别 {target_classes}...") + asset_library = {cls_id: [] for cls_id in target_classes} + + # 只读取特征素材目录中的图片文件 + feature_image_files = [f for f in os.listdir(feature_image_dir) if f.lower().endswith(('.jpg', '.png', '.jpeg'))] + if not feature_image_files: + raise FileNotFoundError(f"特征素材目录 {feature_image_dir} 中未找到图片！") + + for filename in tqdm(feature_image_files, desc="提取目标素材"): + label_file = os.path.splitext(filename)[0] + ".txt" + label_path = os.path.join(feature_label_dir, label_file) + if not os.path.exists(label_path): + continue # 跳过无标签的图片 + + # 读取图片并获取尺寸 + img = cv2.imread(os.path.join(feature_image_dir, filename)) + if img is None: + tqdm.write(f"警告：无法读取图片 {filename}，已跳过") + continue + img_h, img_w, _ = img.shape + + # 解析标签，裁剪目标 + with open(label_path, 'r') as f: + for line in f.readlines(): + line = line.strip() + if not line: + continue + parts = line.split() + # 修改5：处理标签类别ID为浮点数的情况（如6.0→6）：先转float再转int + cls_id = int(float(parts[0])) + if cls_id not in target_classes: + continue # 只保留目标类别 + + # YOLO归一化坐标转像素坐标（x1,y1:左上角；x2,y2:右下角） + x_center, y_center, box_w, box_h = [float(p) for p in parts[1:]] + x1 = int((x_center - box_w / 2) * img_w) + y1 = int((y_center - box_h / 2) * img_h) + x2 = int((x_center + box_w / 2) * img_w) + y2 = int((y_center + box_h / 2) * img_h) + + # 确保坐标在图片范围内，避免裁剪出错 + x1, y1 = max(0, x1), max(0, y1) + x2, y2 = min(img_w, x2), min(img_h, y2) + + # 裁剪目标并加入素材库（排除空图像） + if x1 < x2 and y1 < y2: + cropped_obj = img[y1:y2, x1:x2] + if cropped_obj.size > 0: + asset_library[cls_id].append(cropped_obj) + + # 检查素材库是否为空 + total_assets = sum(len(v) for v in asset_library.values()) + if total_assets == 0: + raise ValueError(f"未从特征素材目录提取到任何目标！请检查类别ID {target_classes} 是否正确") + + print(f"素材库创建完成！共提取 {total_assets} 个目标（类别：{target_classes}）") + return asset_library + + +def paste_objects_to_base(base_image, asset_library): + """ + 将素材库中的目标粘贴到单张底图上 + :param base_image: 输入的底图（cv2读取的BGR图像） + :param asset_library: 目标素材库 + :return: 粘贴后的图像、对应的YOLO格式标签（bboxes + labels） + """ + base_h, base_w, _ = base_image.shape + pasted_bboxes = [] # 存储粘贴目标的YOLO bbox + pasted_labels = [] # 存储粘贴目标的类别ID + + # 随机确定本次要粘贴的目标数量 + num_to_paste = random.randint(*PASTE_COUNT_RANGE) + + for _ in range(num_to_paste): + # 选择要粘贴的目标类别（只从有素材的类别中选） + valid_classes = [cls for cls, assets in asset_library.items() if len(assets) > 0] + if not valid_classes: + break # 极端情况：素材库临时为空（几乎不会发生） + + # 随机选择一个目标类别和该类别下的一个素材 + target_cls = random.choice(valid_classes) + target_obj = random.choice(asset_library[target_cls]) + obj_h, obj_w, _ = target_obj.shape + + # 跳过比底图大的目标（避免粘贴后超出边界） + if obj_h >= base_h or obj_w >= base_w: + continue + + # 随机选择粘贴位置（左上角坐标，确保目标完全在底图内） + paste_x1 = random.randint(0, base_w - obj_w) + paste_y1 = random.randint(0, base_h - obj_h) + paste_x2 = paste_x1 + obj_w + paste_y2 = paste_y1 + obj_h + + # 直接用Numpy切片粘贴目标（覆盖底图对应区域） + base_image[paste_y1:paste_y2, paste_x1:paste_x2] = target_obj + + # 计算粘贴目标的YOLO归一化坐标（x_center, y_center, w, h） + yolo_x_center = (paste_x1 + obj_w / 2) / base_w + yolo_y_center = (paste_y1 + obj_h / 2) / base_h + yolo_w = obj_w / base_w + yolo_h = obj_h / base_h + + # 将标签加入列表 + pasted_bboxes.append([yolo_x_center, yolo_y_center, yolo_w, yolo_h]) + pasted_labels.append(target_cls) + + return base_image, pasted_bboxes, pasted_labels + + +def main(): + # 1. 初始化：创建输出目录 + os.makedirs(OUTPUT_IMAGE_DIR, exist_ok=True) + os.makedirs(OUTPUT_LABEL_DIR, exist_ok=True) + + # 2. 第一步：创建目标素材库（从特征素材目录提取可粘贴的目标） + try: + asset_library = harvest_objects_for_pasting( + feature_image_dir=SOURCE_FEATURE_IMAGE_DIR, + feature_label_dir=SOURCE_FEATURE_LABEL_DIR, + target_classes=SMALL_OBJECT_CLASSES_TO_PASTE + ) + except (FileNotFoundError, ValueError) as e: + print(f"错误：{e}") + return + + # 3. 第二步：获取所有底图（只读取图片文件） + base_image_files = [f for f in os.listdir(BASE_IMAGE_DIR) if f.lower().endswith(('.jpg', '.png', '.jpeg'))] + if not base_image_files: + print(f"错误：底图目录 {BASE_IMAGE_DIR} 中未找到任何图片！") + return + print(f"\n找到 {len(base_image_files)} 张底图，开始生成增强数据（每张底图生成 {AUGMENTATION_FACTOR} 张）") + + # 4. 主循环：遍历每张底图，生成增强数据 + for base_filename in tqdm(base_image_files, desc="处理底图"): + base_name, base_ext = os.path.splitext(base_filename) + base_image_path = os.path.join(BASE_IMAGE_DIR, base_filename) + + # 读取底图（若读取失败则跳过） + base_image = cv2.imread(base_image_path) + if base_image is None: + tqdm.write(f"\n警告：无法读取底图 {base_filename}，已跳过") + continue + + # 为当前底图生成 AUGMENTATION_FACTOR 张增强图 + for aug_idx in range(AUGMENTATION_FACTOR): + # 步骤1：复制底图（避免修改原始底图），并粘贴目标 + base_image_copy = base_image.copy() + pasted_image, pasted_bboxes, pasted_labels = paste_objects_to_base( + base_image=base_image_copy, + asset_library=asset_library + ) + + # 步骤2：对粘贴后的图像应用常规增强（Albumentations需要RGB格式） + pasted_image_rgb = cv2.cvtColor(pasted_image, cv2.COLOR_BGR2RGB) + chosen_transform = random.choice(base_transforms) # 随机选择增强策略 + + try: + # 应用增强（同时处理bbox和label） + augmented_result = chosen_transform( + image=pasted_image_rgb, + bboxes=pasted_bboxes, + class_labels=pasted_labels + ) + final_image_rgb = augmented_result['image'] + final_bboxes = augmented_result['bboxes'] + final_labels = augmented_result['class_labels'] + except Exception as e: + tqdm.write(f"\n警告：底图 {base_filename} 增强失败（序号 {aug_idx}）：{str(e)}") + continue + + # 步骤3：保存增强后的图片和标签 + # 图片命名格式：底图名_aug_序号.jpg（统一转为jpg格式，避免格式混乱） + output_img_name = f"{base_name}_aug_{aug_idx}.jpg" + output_img_path = os.path.join(OUTPUT_IMAGE_DIR, output_img_name) + # RGB转BGR（cv2保存需要BGR格式） + cv2.imwrite(output_img_path, cv2.cvtColor(final_image_rgb, cv2.COLOR_RGB2BGR)) + + # 标签命名格式：与图片同名.txt（YOLO格式） + output_label_name = f"{base_name}_aug_{aug_idx}.txt" + output_label_path = os.path.join(OUTPUT_LABEL_DIR, output_label_name) + + with open(output_label_path, 'w') as f: + for bbox, label in zip(final_bboxes, final_labels): + x_c, y_c, w, h = bbox + # 边界检查：排除增强后可能超出0-1范围的bbox（避免训练报错） + if 0 <= x_c <= 1 and 0 <= y_c <= 1 and 0 <= w <= 1 and 0 <= h <= 1: + f.write(f"{label} {x_c:.6f} {y_c:.6f} {w:.6f} {h:.6f}\n") + + # 5. 完成提示 + total_generated = len(base_image_files) * AUGMENTATION_FACTOR + print(f"\n✅ 数据增强全部完成！") + print(f"📊 生成数据统计：") + print(f" - 底图数量：{len(base_image_files)} 张") + print(f" - 每张底图增强次数：{AUGMENTATION_FACTOR} 次") + print(f" - 总生成图片/标签：{total_generated} 组") + print(f" - 输出路径：") + print(f" 图片 → {OUTPUT_IMAGE_DIR}") + print(f" 标签 → {OUTPUT_LABEL_DIR}") + + +if __name__ == "__main__": + # 运行前务必确认： + # 1. SOURCE_FEATURE_IMAGE_DIR/SOURCE_FEATURE_LABEL_DIR 是「有目标的素材目录」 + # 2. BASE_IMAGE_DIR 是你的「空白底图目录」 + # 3. SMALL_OBJECT_CLASSES_TO_PASTE 是要粘贴的目标类别ID（如no_helmet=2） + main() \ No newline at end of file diff --git a/视频抽帧.py b/视频抽帧.py new file mode 100644 index 0000000..58f15ad --- /dev/null +++ b/视频抽帧.py @@ -0,0 +1,190 @@ +import cv2 +import numpy as np +import os +from tqdm import tqdm +from pathlib import Path + +# -------------------------- 1. 用户配置（根据实际情况修改）-------------------------- +# 视频文件路径（绝对路径或相对路径均可，支持mp4/avi/mov等常见格式） +VIDEO_PATH = r"E:\geminicli\yolo\images\44.mp4" # 替换为你的视频路径 + +# 输出根目录（脚本会自动在该目录下创建train/valid/test子文件夹） +OUTPUT_ROOT_DIR = r"D:\DataPreHandler\images" # 替换为你的输出根路径 + +# 各文件夹需抽取的图片数量（总数量=2800+1400+700=4900） +FRAME_COUNTS = { + "train": 28, + "valid": 14, + "test": 7 +} + +# 图片保存格式（建议用jpg，兼容性更好；也可改为png） +SAVE_FORMAT = "jpg" + +# 图片文件名编号位数（如0001.jpg，避免文件名乱序） +FILE_NUM_DIGITS = 4 # 对应最大编号9999，满足4900张需求 + + +# ----------------------------------------------------------------------------------- + + +def check_video_validity(cap): + """检查视频是否能正常读取，并返回视频总帧数和帧率""" + if not cap.isOpened(): + raise ValueError(f"无法打开视频文件！请检查路径：{VIDEO_PATH}") + + # 获取视频总帧数（注意：部分视频可能返回-1，需特殊处理） + total_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) + if total_frames == -1: + # 若无法直接获取总帧数，通过读取最后一帧间接计算 + cap.set(cv2.CAP_PROP_POS_AVI_RATIO, 1) # 跳转到视频末尾 + total_frames = cap.get(cv2.CAP_PROP_POS_FRAMES) + cap.set(cv2.CAP_PROP_POS_FRAMES, 0) # 跳回视频开头 + + total_frames = int(total_frames) + required_total = sum(FRAME_COUNTS.values()) # 需抽取的总帧数 + + # 检查视频总帧数是否满足需求 + if total_frames < required_total: + raise ValueError( + f"视频总帧数不足！\n" + f"视频实际帧数：{total_frames}，需抽取帧数：{required_total}\n" + f"建议更换更长的视频，或减少各文件夹的抽取数量。" + ) + + # 获取视频帧率（仅用于打印信息，不影响抽取逻辑） + fps = cap.get(cv2.CAP_PROP_FPS) + video_duration = total_frames / fps # 视频总时长（秒） + + print(f"✅ 视频信息读取成功：") + print(f" - 视频路径：{VIDEO_PATH}") + print(f" - 总帧数：{total_frames}") + print(f" - 帧率(FPS)：{fps:.1f}") + print(f" - 总时长：{video_duration // 60:.0f}分{video_duration % 60:.1f}秒") + print( + f" - 需抽取总帧数：{required_total}（train:{FRAME_COUNTS['train']}, valid:{FRAME_COUNTS['valid']}, test:{FRAME_COUNTS['test']}）") + + return total_frames + + +def create_output_dirs(): + """创建输出根目录及train/valid/test子文件夹""" + # 转换为Path对象，适配Windows/Linux/macOS路径格式 + output_root = Path(OUTPUT_ROOT_DIR) + subdirs = FRAME_COUNTS.keys() + + for subdir in subdirs: + subdir_path = output_root / subdir + subdir_path.mkdir(parents=True, exist_ok=True) # parents=True创建父目录，exist_ok=True避免已存在时报错 + print(f"📂 输出文件夹已创建/确认：{subdir_path}") + + return {subdir: output_root / subdir for subdir in subdirs} # 返回各子文件夹路径字典 + + +def generate_random_frame_indices(total_frames, required_total): + """生成无重复的随机帧索引（范围：0 ~ total_frames-1）""" + print(f"\n🎲 正在生成{required_total}个无重复随机帧索引...") + # 使用numpy生成无重复随机整数（replace=False确保不重复） + random_indices = np.random.choice( + a=range(total_frames), + size=required_total, + replace=False + ) + # 排序（可选，使抽取的帧按时间顺序保存，不排序则完全随机） + random_indices.sort() + print(f"✅ 随机帧索引生成完成（共{len(random_indices)}个）") + return random_indices + + +def split_indices_by_dataset(random_indices): + """将总随机索引按train/valid/test的数量拆分""" + train_count = FRAME_COUNTS["train"] + valid_count = FRAME_COUNTS["valid"] + + # 拆分逻辑：前N个给train，中间M个给valid，剩余给test + indices_split = { + "train": random_indices[:train_count], + "valid": random_indices[train_count:train_count + valid_count], + "test": random_indices[train_count + valid_count:] + } + + # 验证拆分数量是否正确（避免配置错误） + for dataset, indices in indices_split.items(): + assert len(indices) == FRAME_COUNTS[dataset], \ + f"{dataset}索引拆分错误！预期{FRAME_COUNTS[dataset]}个，实际{len(indices)}个" + + print(f"\n📊 索引拆分完成：") + for dataset, indices in indices_split.items(): + print(f" - {dataset}：{len(indices)}个帧索引（范围：{indices[0]} ~ {indices[-1]}）") + + return indices_split + + +def extract_and_save_frames(cap, indices_split, output_dirs): + """根据拆分后的索引，抽取视频帧并保存到对应文件夹""" + print(f"\n🚀 开始抽取并保存视频帧...") + + for dataset, indices in indices_split.items(): + dataset_dir = output_dirs[dataset] + print(f"\n--- 正在处理 {dataset} 集（共{len(indices)}张）---") + + # 用tqdm显示进度条 + for idx, frame_idx in tqdm(enumerate(indices, 1), total=len(indices), desc=f"{dataset}进度"): + # 跳转到指定帧（关键步骤：确保读取到正确的帧） + cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx) + + # 读取帧（ret为True表示读取成功，frame为帧数据） + ret, frame = cap.read() + if not ret: + print(f"⚠️ 警告：无法读取帧索引{frame_idx}，已跳过该帧") + continue + + # 生成文件名（如train_0001.jpg） + file_name = f"{dataset}_{idx:0{FILE_NUM_DIGITS}d}.{SAVE_FORMAT}" # 0{FILE_NUM_DIGITS}d表示补零到指定位数 + save_path = dataset_dir / file_name + + # 保存帧为图片（cv2.imwrite默认保存为BGR格式，符合图片存储标准） + cv2.imwrite(str(save_path), frame) + + print(f"\n🎉 所有帧抽取与保存完成！") + # 打印最终结果汇总 + print(f"\n📋 结果汇总：") + for dataset, dataset_dir in output_dirs.items(): + # 统计实际保存的图片数量（避免因读取失败导致数量不足） + actual_count = len([f for f in dataset_dir.glob(f"*.{SAVE_FORMAT}")]) + print(f" - {dataset}集：预期{FRAME_COUNTS[dataset]}张，实际保存{actual_count}张，路径：{dataset_dir}") + + +def main(): + try: + # 1. 初始化视频读取器 + cap = cv2.VideoCapture(VIDEO_PATH) + + # 2. 检查视频有效性并获取总帧数 + total_frames = check_video_validity(cap) + + # 3. 创建输出文件夹 + output_dirs = create_output_dirs() + + # 4. 生成无重复随机帧索引 + required_total = sum(FRAME_COUNTS.values()) + random_indices = generate_random_frame_indices(total_frames, required_total) + + # 5. 按数据集拆分索引 + indices_split = split_indices_by_dataset(random_indices) + + # 6. 抽取并保存帧 + extract_and_save_frames(cap, indices_split, output_dirs) + + except Exception as e: + # 捕获所有异常并友好提示 + print(f"\n❌ 脚本执行失败：{str(e)}") + finally: + # 无论是否报错，都关闭视频读取器（释放资源） + if 'cap' in locals() and cap.isOpened(): + cap.release() + print(f"\n🔌 视频资源已释放") + + +if __name__ == "__main__": + main() \ No newline at end of file