""" 批量下载所有图片到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 )