feat: add WD tagger with CSV tag support and model download via ureq/zip
Introduces tagger.rs for WD-based image tagging with CSV label support. Adds csv, ureq, and zip dependencies; updates gitignore to exclude Python/JSON files.
This commit is contained in:
+37
-5
@@ -80,7 +80,12 @@ pub fn upsert_caption_embedding(conn: &Connection, image_id: i64, embedding: &[f
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) -> Result<Vec<i64>> {
|
||||
pub fn find_similar_image_ids(
|
||||
conn: &Connection,
|
||||
image_id: i64,
|
||||
limit: usize,
|
||||
folder_id: Option<i64>,
|
||||
) -> Result<Vec<i64>> {
|
||||
let embedding: Vec<u8> = match conn.query_row(
|
||||
"SELECT embedding FROM image_vec WHERE image_id = ?1",
|
||||
[image_id],
|
||||
@@ -91,19 +96,37 @@ pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) ->
|
||||
Err(error) => return Err(error.into()),
|
||||
};
|
||||
|
||||
let allowed_folder_ids = match folder_id {
|
||||
Some(folder_id) => Some(image_ids_for_folder(conn, folder_id)?),
|
||||
None => None,
|
||||
};
|
||||
let search_limit = if allowed_folder_ids.is_some() {
|
||||
count_image_vectors(conn)?.max(1) as usize
|
||||
} else {
|
||||
limit + 1
|
||||
};
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT image_id
|
||||
FROM image_vec
|
||||
WHERE embedding MATCH vec_f32(?1)
|
||||
AND k = ?2",
|
||||
)?;
|
||||
let rows = stmt.query_map((&embedding, (limit as i64) + 1), |row| row.get::<_, i64>(0))?;
|
||||
let rows = stmt
|
||||
.query_map((&embedding, search_limit as i64), |row| {
|
||||
row.get::<_, i64>(0)
|
||||
})?
|
||||
.collect::<rusqlite::Result<Vec<_>>>()?;
|
||||
|
||||
let mut ids = Vec::new();
|
||||
for row in rows {
|
||||
let candidate_id = row?;
|
||||
if candidate_id != image_id {
|
||||
ids.push(candidate_id);
|
||||
if row != image_id {
|
||||
if allowed_folder_ids
|
||||
.as_ref()
|
||||
.is_some_and(|folder_ids| !folder_ids.contains(&row))
|
||||
{
|
||||
continue;
|
||||
}
|
||||
ids.push(row);
|
||||
}
|
||||
if ids.len() >= limit {
|
||||
break;
|
||||
@@ -112,6 +135,15 @@ pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) ->
|
||||
Ok(ids)
|
||||
}
|
||||
|
||||
fn image_ids_for_folder(
|
||||
conn: &Connection,
|
||||
folder_id: i64,
|
||||
) -> Result<std::collections::HashSet<i64>> {
|
||||
let mut stmt = conn.prepare("SELECT id FROM images WHERE folder_id = ?1")?;
|
||||
let rows = stmt.query_map([folder_id], |row| row.get::<_, i64>(0))?;
|
||||
Ok(rows.collect::<rusqlite::Result<std::collections::HashSet<_>>>()?)
|
||||
}
|
||||
|
||||
/// Returns all stored image embeddings with their image IDs, optionally filtered to one folder.
|
||||
/// Each entry is `(image_id, normalized_f32_embedding)`.
|
||||
pub fn get_all_image_embeddings_with_ids(
|
||||
|
||||
Reference in New Issue
Block a user