diff --git a/data_management/remove_duplicate_images.py b/data_management/remove_duplicate_images.py new file mode 100644 index 0000000..325eadf --- /dev/null +++ b/data_management/remove_duplicate_images.py @@ -0,0 +1,245 @@ +""" +检查并删除重复的照片 +通过文件内容（MD5）或文件名来识别重复 +""" +import os +import hashlib +from pathlib import Path +from collections import defaultdict +import shutil + + +def calculate_md5(file_path: Path) -> str: + """计算文件的MD5值""" + hash_md5 = hashlib.md5() + with open(file_path, "rb") as f: + for chunk in iter(lambda: f.read(4096), b""): + hash_md5.update(chunk) + return hash_md5.hexdigest() + + +def find_duplicates_by_content(dataset_root: str, dry_run: bool = True): + """ + 通过文件内容（MD5）查找重复照片 + + Args: + dataset_root: 数据集根目录 + dry_run: 是否只是预览，不实际删除 + """ + dataset_root = Path(dataset_root) + + print("="*80) + print("检查重复照片（基于文件内容MD5）") + print("="*80) + print(f"数据集根目录: {dataset_root}") + print(f"模式: {'预览模式（不实际删除）' if dry_run else '执行模式'}") + print("="*80 + "\n") + + # 存储所有文件的MD5值 + md5_to_files = defaultdict(list) + + # 遍历所有图片文件 + image_extensions = ['.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG'] + total_files = 0 + + print("正在计算文件MD5值...") + for root, dirs, files in os.walk(dataset_root): + for file in files: + if any(file.endswith(ext) for ext in image_extensions): + file_path = Path(root) / file + try: + md5_hash = calculate_md5(file_path) + md5_to_files[md5_hash].append(file_path) + total_files += 1 + if total_files % 100 == 0: + print(f" 已处理: {total_files} 个文件...") + except Exception as e: + print(f" ⚠ 处理文件失败 {file_path}: {e}") + + print(f"\n✓ 共处理 {total_files} 个文件\n") + + # 找出重复的文件 + duplicates = {md5: files for md5, files in md5_to_files.items() if len(files) > 1} + + if not duplicates: + print("✓ 没有发现重复的照片") + return + + print(f"发现 {len(duplicates)} 组重复照片:\n") + + total_duplicates = 0 + total_to_delete = 0 + + for md5_hash, files in sorted(duplicates.items(), key=lambda x: len(x[1]), reverse=True): + # 按路径排序，保留第一个（通常保留train目录的） + files_sorted = sorted(files) + + # 优先保留train目录的文件 + keep_file = None + for f in files_sorted: + if 'train' in str(f): + keep_file = f + break + + if keep_file is None: + keep_file = files_sorted[0] # 如果没有train目录的，保留第一个 + + files_to_delete = [f for f in files_sorted if f != keep_file] + + print(f"MD5: {md5_hash[:8]}...") + print(f" 保留: {keep_file}") + for f in files_to_delete: + print(f" 删除: {f}") + + total_duplicates += len(files) + total_to_delete += len(files_to_delete) + + # 实际删除 + if not dry_run: + for f in files_to_delete: + try: + f.unlink() + print(f" ✓ 已删除: {f.name}") + except Exception as e: + print(f" ✗ 删除失败: {e}") + print() + + print("="*80) + print("统计:") + print("="*80) + print(f"重复文件组数: {len(duplicates)}") + print(f"重复文件总数: {total_duplicates}") + print(f"将删除文件数: {total_to_delete}") + print(f"将保留文件数: {len(duplicates)}") + + if dry_run: + print("\n使用 --execute 参数来实际执行删除操作") + else: + print("\n✓ 删除完成!") + + +def find_duplicates_by_name(dataset_root: str, dry_run: bool = True): + """ + 通过文件名查找重复照片（同一文件名出现在多个目录） + + Args: + dataset_root: 数据集根目录 + dry_run: 是否只是预览，不实际删除 + """ + dataset_root = Path(dataset_root) + + print("="*80) + print("检查重复照片（基于文件名）") + print("="*80) + print(f"数据集根目录: {dataset_root}") + print(f"模式: {'预览模式（不实际删除）' if dry_run else '执行模式'}") + print("="*80 + "\n") + + # 存储文件名到文件路径的映射 + name_to_files = defaultdict(list) + + # 遍历所有图片文件 + image_extensions = ['.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG'] + + for root, dirs, files in os.walk(dataset_root): + for file in files: + if any(file.endswith(ext) for ext in image_extensions): + file_path = Path(root) / file + name_to_files[file].append(file_path) + + # 找出重复的文件名 + duplicates = {name: files for name, files in name_to_files.items() if len(files) > 1} + + if not duplicates: + print("✓ 没有发现同名重复的照片") + return + + print(f"发现 {len(duplicates)} 组同名重复照片:\n") + + total_to_delete = 0 + + for file_name, files in sorted(duplicates.items(), key=lambda x: len(x[1]), reverse=True): + # 按路径排序，优先保留train目录的 + files_sorted = sorted(files) + + keep_file = None + for f in files_sorted: + if 'train' in str(f): + keep_file = f + break + + if keep_file is None: + keep_file = files_sorted[0] + + files_to_delete = [f for f in files_sorted if f != keep_file] + + print(f"文件名: {file_name}") + print(f" 保留: {keep_file}") + for f in files_to_delete: + print(f" 删除: {f}") + + total_to_delete += len(files_to_delete) + + # 实际删除 + if not dry_run: + for f in files_to_delete: + try: + f.unlink() + print(f" ✓ 已删除: {f.name}") + except Exception as e: + print(f" ✗ 删除失败: {e}") + print() + + print("="*80) + print(f"将删除文件数: {total_to_delete}") + + if dry_run: + print("\n使用 --execute 参数来实际执行删除操作") + else: + print("\n✓ 删除完成!") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description='检查并删除重复的照片') + parser.add_argument('--dataset', type=str, + default='dataset/WholeIngredientRecognition', + help='数据集根目录（默认: dataset/WholeIngredientRecognition）') + parser.add_argument('--method', type=str, + choices=['content', 'name', 'both'], + default='both', + help='检查方法: content(基于MD5), name(基于文件名), both(两种都检查)') + parser.add_argument('--dry-run', action='store_true', + help='预览模式，不实际删除') + parser.add_argument('--execute', action='store_true', + help='实际执行删除（需要明确指定）') + + args = parser.parse_args() + + # 获取项目根目录 + script_dir = Path(__file__).parent + project_root = script_dir.parent + + # 处理相对路径 + if not os.path.isabs(args.dataset): + if args.dataset.startswith('../'): + dataset_path = (script_dir / args.dataset).resolve() + else: + dataset_path = (project_root / args.dataset).resolve() + else: + dataset_path = Path(args.dataset).resolve() + + # 如果没有指定execute，默认是dry-run + if not args.execute and not args.dry_run: + print("⚠ 警告: 默认是预览模式，不会实际删除") + print("使用 --execute 参数来实际执行删除操作\n") + args.dry_run = True + + if args.method in ['content', 'both']: + find_duplicates_by_content(str(dataset_path), dry_run=args.dry_run) + if args.method == 'both': + print("\n" + "="*80 + "\n") + + if args.method in ['name', 'both']: + find_duplicates_by_name(str(dataset_path), dry_run=args.dry_run)