diff --git a/data_management/batch_download_all.py b/data_management/batch_download_all.py new file mode 100644 index 0000000..21725f3 --- /dev/null +++ b/data_management/batch_download_all.py @@ -0,0 +1,298 @@ +""" +批量下载所有图片到dataset目录 +自动按类别分类到对应文件夹 +""" +import os +import sys +from datetime import datetime +from tqdm import tqdm +import requests + +# 添加当前目录到路径 +sys.path.append(os.path.dirname(os.path.abspath(__file__))) + +from database import MySQLManager, SQLiteManager +from utils import ImageDownloader +from config import IMAGE_URL_PREFIX, DOWNLOAD_TIMEOUT, DOWNLOAD_THREADS + + +def sanitize_folder_name(name: str) -> str: + """清理文件夹名称,移除非法字符""" + # 保留中文字符、字母、数字、括号、方括号、WP代码 + import re + # 移除或替换非法字符 + name = name.replace('/', '_').replace('\\', '_') + name = name.replace(':', '_').replace('*', '_') + name = name.replace('?', '_').replace('"', '_') + name = name.replace('<', '_').replace('>', '_') + name = name.replace('|', '_') + return name.strip() + + +def get_goods_code(goods_name: str) -> str: + """从goods_name中提取WP代码(如果存在)""" + import re + # 匹配WP开头的代码,如 WP2025110700043 + match = re.search(r'WP\d+', goods_name) + if match: + return match.group(0) + return '' + + +def batch_download_all( + datasource: str = '食材', + output_dir: str = None, + max_images: int = None, + skip_existing: bool = True +): + """ + 批量下载所有图片 + + Args: + datasource: 数据源('食材'/'菜品'/'半加工') + output_dir: 输出目录,如果为None则使用dataset目录 + max_images: 最大下载数量(用于测试),None表示下载全部 + skip_existing: 是否跳过已存在的文件 + """ + print("="*80) + print(f"批量下载工具 - 数据源: {datasource}") + print("="*80) + + # 初始化数据库管理器 + print("\n[1/5] 连接数据库...") + mysql_manager = MySQLManager() + sqlite_manager = SQLiteManager() + + # 查询所有图片数据 + print("\n[2/5] 查询所有图片数据...") + all_data = [] + page = 1 + page_size = 1000 # 每次查询1000条 + + while True: + data_list, total_count = mysql_manager.query_images( + datasource=datasource, + page=page, + page_size=page_size + ) + + if not data_list: + break + + all_data.extend(data_list) + print(f" 已查询: {len(all_data)} / {total_count} 张图片") + + if len(all_data) >= total_count: + break + + if max_images and len(all_data) >= max_images: + all_data = all_data[:max_images] + break + + page += 1 + + total_images = len(all_data) + print(f"\n✓ 共找到 {total_images} 张图片") + + if total_images == 0: + print("没有找到图片,退出") + return + + # 确定输出目录 + if output_dir is None: + # 根据数据源确定目录 + base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + if datasource == '食材': + output_dir = os.path.join(base_dir, 'dataset', 'WholeIngredientRecognition', 'train') + elif datasource == '菜品': + output_dir = os.path.join(base_dir, 'dataset', 'DishClassification', 'train') + elif datasource == '半加工': + output_dir = os.path.join(base_dir, 'dataset', 'ProcessedIngredientRecognition', 'train') + else: + output_dir = os.path.join(base_dir, 'dataset', datasource, 'train') + + os.makedirs(output_dir, exist_ok=True) + print(f"\n[3/5] 输出目录: {output_dir}") + + # 按类别分组 + print("\n[4/5] 按类别分组...") + category_groups = {} + for data in all_data: + goods_name = data['goods_name'] + # 清理文件夹名称 + folder_name = sanitize_folder_name(goods_name) + + if folder_name not in category_groups: + category_groups[folder_name] = [] + category_groups[folder_name].append(data) + + print(f"✓ 共 {len(category_groups)} 个类别") + + # 显示类别统计 + print("\n类别统计(前10个):") + sorted_categories = sorted(category_groups.items(), key=lambda x: len(x[1]), reverse=True) + for i, (category, items) in enumerate(sorted_categories[:10], 1): + print(f" {i}. {category}: {len(items)} 张") + + # 下载图片 + print(f"\n[5/5] 开始下载图片(并发数: {DOWNLOAD_THREADS})...") + + success_count = 0 + failed_count = 0 + skipped_count = 0 + error_messages = [] + + # 创建进度条 + pbar = tqdm(total=total_images, desc="下载进度", unit="张") + + def download_single(data: dict) -> tuple[bool, str]: + """下载单张图片""" + try: + goods_name = data['goods_name'] + folder_name = sanitize_folder_name(goods_name) + category_dir = os.path.join(output_dir, folder_name) + os.makedirs(category_dir, exist_ok=True) + + # 构建文件名 + filename = f"img_{data['id']}.jpg" + file_path = os.path.join(category_dir, filename) + + # 检查文件是否已存在 + if skip_existing and os.path.exists(file_path): + return True, "已存在" + + # 下载图片 + image_url = data['image_url'] + if not image_url: + return False, "图片URL为空" + + response = requests.get(image_url, timeout=DOWNLOAD_TIMEOUT, stream=True) + response.raise_for_status() + + # 保存图片 + with open(file_path, 'wb') as f: + for chunk in response.iter_content(chunk_size=8192): + if chunk: + f.write(chunk) + + # 记录下载历史 + sqlite_manager.add_download_record( + datasource=data['datasource'], + goods_id=data.get('goods_id'), + goods_name=goods_name, + image_url=image_url, + local_path=file_path, + file_size=os.path.getsize(file_path) + ) + + return True, None + + except Exception as e: + return False, str(e) + + # 使用线程池下载 + from concurrent.futures import ThreadPoolExecutor, as_completed + + with ThreadPoolExecutor(max_workers=DOWNLOAD_THREADS) as executor: + # 提交所有任务 + futures = {executor.submit(download_single, data): data for data in all_data} + + # 处理结果 + for future in as_completed(futures): + data = futures[future] + try: + success, error = future.result() + if success: + if error == "已存在": + skipped_count += 1 + else: + success_count += 1 + else: + failed_count += 1 + error_messages.append(f"{data['goods_name']} (ID: {data['id']}): {error}") + + pbar.update(1) + + except Exception as e: + failed_count += 1 + error_messages.append(f"{data['goods_name']} (ID: {data['id']}): {str(e)}") + pbar.update(1) + + pbar.close() + + # 打印结果 + print("\n" + "="*80) + print("下载完成!") + print("="*80) + print(f"成功: {success_count} 张") + print(f"跳过: {skipped_count} 张(已存在)") + print(f"失败: {failed_count} 张") + print(f"总计: {total_images} 张") + print(f"\n输出目录: {output_dir}") + + if error_messages: + print(f"\n错误信息(前10条):") + for msg in error_messages[:10]: + print(f" - {msg}") + if len(error_messages) > 10: + print(f" ... 还有 {len(error_messages) - 10} 条错误") + + # 关闭数据库连接 + mysql_manager.close() + + print("\n✓ 完成!") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description='批量下载所有图片到dataset目录') + parser.add_argument('--datasource', type=str, default='食材', + choices=['食材', '菜品', '半加工', 'all'], + help='数据源(默认: 食材,使用 all 下载所有三种)') + parser.add_argument('--output', type=str, default=None, + help='输出目录(默认: dataset/{datasource}/train)') + parser.add_argument('--max', type=int, default=None, + help='最大下载数量(用于测试)') + parser.add_argument('--no-skip', action='store_true', + help='不跳过已存在的文件(重新下载)') + + args = parser.parse_args() + + # 如果选择 all,则依次下载所有三种数据源 + if args.datasource == 'all': + datasources = ['食材', '菜品', '半加工'] + print("\n" + "="*80) + print("将依次下载所有三种数据源:") + print(" 1. 食材 -> WholeIngredientRecognition") + print(" 2. 菜品 -> DishClassification") + print(" 3. 半加工 -> ProcessedIngredientRecognition") + print("="*80 + "\n") + + for i, ds in enumerate(datasources, 1): + print(f"\n{'='*80}") + print(f"[{i}/3] 开始下载: {ds}") + print(f"{'='*80}\n") + + batch_download_all( + datasource=ds, + output_dir=args.output, + max_images=args.max, + skip_existing=not args.no_skip + ) + + if i < len(datasources): + print(f"\n等待5秒后继续下一个数据源...") + import time + time.sleep(5) + + print("\n" + "="*80) + print("✓ 所有数据源下载完成!") + print("="*80) + else: + batch_download_all( + datasource=args.datasource, + output_dir=args.output, + max_images=args.max, + skip_existing=not args.no_skip + ) diff --git a/data_management/config.py b/data_management/config.py index f92fdec..e34d863 100644 --- a/data_management/config.py +++ b/data_management/config.py @@ -9,7 +9,7 @@ MYSQL_CONFIG = { 'port': 3308, 'user': 'root', 'password': 'Aa135790123', - 'database': 'yx_terminal', + 'database': 'yx_terminal_uat', 'charset': 'utf8mb4' } @@ -21,7 +21,7 @@ TABLE_MAPPING = { } # ======================== 图片URL前缀 ======================== -IMAGE_URL_PREFIX = 'http://dev.yixiong-tech.com:8081' +IMAGE_URL_PREFIX = 'https://dev.yixiong-tech.com:8081' # ======================== SQLite下载历史数据库 ======================== # 存储在项目的data_management目录下 diff --git a/data_management/count_images.py b/data_management/count_images.py new file mode 100644 index 0000000..a64927d --- /dev/null +++ b/data_management/count_images.py @@ -0,0 +1,171 @@ +""" +统计每个类别的图片数量(包含train/val/test) +""" +import os +from pathlib import Path +from collections import defaultdict + + +def count_images_by_class(dataset_root: str, output_file: str = None): + """ + 统计每个类别的图片数量(包含train/val/test) + + Args: + dataset_root: 数据集根目录(包含train/val/test目录) + output_file: 输出文件路径(可选,如果提供则保存到文件) + """ + dataset_root = Path(dataset_root) + train_dir = dataset_root / 'train' + val_dir = dataset_root / 'val' + test_dir = dataset_root / 'test' + + if not train_dir.exists(): + print(f"错误: 训练目录不存在: {train_dir}") + return + + print("="*80) + print("每个类别的图片数量统计(包含train/val/test)") + print("="*80) + + class_counts = {} + class_counts_detail = {} # 详细统计:train/val/test分别的数量 + total_images = 0 + total_train = 0 + total_val = 0 + total_test = 0 + + # 获取所有类别文件夹(从train目录) + class_folders = sorted([f for f in train_dir.iterdir() if f.is_dir()]) + + # 也检查val和test目录中是否有train中没有的类别 + all_class_names = set() + for folder in class_folders: + all_class_names.add(folder.name) + + if val_dir.exists(): + for folder in val_dir.iterdir(): + if folder.is_dir(): + all_class_names.add(folder.name) + + if test_dir.exists(): + for folder in test_dir.iterdir(): + if folder.is_dir(): + all_class_names.add(folder.name) + + # 统计每个类别的图片数量 + for class_name in sorted(all_class_names): + train_count = 0 + val_count = 0 + test_count = 0 + + # 统计train目录 + train_class_dir = train_dir / class_name + if train_class_dir.exists(): + for ext in ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']: + train_count += len(list(train_class_dir.glob(ext))) + + # 统计val目录 + val_class_dir = val_dir / class_name + if val_class_dir.exists(): + for ext in ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']: + val_count += len(list(val_class_dir.glob(ext))) + + # 统计test目录 + test_class_dir = test_dir / class_name + if test_class_dir.exists(): + for ext in ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']: + test_count += len(list(test_class_dir.glob(ext))) + + total = train_count + val_count + test_count + class_counts[class_name] = total + class_counts_detail[class_name] = { + 'train': train_count, + 'val': val_count, + 'test': test_count, + 'total': total + } + total_images += total + total_train += train_count + total_val += val_count + total_test += test_count + + # 按数量排序 + sorted_classes = sorted(class_counts.items(), key=lambda x: x[1], reverse=True) + + # 修复文件开头的注释 + if output_file: + # 确保sorted_classes已定义 + pass + + # 打印结果 + print(f"\n总类别数: {len(class_counts)}") + print(f"总图片数: {total_images} (训练集: {total_train}, 验证集: {total_val}, 测试集: {total_test})") + print(f"平均每类: {total_images // len(class_counts) if len(class_counts) > 0 else 0} 张\n") + print("-"*100) + print(f"{'类别名称':<50} {'总计':<8} {'训练':<8} {'验证':<8} {'测试':<8} {'状态':<10}") + print("-"*100) + + # 统计样本不足的类别(假设阈值是10) + min_threshold = 10 + insufficient_count = 0 + + for class_name, count in sorted_classes: + detail = class_counts_detail[class_name] + status = "⚠样本不足" if count < min_threshold else "✓" + if count < min_threshold: + insufficient_count += 1 + print(f"{class_name:<50} {detail['total']:<8} {detail['train']:<8} {detail['val']:<8} {detail['test']:<8} {status:<10}") + + print("-"*100) + print(f"\n样本不足的类别数(<{min_threshold}张): {insufficient_count}") + print(f"样本充足的类别数(>={min_threshold}张): {len(class_counts) - insufficient_count}") + + # 保存到文件(如果指定) + if output_file: + output_path = Path(output_file) + with open(output_path, 'w', encoding='utf-8') as f: + f.write("="*100 + "\n") + f.write("每个类别的图片数量统计(包含train/val/test)\n") + f.write("="*100 + "\n\n") + f.write(f"总类别数: {len(class_counts)}\n") + f.write(f"总图片数: {total_images} (训练集: {total_train}, 验证集: {total_val}, 测试集: {total_test})\n") + f.write(f"平均每类: {total_images // len(class_counts) if len(class_counts) > 0 else 0} 张\n\n") + f.write("-"*100 + "\n") + f.write(f"{'类别名称':<50} {'总计':<8} {'训练':<8} {'验证':<8} {'测试':<8} {'状态':<10}\n") + f.write("-"*100 + "\n") + + for class_name, count in sorted_classes: + detail = class_counts_detail[class_name] + status = "⚠样本不足" if count < min_threshold else "✓" + f.write(f"{class_name:<50} {detail['total']:<8} {detail['train']:<8} {detail['val']:<8} {detail['test']:<8} {status:<10}\n") + + f.write("-"*100 + "\n") + f.write(f"\n样本不足的类别数(<{min_threshold}张): {insufficient_count}\n") + f.write(f"样本充足的类别数(>={min_threshold}张): {len(class_counts) - insufficient_count}\n") + + print(f"\n✓ 统计结果已保存到: {output_path}") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description='统计每个类别的图片数量') + parser.add_argument('--dataset', type=str, + default='dataset/WholeIngredientRecognition', + help='数据集根目录(默认: dataset/WholeIngredientRecognition)') + parser.add_argument('--output', type=str, default=None, + help='输出文件路径(可选,保存统计结果到文件)') + + args = parser.parse_args() + + # 获取项目根目录 + script_dir = Path(__file__).parent + project_root = script_dir.parent + + # 处理相对路径 + if not os.path.isabs(args.dataset): + dataset_path = project_root / args.dataset + else: + dataset_path = Path(args.dataset) + + count_images_by_class(str(dataset_path), args.output) 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) diff --git a/data_management/split_dataset.py b/data_management/split_dataset.py new file mode 100644 index 0000000..a4b96c3 --- /dev/null +++ b/data_management/split_dataset.py @@ -0,0 +1,367 @@ +"""cui +数据集分割脚本 +1. 将train目录下的图片按70:15:15分割到train/val/test +2. 对样本量过少的类别进行标记 +""" +import os +import shutil +import random +from pathlib import Path +from typing import List, Tuple + + +def split_dataset( + dataset_root: str, + train_ratio: float = 0.70, + val_ratio: float = 0.15, + test_ratio: float = 0.15, + min_samples_threshold: int = 10, + mark_insufficient: bool = True, + max_total: int = 120, + seed: int = 42 +): + """ + 分割数据集 + + Args: + dataset_root: 数据集根目录(包含train目录) + train_ratio: 训练集比例(默认0.70) + val_ratio: 验证集比例(默认0.15) + test_ratio: 测试集比例(默认0.15) + min_samples_threshold: 最小样本数阈值,低于此值的类别会被标记(默认10) + mark_insufficient: 是否在文件夹名称后添加"样本数太少"标记(默认True) + max_total: 如果总数超过100,限制到max_total张(默认120) + seed: 随机种子(默认42) + """ + # 验证比例 + if abs(train_ratio + val_ratio + test_ratio - 1.0) > 0.001: + raise ValueError(f"比例总和必须为1.0,当前为: {train_ratio + val_ratio + test_ratio}") + + dataset_root = Path(dataset_root) + train_dir = dataset_root / 'train' + val_dir = dataset_root / 'val' + test_dir = dataset_root / 'test' + + if not train_dir.exists(): + raise ValueError(f"训练目录不存在: {train_dir}") + + # 创建val和test目录 + val_dir.mkdir(parents=True, exist_ok=True) + test_dir.mkdir(parents=True, exist_ok=True) + + # 设置随机种子 + random.seed(seed) + + print("="*80) + print(f"数据集分割工具") + print("="*80) + print(f"数据集根目录: {dataset_root}") + print(f"分割比例: 训练集={train_ratio*100:.0f}%, 验证集={val_ratio*100:.0f}%, 测试集={test_ratio*100:.0f}%") + print(f"最小样本数阈值: {min_samples_threshold}") + print(f"样本不足标记: {'启用' if mark_insufficient else '禁用'}") + print("="*80 + "\n") + + # 检查val和test目录是否已有数据 + val_has_data = val_dir.exists() and any(val_dir.iterdir()) + test_has_data = test_dir.exists() and any(test_dir.iterdir()) + + if val_has_data or test_has_data: + print("⚠ 检测到验证集或测试集已有数据") + print("将检查所有类别的比例,如果不符合70:15:15将重新分配\n") + + # 获取所有类别文件夹 + class_folders = [f for f in train_dir.iterdir() if f.is_dir()] + total_classes = len(class_folders) + + print(f"找到 {total_classes} 个类别\n") + + stats = { + 'total_classes': total_classes, + 'processed_classes': 0, + 'skipped_classes': 0, + 'sufficient_samples': 0, + 'insufficient_samples': 0, + 'total_images': 0, + 'train_images': 0, + 'val_images': 0, + 'test_images': 0, + 'marked_classes': [] + } + + # 处理每个类别 + for class_folder in sorted(class_folders): + class_name = class_folder.name + + # 收集所有图片(从train/val/test) + train_images = [] + for ext in ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']: + train_images.extend(list(class_folder.glob(ext))) + + val_class_dir = val_dir / class_name + val_images = [] + if val_class_dir.exists(): + for ext in ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']: + val_images.extend(list(val_class_dir.glob(ext))) + + test_class_dir = test_dir / class_name + test_images = [] + if test_class_dir.exists(): + for ext in ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']: + test_images.extend(list(test_class_dir.glob(ext))) + + # 合并所有图片 + all_images = train_images + val_images + test_images + num_images = len(all_images) + stats['total_images'] += num_images + + if num_images == 0: + print(f"⚠ {class_name}: 0张图片,跳过") + continue + + # 检查当前比例是否正确 + current_train = len(train_images) + current_val = len(val_images) + current_test = len(test_images) + + # 计算期望的数量 + expected_train = int(num_images * train_ratio) + expected_val = int(num_images * val_ratio) + expected_test = num_images - expected_train - expected_val + + # 检查比例是否接近(允许1-2张的误差) + ratio_correct = ( + abs(current_train - expected_train) <= 2 and + abs(current_val - expected_val) <= 2 and + abs(current_test - expected_test) <= 2 + ) + + # 如果总数超过100,即使比例正确也需要限制 + if ratio_correct and num_images > 0 and num_images <= 100: + # 比例正确且总数<=100,跳过 + stats['skipped_classes'] += 1 + is_insufficient = num_images < min_samples_threshold + if is_insufficient: + stats['insufficient_samples'] += 1 + else: + stats['sufficient_samples'] += 1 + print(f"✓ {class_name}: 总数={num_images}, 比例正确 (train={current_train}, val={current_val}, test={current_test}),跳过") + continue + + # 如果比例正确但总数>100,需要限制 + if ratio_correct and num_images > 100: + print(f"处理 {class_name}: 总数={num_images} > 100, 比例正确但需要限制到{max_total}张") + stats['processed_classes'] += 1 + # 随机选择max_total张 + random.shuffle(all_images) + image_files = all_images[:max_total] + num_images = len(image_files) + # 重新计算期望数量 + expected_train = int(num_images * train_ratio) + expected_val = int(num_images * val_ratio) + expected_test = num_images - expected_train - expected_val + print(f" 将重新分配为: train={expected_train}, val={expected_val}, test={expected_test}") + elif not ratio_correct: + # 比例不正确,需要重新分配 + stats['processed_classes'] += 1 + print(f"处理 {class_name}: 总数={num_images}, 当前比例不正确 (train={current_train}, val={current_val}, test={current_test})") + + # 如果总数超过100,限制到max_total + if num_images > 100: + print(f" 总数={num_images} > 100, 限制到{max_total}张") + random.shuffle(all_images) + image_files = all_images[:max_total] + num_images = len(image_files) + # 重新计算期望数量 + expected_train = int(num_images * train_ratio) + expected_val = int(num_images * val_ratio) + expected_test = num_images - expected_train - expected_val + else: + image_files = all_images + print(f" 将重新分配为: train={expected_train}, val={expected_val}, test={expected_test}") + + + # 检查样本数是否足够 + is_insufficient = num_images < min_samples_threshold + + if is_insufficient: + stats['insufficient_samples'] += 1 + if mark_insufficient: + # 标记样本不足的类别 + marked_name = f"{class_name}_样本数太少" + stats['marked_classes'].append((class_name, num_images, marked_name)) + else: + stats['sufficient_samples'] += 1 + + # 确定最终的类别名称(如果需要标记) + final_class_name = f"{class_name}_样本数太少" if (is_insufficient and mark_insufficient) else class_name + + # 如果还没打乱(总数<=100的情况),现在打乱 + if num_images == len(all_images): + random.shuffle(image_files) + + # 计算分割点 + n_train = int(num_images * train_ratio) + n_val = int(num_images * val_ratio) + # 剩余的全部给test,确保所有图片都被分配 + + # 分割图片(重新分配) + new_train_images = image_files[:n_train] + new_val_images = image_files[n_train:n_train + n_val] + new_test_images = image_files[n_train + n_val:] + + # 创建目标目录 + train_target = train_dir / final_class_name + val_target = val_dir / final_class_name + test_target = test_dir / final_class_name + + train_target.mkdir(exist_ok=True) + val_target.mkdir(exist_ok=True) + test_target.mkdir(exist_ok=True) + + # 先移动图片到正确的位置(在删除之前) + moved_to_train = 0 + for img in new_train_images: + if not img.exists(): + continue + target = train_target / img.name + if img.parent != train_target: + if target.exists(): + target.unlink() + try: + shutil.move(str(img), str(target)) + moved_to_train += 1 + except Exception as e: + print(f" ⚠ 移动失败 {img}: {e}") + + moved_to_val = 0 + for img in new_val_images: + if not img.exists(): + continue + target = val_target / img.name + if img.parent != val_target: + if target.exists(): + target.unlink() + try: + shutil.move(str(img), str(target)) + moved_to_val += 1 + except Exception as e: + print(f" ⚠ 移动失败 {img}: {e}") + + moved_to_test = 0 + for img in new_test_images: + if not img.exists(): + continue + target = test_target / img.name + if img.parent != test_target: + if target.exists(): + target.unlink() + try: + shutil.move(str(img), str(target)) + moved_to_test += 1 + except Exception as e: + print(f" ⚠ 移动失败 {img}: {e}") + + # 然后清理所有目录(删除不在新分配列表中的图片) + # 收集所有原始图片路径 + original_images = set(train_images + val_images + test_images) + new_images = set(new_train_images + new_val_images + new_test_images) + images_to_delete = original_images - new_images + + for img in images_to_delete: + if img.exists(): + try: + img.unlink() + except Exception as e: + pass + + # 处理样本不足标记(如果需要) + if mark_insufficient and is_insufficient: + # 需要重命名train目录的文件夹 + if class_folder != train_target: + # 如果原文件夹还存在且不为空,需要移动剩余文件 + if class_folder.exists() and any(class_folder.iterdir()): + for item in class_folder.iterdir(): + if item.is_file(): + shutil.move(str(item), train_target / item.name) + # 删除原文件夹(如果为空) + try: + if class_folder.exists() and not any(class_folder.iterdir()): + class_folder.rmdir() + except Exception as e: + print(f" 警告: 删除原文件夹失败: {e}") + + # 更新统计 + stats['train_images'] += len(new_train_images) + stats['val_images'] += len(new_val_images) + stats['test_images'] += len(new_test_images) + + # 打印进度 + status = "⚠样本不足" if is_insufficient else "✓" + print(f"{status} {class_name}: 总计={num_images}, 训练={len(new_train_images)}, 验证={len(new_val_images)}, 测试={len(new_test_images)}") + + # 打印总结 + print("\n" + "="*80) + print("分割完成!") + print("="*80) + print(f"总类别数: {stats['total_classes']}") + print(f"已处理: {stats['processed_classes']} 个类别(比例不正确,已重新分配)") + print(f"已跳过: {stats['skipped_classes']} 个类别(比例正确)") + print(f"样本充足: {stats['sufficient_samples']} 个类别") + print(f"样本不足: {stats['insufficient_samples']} 个类别") + print(f"\n总图片数: {stats['total_images']}") + print(f"训练集: {stats['train_images']} 张 ({stats['train_images']/stats['total_images']*100:.1f}%)") + print(f"验证集: {stats['val_images']} 张 ({stats['val_images']/stats['total_images']*100:.1f}%)") + print(f"测试集: {stats['test_images']} 张 ({stats['test_images']/stats['total_images']*100:.1f}%)") + + if stats['marked_classes']: + print(f"\n已标记的类别(样本数 < {min_samples_threshold}):") + for old_name, count, new_name in stats['marked_classes']: + print(f" {old_name} ({count}张) -> {new_name}") + + 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('--train-ratio', type=float, default=0.70, + help='训练集比例(默认: 0.70)') + parser.add_argument('--val-ratio', type=float, default=0.15, + help='验证集比例(默认: 0.15)') + parser.add_argument('--test-ratio', type=float, default=0.15, + help='测试集比例(默认: 0.15)') + parser.add_argument('--min-samples', type=int, default=10, + help='最小样本数阈值,低于此值会被标记(默认: 10)') + parser.add_argument('--no-mark', action='store_true', + help='不标记样本不足的类别') + parser.add_argument('--seed', type=int, default=42, + help='随机种子(默认: 42)') + parser.add_argument('--max-total', type=int, default=120, + help='如果总数超过100,限制到max_total张(默认: 120)') + + args = parser.parse_args() + + # 获取项目根目录 + script_dir = Path(__file__).parent + project_root = script_dir.parent + + # 处理相对路径 + if not os.path.isabs(args.dataset): + dataset_path = project_root / args.dataset + else: + dataset_path = Path(args.dataset) + + split_dataset( + dataset_root=str(dataset_path), + train_ratio=args.train_ratio, + val_ratio=args.val_ratio, + test_ratio=args.test_ratio, + min_samples_threshold=args.min_samples, + mark_insufficient=not args.no_mark, + max_total=args.max_total, + seed=args.seed + ) diff --git a/faiss_vector_db/build_faiss_index.py b/faiss_vector_db/build_faiss_index.py index adf6315..c9f6ca7 100644 --- a/faiss_vector_db/build_faiss_index.py +++ b/faiss_vector_db/build_faiss_index.py @@ -556,7 +556,14 @@ def main(): # MODEL_PATH = "../model/embedding_20251011_133653/best_embedding_model.pth" # MODEL_PATH = "../model/ProcessedIngredientRecognition/embedding_20251103_172012/best_embedding_model.pth" # MODEL_PATH = "../model/WholeIngredientRecognition/cosface_20251113_160103/best_cosface_model.pth" - MODEL_PATH = "../model/WholeIngredientRecognition/grid_search_20251204_140816/model_s64.0_m0.4.pth" + #老索引 + #MODEL_PATH = "../model/WholeIngredientRecognition/grid_search_20251204_140816/model_s64.0_m0.4.pth" + #新索引 + #MODEL_PATH = "../model/WholeIngredientRecognition/cosface_20260118_220318/best_cosface_model.pth" + #新索引 + #MODEL_PATH = "../model/WholeIngredientRecognition/cosface_20260121_133207/best_cosface_model.pth" + #grid_search最佳模型 s=68, m=0.42 1.22 + MODEL_PATH = "../model/WholeIngredientRecognition/grid_search_20260122_112545/best_model_s68.0_m0.42.pth" # MODEL_PATH = "../model/DishClassification/grid_search_20251121_102723/model_s56.0_m0.4.pth" # MODEL_PATH = "../model/DishClassification/cosface_20251111_153649/best_cosface_model.pth" # TRAIN_DIR = "../dataset/ProcessedIngredientRecognition/train" diff --git a/toAndroid/toAndroidEmbedding.py b/toAndroid/toAndroidEmbedding.py index 4bdafbf..43195f8 100644 --- a/toAndroid/toAndroidEmbedding.py +++ b/toAndroid/toAndroidEmbedding.py @@ -19,7 +19,10 @@ def main(): base_model = create_mobile_resnet50_embedding(embedding_dim=512, pretrained=True) # model_path = "../model/DishClassification/grid_search_20251121_102723/model_s56.0_m0.4.pth" # model_path = "../model/WholeIngredientRecognition/cosface_20251106_134718/best_cosface_model.pth" - model_path = "../model/WholeIngredientRecognition/grid_search_20251204_140816/model_s64.0_m0.4.pth" + #原本路径,原本模型 + #model_path = "../model/WholeIngredientRecognition/grid_search_20251204_140816/model_s64.0_m0.4.pth" + #新路径,新模型 (模型不可用) + model_path = "../model/WholeIngredientRecognition/cosface_20260121_133207/best_cosface_model.pth" # model_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/best_embedding_model.pth" if not os.path.exists(model_path): @@ -94,12 +97,15 @@ def main(): # 保存模型,一定不要乱改,防止模型覆盖!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! # output_path = "../model/DishClassification/grid_search_20251121_102723/best_embedding_model_mobile.pt" - output_path = "../model/WholeIngredientRecognition/grid_search_20251204_140816/best_embedding_model_mobile.pt" + #原本路径 + #output_path = "../model/WholeIngredientRecognition/grid_search_20251204_140816/best_embedding_model_mobile.pt" + #新路径 + output_path = "../model/WholeIngredientRecognition/cosface_20260121_133207/best_cosface_model_mobile.pt" # output_path = "../model/WholeIngredientRecognition/cosface_20251106_134718/best_embedding_model_mobile.pt" # output_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/best_embedding_model_mobile.pt" traced_model.save(output_path) print(f"✓ TorchScript模型保存成功: {output_path}") - + # 验证保存的模型 loaded_model = torch.jit.load(output_path) with torch.no_grad():