316 lines
9.6 KiB
Python
316 lines
9.6 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""
|
|||
|
|
收集并合并 xlsx 文件
|
|||
|
|
|
|||
|
|
功能：
|
|||
|
|
1. 从 data/batch_output 子文件夹收集 results.xlsx（图片分析结果）
|
|||
|
|
2. 与 data/data_all 中对应的原始数据（{name}_text_img.xlsx）合并
|
|||
|
|
3. 通过图片名关联两个数据源
|
|||
|
|
4. 保存合并后的文件到目标目录
|
|||
|
|
|
|||
|
|
用法：
|
|||
|
|
python3 collect_xlsx.py # 默认合并并输出
|
|||
|
|
python3 collect_xlsx.py -o ../data/merged # 指定输出目录
|
|||
|
|
python3 collect_xlsx.py --no-merge # 不合并，只复制
|
|||
|
|
python3 collect_xlsx.py -n # 预览模式
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import argparse
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import Optional, Tuple, List
|
|||
|
|
|
|||
|
|
import pandas as pd
|
|||
|
|
|
|||
|
|
|
|||
|
|
def extract_image_name(path: str) -> str:
|
|||
|
|
"""
|
|||
|
|
从完整路径提取图片文件名
|
|||
|
|
|
|||
|
|
支持 Windows 和 Unix 路径格式
|
|||
|
|
"""
|
|||
|
|
if pd.isna(path):
|
|||
|
|
return ""
|
|||
|
|
|
|||
|
|
path_str = str(path).strip()
|
|||
|
|
if not path_str:
|
|||
|
|
return ""
|
|||
|
|
|
|||
|
|
# 同时处理 Windows (\) 和 Unix (/) 路径分隔符
|
|||
|
|
# 先统一替换为 /，再提取文件名
|
|||
|
|
normalized = path_str.replace("\\", "/")
|
|||
|
|
filename = normalized.split("/")[-1]
|
|||
|
|
|
|||
|
|
return filename
|
|||
|
|
|
|||
|
|
|
|||
|
|
def merge_xlsx_files(
|
|||
|
|
results_file: Path,
|
|||
|
|
original_file: Path,
|
|||
|
|
results_image_col: str = "image_name",
|
|||
|
|
original_image_cols: list = None,
|
|||
|
|
original_text_col: str = "文本"
|
|||
|
|
) -> Tuple[pd.DataFrame, dict]:
|
|||
|
|
"""
|
|||
|
|
合并分析结果和原始数据
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
results_file: 分析结果文件 (batch_output/.../results.xlsx)
|
|||
|
|
original_file: 原始数据文件 (data_all/..._text_img.xlsx)
|
|||
|
|
results_image_col: 结果文件中的图片名列
|
|||
|
|
original_image_cols: 原始文件中可能的图片路径列（按优先级）
|
|||
|
|
original_text_col: 原始文件中的文本列
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
合并后的 DataFrame 和统计信息
|
|||
|
|
"""
|
|||
|
|
if original_image_cols is None:
|
|||
|
|
original_image_cols = ["图片_新", "图片", "图片链接"]
|
|||
|
|
|
|||
|
|
# 读取文件
|
|||
|
|
results_df = pd.read_excel(results_file)
|
|||
|
|
original_df = pd.read_excel(original_file)
|
|||
|
|
|
|||
|
|
stats = {
|
|||
|
|
"results_rows": len(results_df),
|
|||
|
|
"original_rows": len(original_df),
|
|||
|
|
"merged_rows": 0,
|
|||
|
|
"unmatched_results": 0,
|
|||
|
|
"original_columns_added": [],
|
|||
|
|
"image_col_used": None
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 找到可用的图片列
|
|||
|
|
image_col = None
|
|||
|
|
for col in original_image_cols:
|
|||
|
|
if col in original_df.columns:
|
|||
|
|
image_col = col
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
if image_col is None:
|
|||
|
|
raise ValueError(f"原始文件中未找到图片列，尝试过: {original_image_cols}")
|
|||
|
|
|
|||
|
|
stats["image_col_used"] = image_col
|
|||
|
|
|
|||
|
|
# 从原始数据提取图片名作为关联键
|
|||
|
|
original_df["_image_name"] = original_df[image_col].apply(extract_image_name)
|
|||
|
|
|
|||
|
|
# 去重：原始数据可能有重复图片，保留第一条
|
|||
|
|
original_dedup = original_df.drop_duplicates(subset=["_image_name"], keep="first")
|
|||
|
|
|
|||
|
|
# 确定要添加的原始数据列（排除图片路径列和临时列）
|
|||
|
|
exclude_cols = set(original_image_cols + ["_image_name"])
|
|||
|
|
original_cols_to_add = [col for col in original_df.columns
|
|||
|
|
if col not in exclude_cols
|
|||
|
|
and col not in results_df.columns]
|
|||
|
|
|
|||
|
|
stats["original_columns_added"] = original_cols_to_add
|
|||
|
|
|
|||
|
|
# 创建图片名到原始数据的映射
|
|||
|
|
original_map = original_dedup.set_index("_image_name")[original_cols_to_add].to_dict("index")
|
|||
|
|
|
|||
|
|
# 合并：为结果数据添加原始数据列
|
|||
|
|
merged_df = results_df.copy()
|
|||
|
|
|
|||
|
|
# 初始化新列
|
|||
|
|
for col in original_cols_to_add:
|
|||
|
|
merged_df[col] = None
|
|||
|
|
|
|||
|
|
# 逐行匹配并填充
|
|||
|
|
matched_count = 0
|
|||
|
|
for idx, row in merged_df.iterrows():
|
|||
|
|
image_name = row[results_image_col]
|
|||
|
|
if image_name in original_map:
|
|||
|
|
for col in original_cols_to_add:
|
|||
|
|
merged_df.at[idx, col] = original_map[image_name].get(col)
|
|||
|
|
matched_count += 1
|
|||
|
|
|
|||
|
|
stats["merged_rows"] = len(merged_df)
|
|||
|
|
stats["matched_count"] = matched_count
|
|||
|
|
stats["unmatched_results"] = len(merged_df) - matched_count
|
|||
|
|
|
|||
|
|
return merged_df, stats
|
|||
|
|
|
|||
|
|
|
|||
|
|
def collect_and_merge_xlsx(
|
|||
|
|
source_dir: str,
|
|||
|
|
data_all_dir: str,
|
|||
|
|
output_dir: str,
|
|||
|
|
merge: bool = True,
|
|||
|
|
dry_run: bool = False
|
|||
|
|
) -> List[dict]:
|
|||
|
|
"""
|
|||
|
|
收集并合并 xlsx 文件
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
source_dir: batch_output 目录路径
|
|||
|
|
data_all_dir: data_all 目录路径
|
|||
|
|
output_dir: 输出目录路径
|
|||
|
|
merge: 是否合并原始数据
|
|||
|
|
dry_run: 预览模式
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
处理结果列表
|
|||
|
|
"""
|
|||
|
|
source_path = Path(source_dir)
|
|||
|
|
data_all_path = Path(data_all_dir)
|
|||
|
|
output_path = Path(output_dir)
|
|||
|
|
|
|||
|
|
if not source_path.exists():
|
|||
|
|
print(f"错误: 源目录不存在: {source_dir}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 创建输出目录
|
|||
|
|
if not dry_run:
|
|||
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
results = []
|
|||
|
|
|
|||
|
|
# 遍历子文件夹
|
|||
|
|
for folder in sorted(source_path.iterdir()):
|
|||
|
|
if not folder.is_dir():
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
folder_name = folder.name
|
|||
|
|
results_file = folder / "results.xlsx"
|
|||
|
|
|
|||
|
|
if not results_file.exists():
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 输出文件名
|
|||
|
|
output_file = output_path / f"{folder_name}.xlsx"
|
|||
|
|
|
|||
|
|
# 查找对应的原始数据文件
|
|||
|
|
original_file = data_all_path / f"{folder_name}_text_img.xlsx"
|
|||
|
|
|
|||
|
|
result_info = {
|
|||
|
|
"folder": folder_name,
|
|||
|
|
"results_file": str(results_file),
|
|||
|
|
"original_file": str(original_file) if original_file.exists() else None,
|
|||
|
|
"output_file": str(output_file),
|
|||
|
|
"merged": False,
|
|||
|
|
"stats": {}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if dry_run:
|
|||
|
|
if merge and original_file.exists():
|
|||
|
|
print(f"[预览] 合并: {folder_name}/results.xlsx + {folder_name}_text_img.xlsx -> {folder_name}.xlsx")
|
|||
|
|
else:
|
|||
|
|
print(f"[预览] 复制: {folder_name}/results.xlsx -> {folder_name}.xlsx")
|
|||
|
|
results.append(result_info)
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 执行合并或复制
|
|||
|
|
if merge and original_file.exists():
|
|||
|
|
try:
|
|||
|
|
merged_df, stats = merge_xlsx_files(results_file, original_file)
|
|||
|
|
merged_df.to_excel(output_file, index=False, engine="openpyxl")
|
|||
|
|
|
|||
|
|
result_info["merged"] = True
|
|||
|
|
result_info["stats"] = stats
|
|||
|
|
|
|||
|
|
print(f"已合并: {folder_name}")
|
|||
|
|
print(f" - 分析结果: {stats['results_rows']} 行")
|
|||
|
|
print(f" - 原始数据: {stats['original_rows']} 行")
|
|||
|
|
print(f" - 匹配成功: {stats['matched_count']} 行")
|
|||
|
|
print(f" - 添加列: {stats['original_columns_added']}")
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"合并失败 {folder_name}: {e}")
|
|||
|
|
# 回退到复制模式
|
|||
|
|
import shutil
|
|||
|
|
shutil.copy2(results_file, output_file)
|
|||
|
|
print(f" 已回退到复制模式")
|
|||
|
|
else:
|
|||
|
|
# 只复制，不合并
|
|||
|
|
import shutil
|
|||
|
|
shutil.copy2(results_file, output_file)
|
|||
|
|
|
|||
|
|
if merge and not original_file.exists():
|
|||
|
|
print(f"已复制: {folder_name} (原始数据不存在: {folder_name}_text_img.xlsx)")
|
|||
|
|
else:
|
|||
|
|
print(f"已复制: {folder_name}")
|
|||
|
|
|
|||
|
|
results.append(result_info)
|
|||
|
|
|
|||
|
|
return results
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser(
|
|||
|
|
description="收集并合并 batch_output 和 data_all 中的 xlsx 文件",
|
|||
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|||
|
|
epilog="""
|
|||
|
|
示例:
|
|||
|
|
python3 collect_xlsx.py # 默认合并并输出
|
|||
|
|
python3 collect_xlsx.py -o ../data/merged # 指定输出目录
|
|||
|
|
python3 collect_xlsx.py --no-merge # 不合并，只复制
|
|||
|
|
python3 collect_xlsx.py -n # 预览模式
|
|||
|
|
"""
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
parser.add_argument(
|
|||
|
|
"-s", "--source",
|
|||
|
|
default="../data/batch_output",
|
|||
|
|
help="batch_output 目录路径 (默认: ../data/batch_output)"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
parser.add_argument(
|
|||
|
|
"-d", "--data-all",
|
|||
|
|
default="../data/data_all",
|
|||
|
|
help="data_all 目录路径 (默认: ../data/data_all)"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
parser.add_argument(
|
|||
|
|
"-o", "--output",
|
|||
|
|
default="../data/collected_xlsx",
|
|||
|
|
help="输出目录路径 (默认: ../data/collected_xlsx)"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
parser.add_argument(
|
|||
|
|
"--no-merge",
|
|||
|
|
action="store_true",
|
|||
|
|
help="不合并原始数据，只复制分析结果"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
parser.add_argument(
|
|||
|
|
"-n", "--dry-run",
|
|||
|
|
action="store_true",
|
|||
|
|
help="预览模式，只打印不执行"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
# 转换为绝对路径
|
|||
|
|
script_dir = Path(__file__).parent
|
|||
|
|
source_dir = (script_dir / args.source).resolve()
|
|||
|
|
data_all_dir = (script_dir / args.data_all).resolve()
|
|||
|
|
output_dir = (script_dir / args.output).resolve()
|
|||
|
|
|
|||
|
|
print("=" * 60)
|
|||
|
|
print("收集并合并 xlsx 文件")
|
|||
|
|
print("=" * 60)
|
|||
|
|
print(f"分析结果目录: {source_dir}")
|
|||
|
|
print(f"原始数据目录: {data_all_dir}")
|
|||
|
|
print(f"输出目录: {output_dir}")
|
|||
|
|
print(f"合并模式: {'否' if args.no_merge else '是'}")
|
|||
|
|
print("-" * 60)
|
|||
|
|
|
|||
|
|
results = collect_and_merge_xlsx(
|
|||
|
|
str(source_dir),
|
|||
|
|
str(data_all_dir),
|
|||
|
|
str(output_dir),
|
|||
|
|
merge=not args.no_merge,
|
|||
|
|
dry_run=args.dry_run
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
print("-" * 60)
|
|||
|
|
merged_count = sum(1 for r in results if r.get("merged"))
|
|||
|
|
print(f"共处理 {len(results)} 个文件")
|
|||
|
|
if not args.no_merge:
|
|||
|
|
print(f" - 合并成功: {merged_count}")
|
|||
|
|
print(f" - 仅复制: {len(results) - merged_count}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|