feat: add region-based similarity search

Adds a "Search within image" button to the lightbox that lets the user
draw a crop region on an image and find visually similar results using
that crop's embedding. The crop is embedded in-memory without a temp
file via a new embed_image_crop method on ClipImageEmbedder.

Introduces a folder-scoped cosine search (search_image_ids_by_embedding_in_folder)
to support the current-folder scope option, and wires up the new
find_similar_by_region Tauri command end-to-end from Rust through to
the Zustand store and Lightbox UI.
This commit is contained in:
2026-06-02 19:39:31 +01:00
parent ff4a568b57
commit df17497808
6 changed files with 477 additions and 24 deletions
+73
View File
@@ -59,6 +59,19 @@ pub struct FindSimilarImagesParams {
pub threshold: Option<f32>,
}
#[derive(Deserialize)]
pub struct FindSimilarByRegionParams {
pub image_id: i64,
/// Normalized crop rect (0.01.0).
pub crop_x: f32,
pub crop_y: f32,
pub crop_w: f32,
pub crop_h: f32,
pub folder_id: Option<i64>,
pub offset: Option<usize>,
pub limit: Option<usize>,
}
#[derive(Deserialize)]
pub struct DebugSimilarImagesParams {
pub image_id: i64,
@@ -354,6 +367,66 @@ pub async fn find_similar_images(
})
}
#[tauri::command]
pub async fn find_similar_by_region(
db: State<'_, DbState>,
params: FindSimilarByRegionParams,
) -> Result<SimilarImagesPage, String> {
let conn = db.get().map_err(|e| e.to_string())?;
let limit = params.limit.unwrap_or(32);
let offset = params.offset.unwrap_or(0);
// Look up the source image path
let image = db::get_image_by_id(&conn, params.image_id).map_err(|e| e.to_string())?;
let image_path = std::path::Path::new(&image.path);
// Embed the cropped region in-memory (no temp file needed)
let embedder = embedder::ClipImageEmbedder::new().map_err(|e| e.to_string())?;
let embedding = embedder
.embed_image_crop(
image_path,
params.crop_x,
params.crop_y,
params.crop_w,
params.crop_h,
)
.map_err(|e| e.to_string())?;
// Search for similar images using the crop embedding
let image_ids = match params.folder_id {
Some(folder_id) => vector::search_image_ids_by_embedding_in_folder(
&conn,
&embedding,
folder_id,
Some(params.image_id),
offset + limit + 1,
)
.map_err(|e| e.to_string())?,
None => {
let mut ids = vector::search_image_ids_by_embedding(&conn, &embedding, offset + limit + 1)
.map_err(|e| e.to_string())?;
// Exclude the source image from global results
ids.retain(|&id| id != params.image_id);
ids
}
};
let has_more = image_ids.len() > offset + limit;
let page_ids = image_ids
.into_iter()
.skip(offset)
.take(limit)
.collect::<Vec<_>>();
let images = db::get_images_by_ids(&conn, &page_ids).map_err(|e| e.to_string())?;
Ok(SimilarImagesPage {
images,
offset,
limit,
has_more,
})
}
#[derive(Serialize)]
pub struct SimilarImagesDebug {
pub image_id: i64,
+45
View File
@@ -74,6 +74,51 @@ impl ClipImageEmbedder {
Ok(self.embed_images(&[path.to_path_buf()])?.remove(0))
}
/// Embed a cropped region of an image without writing a temp file to disk.
/// `crop_x`, `crop_y`, `crop_w`, `crop_h` are normalized 0.01.0 coordinates.
pub fn embed_image_crop(
&self,
path: &Path,
crop_x: f32,
crop_y: f32,
crop_w: f32,
crop_h: f32,
) -> Result<Vec<f32>> {
let img = image::ImageReader::open(path)?
.with_guessed_format()?
.decode()?;
let img_w = img.width() as f32;
let img_h = img.height() as f32;
let x = ((crop_x * img_w) as u32).min(img.width().saturating_sub(1));
let y = ((crop_y * img_h) as u32).min(img.height().saturating_sub(1));
let w = ((crop_w * img_w) as u32).max(1).min(img.width() - x);
let h = ((crop_h * img_h) as u32).max(1).min(img.height() - y);
let cropped = img.crop_imm(x, y, w, h);
let resized = cropped.resize_to_fill(
self.image_size as u32,
self.image_size as u32,
image::imageops::FilterType::Triangle,
);
let raw = resized.to_rgb8().into_raw();
let tensor = candle_core::Tensor::from_vec(
raw,
(self.image_size, self.image_size, 3),
&candle_core::Device::Cpu,
)?
.permute((2, 0, 1))?
.to_dtype(candle_core::DType::F32)?
.affine(2.0 / 255.0, -1.0)?;
let batch = tensor.unsqueeze(0)?.to_device(&self.device)?;
let features = self.model.get_image_features(&batch)?;
let normalized = candle_transformers::models::clip::div_l2_norm(&features)?;
Ok(normalized.get(0)?.flatten_all()?.to_vec1::<f32>()?)
}
pub fn embed_images(&self, paths: &[PathBuf]) -> Result<Vec<Vec<f32>>> {
let images = load_images(paths, self.image_size)?.to_device(&self.device)?;
let features = self.model.get_image_features(&images)?;
+1
View File
@@ -87,6 +87,7 @@ pub fn run() {
commands::reindex_folder,
commands::update_image_details,
commands::find_similar_images,
commands::find_similar_by_region,
commands::debug_similar_images,
commands::retry_failed_embeddings,
commands::semantic_search_images,
+34
View File
@@ -310,6 +310,40 @@ pub fn search_image_ids_by_embedding(
Ok(ids)
}
/// Brute-force cosine search scoped to a single folder, ordered by ascending distance.
/// Used for region-based similarity search where we want folder-scoped results.
pub fn search_image_ids_by_embedding_in_folder(
conn: &Connection,
embedding: &[f32],
folder_id: i64,
exclude_image_id: Option<i64>,
limit: usize,
) -> Result<Vec<i64>> {
if embedding.len() != CLIP_VECTOR_DIM {
return Err(anyhow!(
"expected {}-dimensional embedding, got {}",
CLIP_VECTOR_DIM,
embedding.len()
));
}
let packed = pack_f32(embedding);
let exclude_id = exclude_image_id.unwrap_or(-1);
let mut stmt = conn.prepare(
"SELECT v.image_id
FROM image_vec v
JOIN images i ON i.id = v.image_id
WHERE i.folder_id = ?2
AND v.image_id != ?3
ORDER BY vec_distance_cosine(v.embedding, vec_f32(?1)) ASC
LIMIT ?4",
)?;
let rows = stmt.query_map((&packed, folder_id, exclude_id, limit as i64), |row| {
row.get::<_, i64>(0)
})?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
#[allow(dead_code)]
pub fn search_caption_ids_by_embedding(
conn: &Connection,