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:
@@ -59,6 +59,19 @@ pub struct FindSimilarImagesParams {
|
||||
pub threshold: Option<f32>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct FindSimilarByRegionParams {
|
||||
pub image_id: i64,
|
||||
/// Normalized crop rect (0.0–1.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,
|
||||
|
||||
@@ -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.0–1.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)?;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user