将后端的图片下载到数据集中

2026-02-06 02:57:09 -05:00
parent f60c4d9883
commit de106f8ecb
1 changed files with 298 additions and 0 deletions
@@ -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
+        )