Files
rust-imager/src/engine/pipeline.rs
lychang 6b2f19ef47 update
2026-07-20 01:01:59 +08:00

492 lines
23 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

//! 处理管线 — 批量处理相册元数据
//!
//! 核心架构:脏集追踪 + 增量解析
//!
//! Phase 0 清洗标签标点
//! Phase A1 检测标签中的人名并直接提取到 person 字段 → 加入脏集
//! Phase A2 清理非人物 → 加入脏集
//! Phase A3 收集新相册/半成品 → 加入脏集
//! Phase 1 解析脏集相册(增量解析,不扫全库)
//! Phase 2 人物合并
//! Phase B 后校验标签中的人名
//!
//! 脏集(DirtySet)贯穿各阶段,每个阶段发现的"需要重新解析"的相册都会加入脏集,
//! 最终只处理脏集里的相册,而非全库扫描,大幅提升效率。
use crate::db::Database;
use crate::engine::title_parser;
use crate::engine::tag_extractor;
use crate::engine::person_merger;
use std::collections::HashSet;
// ════════════════════════════════════════════════
// 脏集 — 追踪需要重新解析的相册 ID
// ════════════════════════════════════════════════
/// 脏集:记录哪些相册需要重新解析
pub struct DirtySet {
ids: HashSet<String>,
}
impl DirtySet {
pub fn new() -> Self { Self { ids: HashSet::new() } }
pub fn add(&mut self, id: &str) { self.ids.insert(id.to_string()); }
pub fn extend(&mut self, ids: impl IntoIterator<Item = String>) { self.ids.extend(ids); }
pub fn contains(&self, id: &str) -> bool { self.ids.contains(id) }
pub fn len(&self) -> usize { self.ids.len() }
pub fn is_empty(&self) -> bool { self.ids.is_empty() }
pub fn into_inner(self) -> HashSet<String> { self.ids }
}
// ════════════════════════════════════════════════
// Phase 0:清洗已有标签中的头尾标点
// ════════════════════════════════════════════════
pub async fn phase_clean_tags(db: &Database) -> usize {
let albums = db.get_all_albums_full().await;
let mut cleaned_count = 0usize;
let mut removed_tags = HashSet::new();
for album in &albums {
let tags_json = match album.tags.as_deref() {
Some(t) if !t.is_empty() && t != "[]" => t,
_ => continue,
};
let old_tags: Vec<String> = match serde_json::from_str(tags_json) {
Ok(t) => t, Err(_) => continue,
};
if old_tags.is_empty() { continue; }
let mut changed = false;
let mut new_tags = Vec::new();
for tag in old_tags {
let cleaned = match tag_extractor::clean_tag(&tag) {
Some(t) => t,
None => { removed_tags.insert(tag); changed = true; continue; }
};
if cleaned != tag { removed_tags.insert(tag); changed = true; }
new_tags.push(cleaned);
}
if !changed { continue; }
cleaned_count += 1;
let tags_json_new = serde_json::to_string(&new_tags).unwrap_or_else(|_| "[]".to_string());
let all_words_json = album.all_words.as_deref().unwrap_or("[]");
let all_words_new: Vec<[String; 2]> = serde_json::from_str::<Vec<[String; 2]>>(all_words_json)
.unwrap_or_default()
.into_iter()
.filter_map(|pair| { let cleaned = tag_extractor::clean_tag(&pair[0])?; Some([cleaned, pair[1].clone()]) })
.collect();
let all_words_json_new = serde_json::to_string(&all_words_new).unwrap_or_else(|_| "[]".to_string());
db.update_album_tags(&album.id, &tags_json_new, &all_words_json_new).await;
}
if cleaned_count > 0 {
println!("[管线·Phase0] 标签清洗完成:清洗 {} 个相册,移除标签: {:?}",
cleaned_count, removed_tags.iter().collect::<Vec<_>>());
} else {
println!("[管线·Phase0] 标签清洗完成:无需要清洗的标签");
}
cleaned_count
}
// ════════════════════════════════════════════════
// Phase A1:检测标签中的人名并直接提取到 person 字段
// ════════════════════════════════════════════════
/// Phase A1:检测标签中可能的人名并直接提取到 person_name
///
/// 核心变化:不再只标记脏集等待 Phase 1,而是直接从 tags 中提取人名写入 person 字段。
/// 因为 tags 中的人名比从 title 重解析更直接、更准确。
///
/// 两层检测 + 最终写入:
/// 1. 逐个标签 — 标签是已知人物或像人名 → 收集为候选人名(中文加权)
/// 2. jieba 合并(空格分隔)— 发现跨标签拆分的人名(["杨晨","晨Yome"])
/// 3. 最佳候选人名 → 写入 person_name,从标签中移除关联标签
///
/// 返回 (脏集, 从标签中移除的人名标签数)
pub async fn phase_detect_person_in_tags(db: &Database) -> (DirtySet, usize) {
let albums = db.get_all_albums_full().await;
let mut dirty = DirtySet::new();
let mut person_like_removed = 0usize;
let known_persons: HashSet<String> = db.get_all_person_names().await
.into_iter().map(|(n, _)| n.to_lowercase()).collect();
for album in &albums {
let tags_json = match album.tags.as_deref() {
Some(t) if !t.is_empty() && t != "[]" => t,
_ => continue,
};
let old_tags: Vec<String> = match serde_json::from_str(tags_json) {
Ok(t) => t, Err(_) => continue,
};
let has_person = album.person_name_original.as_deref()
.map(|s| !s.is_empty()).unwrap_or(false);
// 快速判断是否需要处理
let has_known_tag = old_tags.iter().any(|t| known_persons.contains(&t.to_lowercase()));
let has_like_tag = !has_person && old_tags.iter().any(|t|
title_parser::is_likely_person_name(t) && !title_parser::is_excluded_name(t));
if has_person && !has_known_tag && !has_like_tag { continue; }
// ── 第一层:逐个标签收集候选人名 ──
let mut known_person_found: Option<String> = None;
let mut person_like_tags: Vec<String> = Vec::new();
for tag in &old_tags {
let tag_lower = tag.to_lowercase();
if known_persons.contains(&tag_lower) {
known_person_found = Some(tag.clone());
continue;
}
if !has_person && title_parser::is_likely_person_name(tag) && !title_parser::is_excluded_name(tag) {
person_like_tags.push(tag.clone());
continue;
}
}
// 从 person_like_tags 中选最优(中文加权)
let best_layer1 = person_like_tags.into_iter().max_by_key(|t| {
if let Some(cn) = title_parser::extract_cn(t) {
if cn.len() >= 2 { cn.len() * 10 + t.len() } else { t.len() }
} else { t.len() }
});
// ── 第二层:空格连接所有标签,用 jieba 提取人名 ──
let combined = old_tags.join(" ");
let (jieba_name, _) = title_parser::extract_person_and_tags(&combined);
// ── 确定最终人名 ──
let final_person = known_person_found.or_else(|| {
match (&jieba_name, &best_layer1) {
(Some(j), Some(b)) => {
// jieba 结果更完整 → 用 jieba
if j.len() > b.len() || title_parser::is_excluded_name(b) { Some(j.clone()) } else { Some(b.clone()) }
}
(Some(j), None) => Some(j.clone()),
(None, Some(b)) => Some(b.clone()),
(None, None) => None,
}
});
if let Some(ref person_name) = final_person {
// ── 从标签中移除所有人名相关标签 ──
let mut changed = false;
let mut clean_tags: Vec<String> = Vec::new();
let person_lower = person_name.to_lowercase();
let (primary_cleaned, _) = title_parser::split_person_name_with_aliases(person_name);
let primary_lower = primary_cleaned.to_lowercase();
for tag in &old_tags {
let tag_lower = tag.to_lowercase();
let should_remove =
tag_lower == person_lower // 完全匹配
|| tag_lower == primary_lower // 中文核心匹配
|| (tag.len() >= 2 && tag.chars().all(|c| c.is_ascii_alphabetic())
&& !tag_lower.starts_with("cos")
&& person_lower.contains(&tag_lower)) // 英文后缀
|| (tag_lower.len() >= 2 && person_lower.contains(&tag_lower)); // 人名包含
if should_remove {
person_like_removed += 1;
changed = true;
continue;
}
clean_tags.push(tag.clone());
}
if changed || !has_person {
// ── 写入 person_name(直接提取,无需等待 Phase 1)──
let (primary, aliases) = title_parser::split_person_name_with_aliases(person_name);
let final_name = if primary.is_empty() { person_name.clone() } else { primary };
// 用事务包裹所有写入,避免逐条 autocommit 导致巨慢
let mut tx = db.pool.begin().await.expect("begin tx");
db.update_album_person_with_tx(&mut tx, &album.id, person_name, &final_name).await;
if !title_parser::is_excluded_name(&final_name) {
let pid = db.insert_person_with_tx(&mut tx, &final_name, None).await;
db.insert_person_album_with_tx(&mut tx, pid, &album.id).await;
for (at, an) in &aliases {
if !an.is_empty() && an != &final_name {
db.insert_person_alias_with_tx(&mut tx, pid, at, an).await;
}
}
db.update_person_counts_with_tx(&mut tx, pid).await;
}
// ── 写入清理后的标签 ──
let tags_json_new = serde_json::to_string(&clean_tags).unwrap_or_else(|_| "[]".to_string());
let all_words_json = album.all_words.as_deref().unwrap_or("[]");
let all_words_new: Vec<[String; 2]> = serde_json::from_str::<Vec<[String; 2]>>(all_words_json)
.unwrap_or_default()
.into_iter()
.filter(|pair| !known_persons.contains(&pair[0].to_lowercase()))
.collect();
let words_json_new = serde_json::to_string(&all_words_new).unwrap_or_else(|_| "[]".to_string());
db.update_album_tags_with_tx(&mut tx, &album.id, &tags_json_new, &words_json_new).await;
tx.commit().await.expect("commit tx");
if has_person { dirty.add(&album.id); }
}
}
}
if person_like_removed > 0 {
println!("[管线·PhaseA1] 从标签中检测并提取人名:移除 {} 个人名标签,写入 {} 个相册的 person 字段",
person_like_removed, dirty.len());
} else {
println!("[管线·PhaseA1] 检测标签中的人名:无异常");
}
(dirty, person_like_removed)
}
// ════════════════════════════════════════════════
// Phase A2:清理 persons 表中的非人物
// ════════════════════════════════════════════════
pub async fn phase_clean_persons(db: &Database) -> (DirtySet, usize) {
let persons = db.get_all_persons_with_counts().await;
let mut dirty = DirtySet::new();
let mut deleted = 0usize;
for (pid, name, _albums, _images) in &persons {
if title_parser::is_excluded_name(name) || !title_parser::is_likely_person_name(name) {
let album_list = db.get_person_albums(*pid).await;
let album_ids: Vec<String> = album_list.into_iter().map(|a| a.id.to_string()).collect();
for aid in &album_ids { dirty.add(aid); }
println!("[管线·PhaseA2] 删除非人物:{} (id={}),释放 {} 个相册", name, pid, album_ids.len());
db.update_album_person_by_person_id(*pid).await;
db.delete_person_completely(*pid).await;
deleted += 1;
}
}
if deleted > 0 {
println!("[管线·PhaseA2] 清理完成:删除 {} 个非人物,标记 {} 个脏相册", deleted, dirty.len());
} else {
println!("[管线·PhaseA2] 清理完成:无非人物需要删除");
}
(dirty, deleted)
}
// ════════════════════════════════════════════════
// Phase A3:收集新相册 + 半成品
// ════════════════════════════════════════════════
pub async fn phase_collect_new_albums(db: &Database) -> DirtySet {
let albums = db.get_all_albums_full().await;
let mut dirty = DirtySet::new();
let mut new_count = 0usize;
let mut partial_count = 0usize;
for album in &albums {
let tags_empty = album.tags.as_deref()
.map(|t| t.is_empty() || t == "[]" || t == "null").unwrap_or(true);
if tags_empty { dirty.add(&album.id); new_count += 1; continue; }
let has_person = album.person_name_original.as_deref()
.map(|s| !s.is_empty()).unwrap_or(false);
if !has_person { dirty.add(&album.id); partial_count += 1; }
}
println!("[管线·PhaseA3] 收集相册:{} 个全新,{} 个半成品,脏集合计 {} 个", new_count, partial_count, dirty.len());
dirty
}
// ════════════════════════════════════════════════
// Phase 1:增量解析 — 只处理脏集中的相册
// ════════════════════════════════════════════════
pub async fn phase_parse_dirty(db: &Database, dirty_set: &DirtySet) -> (usize, usize) {
if dirty_set.is_empty() {
println!("[管线·Phase1] 增量解析跳过:脏集为空");
return (0, 0);
}
let albums = db.get_all_albums_full().await;
let patterns = person_merger::collect_person_patterns(db).await;
let use_patterns = !patterns.is_empty();
let mut person_updated = 0usize;
let mut tag_updated = 0usize;
let mut skipped = 0usize;
for album in &albums {
if !dirty_set.contains(&album.id) { continue; }
let title = album.title.as_deref().unwrap_or("");
if title.is_empty() { skipped += 1; continue; }
let (jieba_person, jieba_tags) = title_parser::extract_person_and_tags(title);
let best_person = if use_patterns {
person_merger::match_title(title, &patterns)
.map(|(_, name, _)| name).or(jieba_person)
} else {
jieba_person
};
// 用事务包裹所有写入,避免逐条 autocommit 导致巨慢
let mut pending_tx: Option<sqlx::Transaction<'_, sqlx::Sqlite>> = None;
if let Some(name) = &best_person {
let (primary, aliases) = title_parser::split_person_name_with_aliases(name);
let final_name = if primary.is_empty() { name.clone() } else { primary };
let mut tx = db.pool.begin().await.expect("begin tx");
db.update_album_person_with_tx(&mut tx, &album.id, name, &final_name).await;
let pid = db.insert_person_with_tx(&mut tx, &final_name, None).await;
db.insert_person_album_with_tx(&mut tx, pid, &album.id).await;
for (at, an) in &aliases {
if !an.is_empty() && an != &final_name {
db.insert_person_alias_with_tx(&mut tx, pid, at, an).await;
}
}
db.update_person_counts_with_tx(&mut tx, pid).await;
pending_tx = Some(tx);
person_updated += 1;
}
// 确定标签排除用的人名 — 用最完整的名字
// 优先级:person_name_original(完整含英文后缀) > best_person > person_name_cleaned(仅中文)
let exclude_person = album.person_name_original.as_deref()
.filter(|s| !s.is_empty())
.or_else(|| best_person.as_deref())
.or_else(|| album.person_name_cleaned.as_deref().filter(|s| !s.is_empty()));
let mut raw_tags = tag_extractor::extract_tags_from_title(title, exclude_person);
raw_tags.extend(jieba_tags);
let unique = tag_extractor::unique_tags(raw_tags);
let tag_names: Vec<String> = unique.iter().map(|(t, _)| t.clone()).collect();
if !tag_names.is_empty() {
let tags_json = serde_json::to_string(&tag_names).unwrap_or_else(|_| "[]".to_string());
let all_words: Vec<[String; 2]> = unique.iter().map(|(t, tp)| [t.clone(), tp.clone()]).collect();
let words_json = serde_json::to_string(&all_words).unwrap_or_else(|_| "[]".to_string());
let mut tx = match pending_tx.take() {
Some(t) => t,
None => db.pool.begin().await.expect("begin tx"),
};
db.update_album_tags_with_tx(&mut tx, &album.id, &tags_json, &words_json).await;
tx.commit().await.expect("commit tx");
tag_updated += 1;
} else if let Some(tx) = pending_tx.take() {
// 只有 person 写入,没有标签写入
tx.commit().await.expect("commit tx");
}
}
println!("[管线·Phase1] 增量解析完成:脏集 {} 个,处理 {} 个,跳过 {} 个,更新人名 {} 个,更新标签 {} 个",
dirty_set.len(), person_updated.max(tag_updated), skipped, person_updated, tag_updated);
(person_updated, tag_updated)
}
// ════════════════════════════════════════════════
// Phase 2:人物合并
// ════════════════════════════════════════════════
pub async fn phase_merge_persons(db: &Database) -> (usize, usize) {
let result = person_merger::execute_merge(db).await;
if result.persons_merged > 0 {
println!("[管线·Phase2] 人物合并完成:分析 {} 组,合并 {} 人,迁移 {} 个别名",
result.groups_analyzed, result.persons_merged, result.aliases_moved);
} else {
println!("[管线·Phase2] 人物合并完成:无需要合并的人物");
}
(result.persons_merged, result.aliases_moved)
}
// ════════════════════════════════════════════════
// Phase B:后校验标签 — 用合并后的人名检查
// ════════════════════════════════════════════════
pub async fn phase_clean_tags_after_parse(db: &Database) -> usize {
let albums = db.get_all_albums_full().await;
let all_person_names: HashSet<String> = db.get_all_person_names().await
.into_iter().map(|(n, _)| n.to_lowercase()).collect();
let all_aliases: HashSet<String> = db.get_all_aliases().await
.into_iter().map(|(a, _)| a.to_lowercase()).collect();
let person_or_alias: HashSet<String> = all_person_names.into_iter().chain(all_aliases).collect();
if person_or_alias.is_empty() {
println!("[管线·PhaseB] 后校验标签完成:无人物数据,跳过");
return 0;
}
let mut tag_removed = 0usize;
for album in &albums {
let tags_json = match album.tags.as_deref() {
Some(t) if !t.is_empty() && t != "[]" => t, _ => continue,
};
let old_tags: Vec<String> = match serde_json::from_str(tags_json) { Ok(t) => t, Err(_) => continue, };
let mut album_removed = 0usize;
let good_tags: Vec<String> = old_tags.into_iter().filter(|t| {
if person_or_alias.contains(&t.to_lowercase()) { album_removed += 1; false } else { true }
}).collect();
if album_removed == 0 { continue; }
tag_removed += album_removed;
let all_words_json = album.all_words.as_deref().unwrap_or("[]");
let all_words_new: Vec<[String; 2]> = serde_json::from_str::<Vec<[String; 2]>>(all_words_json)
.unwrap_or_default().into_iter().filter(|pair| !person_or_alias.contains(&pair[0].to_lowercase())).collect();
let tags_json_new = serde_json::to_string(&good_tags).unwrap_or_else(|_| "[]".to_string());
let words_json_new = serde_json::to_string(&all_words_new).unwrap_or_else(|_| "[]".to_string());
db.update_album_tags(&album.id, &tags_json_new, &words_json_new).await;
}
if tag_removed > 0 {
println!("[管线·PhaseB] 后校验标签完成:移除 {} 个人名/别名", tag_removed);
} else {
println!("[管线·PhaseB] 后校验标签完成:无需要修复的标签");
}
tag_removed
}
// ════════════════════════════════════════════════
// Pipeline 一键执行(7 阶段)
// ════════════════════════════════════════════════
#[derive(Debug, Clone, Default)]
pub struct PipelineResult {
pub tags_cleaned: usize, // Phase 0
pub person_like_tags_removed: usize, // Phase A1
pub dirty_albums_detected: usize, // A1+A2+A3
pub non_persons_deleted: usize, // Phase A2
pub new_albums_found: usize, // Phase A3
pub persons_updated: usize, // Phase 1
pub tags_updated: usize, // Phase 1
pub persons_merged: usize, // Phase 2
pub aliases_moved: usize, // Phase 2
pub tags_removed_person: usize, // Phase B
}
/// 一键执行完整管线(7 阶段顺序串行)
pub async fn run_pipeline(db: &Database) -> PipelineResult {
let tags_cleaned = phase_clean_tags(db).await;
let (dirty_a1, person_like_tags_removed) = phase_detect_person_in_tags(db).await;
let (dirty_a2, non_persons_deleted) = phase_clean_persons(db).await;
let dirty_a3 = phase_collect_new_albums(db).await;
let new_albums_found = dirty_a3.len();
let mut master_dirty = DirtySet::new();
master_dirty.extend(dirty_a1.into_inner());
master_dirty.extend(dirty_a2.into_inner());
master_dirty.extend(dirty_a3.into_inner());
let dirty_albums_detected = master_dirty.len();
let (persons_updated, tags_updated) = phase_parse_dirty(db, &master_dirty).await;
let (persons_merged, aliases_moved) = phase_merge_persons(db).await;
let tags_removed_person = phase_clean_tags_after_parse(db).await;
PipelineResult {
tags_cleaned, person_like_tags_removed, dirty_albums_detected,
non_persons_deleted, new_albums_found,
persons_updated, tags_updated, persons_merged,
aliases_moved, tags_removed_person,
}
}