master #1

Merged
giteaadmin merged 6 commits from (deleted):master into master 2026-02-09 16:01:24 +08:00
7 changed files with 1100 additions and 6 deletions
+298
View File
@@ -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
)
+2 -2
View File
@@ -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目录下
+171
View File
@@ -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)
+245
View File
@@ -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)
+367
View File
@@ -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
)
+8 -1
View File
@@ -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"
+8 -2
View File
@@ -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,7 +97,10 @@ 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)