diff --git a/src-tauri/src/ai_tag_filter.rs b/src-tauri/src/ai_tag_filter.rs new file mode 100644 index 0000000..d683235 --- /dev/null +++ b/src-tauri/src/ai_tag_filter.rs @@ -0,0 +1,45 @@ +/// AI-generated tags that are too broad/noisy to be useful in this gallery. +/// Edit this list to change what the tagger removes. Manual user tags are not +/// affected. +pub const AI_TAG_REMOVAL_LIST: &[&str] = &["1girl", "1boy", "no humans", "2girls", "2boys"]; + +fn normalize_tag_for_removal(tag: &str) -> String { + tag.trim() + .chars() + .filter(|c| !matches!(c, ' ' | '_' | '-')) + .flat_map(char::to_lowercase) + .collect() +} + +pub fn is_removed_ai_tag(tag: &str) -> bool { + let normalized = normalize_tag_for_removal(tag); + AI_TAG_REMOVAL_LIST + .iter() + .any(|blocked| normalize_tag_for_removal(blocked) == normalized) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn removed_ai_tags_match_common_spellings() { + for tag in [ + "1girl", + "1 girl", + "1_girl", + "1-girl", + "NO_HUMANS", + "no humans", + ] { + assert!(is_removed_ai_tag(tag), "{tag} should be removed"); + } + } + + #[test] + fn removed_ai_tags_do_not_match_unrelated_tags() { + for tag in ["girl", "boy", "humans", "solo", "landscape"] { + assert!(!is_removed_ai_tag(tag), "{tag} should be kept"); + } + } +} diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index f6b2bd6..361fa85 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -1,4 +1,4 @@ -use crate::vector; +use crate::{ai_tag_filter, vector}; use anyhow::Result; use r2d2::Pool; use r2d2_sqlite::SqliteConnectionManager; @@ -378,6 +378,42 @@ pub fn migrate(conn: &Connection) -> Result<()> { )?; vector::migrate(conn)?; + remove_filtered_ai_tags(conn)?; + Ok(()) +} + +fn remove_filtered_ai_tags(conn: &Connection) -> Result<()> { + let mut stmt = conn.prepare("SELECT id, tag FROM image_tags WHERE source = 'ai'")?; + let filtered_ids = stmt + .query_map([], |row| { + Ok((row.get::<_, i64>(0)?, row.get::<_, String>(1)?)) + })? + .filter_map(|row| match row { + Ok((id, tag)) if ai_tag_filter::is_removed_ai_tag(&tag) => Some(Ok(id)), + Ok(_) => None, + Err(error) => Some(Err(error)), + }) + .collect::>>()?; + drop(stmt); + + if filtered_ids.is_empty() { + return Ok(()); + } + + let tx = conn.unchecked_transaction()?; + { + let mut delete_stmt = tx.prepare("DELETE FROM image_tags WHERE id = ?1")?; + for id in &filtered_ids { + delete_stmt.execute([id])?; + } + } + tx.execute("DELETE FROM tag_cloud_cache", [])?; + tx.commit()?; + + log::info!( + "Removed {} filtered AI tag(s) from existing library data", + filtered_ids.len() + ); Ok(()) } diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 18a7eac..880c95d 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -1,3 +1,4 @@ +mod ai_tag_filter; mod captioner; mod color; mod commands; diff --git a/src-tauri/src/tagger.rs b/src-tauri/src/tagger.rs index 4923139..2b9616e 100644 --- a/src-tauri/src/tagger.rs +++ b/src-tauri/src/tagger.rs @@ -1,3 +1,4 @@ +use crate::ai_tag_filter; use anyhow::Result; use hf_hub::{api::sync::Api, Repo, RepoType}; use image::{imageops::FilterType, DynamicImage, ImageReader}; @@ -759,6 +760,7 @@ impl WdTagger { .filter(|(entry, prob)| { (entry.category == GENERAL_CATEGORY || entry.category == CHARACTER_CATEGORY) && **prob >= threshold + && !ai_tag_filter::is_removed_ai_tag(&entry.name) }) .map(|(entry, prob)| TagResult { tag: entry.name.clone(), @@ -1137,9 +1139,11 @@ fn joytag_tags_from_logits( .zip(logits.iter()) .filter_map(|(name, logit)| { let confidence = sigmoid(*logit); - (confidence >= threshold).then(|| TagResult { - tag: name.clone(), - confidence, + (confidence >= threshold && !ai_tag_filter::is_removed_ai_tag(name)).then(|| { + TagResult { + tag: name.clone(), + confidence, + } }) }) .collect();