Files
FoodClassifier/data_management/remove_duplicate_images.py
2026-02-06 02:59:10 -05:00

246 lines
8.2 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
检查并删除重复的照片
通过文件内容(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)