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 + )