aboutsummaryrefslogtreecommitdiff
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/embedders/ai.rs11
-rw-r--r--src/main.rs19
2 files changed, 21 insertions, 9 deletions
diff --git a/src/embedders/ai.rs b/src/embedders/ai.rs
index 8c9de11..7d31a6b 100644
--- a/src/embedders/ai.rs
+++ b/src/embedders/ai.rs
@@ -13,7 +13,10 @@ pub(crate) struct ContentEmbedder<'a, Metric> {
}
impl<'a, Metric> ContentEmbedder<'a, Metric> {
pub(crate) fn new(cfg: &'a Config) -> Self {
- ContentEmbedder { cfg, _sim: PhantomData }
+ ContentEmbedder {
+ cfg,
+ _sim: PhantomData,
+ }
}
}
@@ -23,11 +26,13 @@ impl<Metric: VecMetric> ContentEmbedder<'_, Metric> {
paths: &[PathBuf],
) -> Result<Vec<Result<<Self as BatchEmbedder>::Embedding>>> {
let mut options = ImageInitOptions::default();
- options.cache_dir = self.cfg.base_dirs.get_cache_home();
+ options.cache_dir = self.cfg.cache_dir.join("models");
let embedder = ImageEmbedding::try_new(options)?;
let bar = ProgressBar::new(paths.len() as u64);
- bar.set_style(ProgressStyle::with_template("{bar:20.cyan/blue} {pos}/{len} {msg}")?);
+ bar.set_style(ProgressStyle::with_template(
+ "{bar:20.cyan/blue} {pos}/{len} {msg}",
+ )?);
bar.enable_steady_tick(std::time::Duration::from_millis(100));
bar.set_message("Embedding images...");
diff --git a/src/main.rs b/src/main.rs
index 92f51f1..7652166 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,6 @@
#![feature(iterator_try_collect)]
-use anyhow::Result;
+use anyhow::{anyhow, Result};
use clap::Parser;
use sha2::{Digest, Sha512_256};
use std::{
@@ -37,6 +37,7 @@ struct Args {
embedder: Embedder,
/// Symlink the sorted images into this directory
+ #[cfg(unix)]
#[arg(short = 's', long)]
symlink_dir: Option<PathBuf>,
@@ -77,13 +78,14 @@ struct Args {
#[derive(Debug)]
struct Config {
- base_dirs: xdg::BaseDirectories,
+ cache_dir: PathBuf,
}
fn get_config() -> Result<Config> {
- let dirs = xdg::BaseDirectories::with_prefix("embeddings-sort")?;
-
- Ok(Config { base_dirs: dirs })
+ let glob_cache_dir = dirs::cache_dir().ok_or(anyhow!("Could not get cache directory"))?;
+ Ok(Config {
+ cache_dir: glob_cache_dir.join("embeddings-sort"),
+ })
}
fn hash_file(p: &PathBuf) -> Result<[u8; 32]> {
@@ -102,7 +104,7 @@ fn process_embedder<E>(mut e: E, args: &Args, cfg: &Config) -> Result<(Vec<PathB
where
E: BatchEmbedder,
{
- let db = sled::open(cfg.base_dirs.place_cache_file("embeddings.db")?)?;
+ let db = sled::open(cfg.cache_dir.join("embeddings.db"))?;
let tree = typed_sled::Tree::<[u8; 32], E::Embedding>::open(&db, E::NAME);
// find cached embeddings
@@ -180,6 +182,7 @@ where
))
}
+#[allow(unused_variables)] // use_symlinks on windows
fn copy_into(tsp: &[PathBuf], target: &PathBuf, use_symlinks: bool) -> Result<()> {
fs::create_dir_all(target)?;
@@ -191,6 +194,7 @@ fn copy_into(tsp: &[PathBuf], target: &PathBuf, use_symlinks: bool) -> Result<()
};
let tp = target.join(format!("{i:0pad_len$}{ext}"));
+ #[cfg(unix)]
if use_symlinks {
let rel_path =
pathdiff::diff_paths(path::absolute(p)?, path::absolute(target)?).unwrap();
@@ -199,6 +203,8 @@ fn copy_into(tsp: &[PathBuf], target: &PathBuf, use_symlinks: bool) -> Result<()
} else {
reflink_copy::reflink_or_copy(p, tp)?;
}
+ #[cfg(not(unix))]
+ reflink_copy::reflink_or_copy(p, tp)?;
}
Ok(())
}
@@ -226,6 +232,7 @@ fn main() -> Result<()> {
eprintln!("Found tour with length: {}", total_dist);
}
+ #[cfg(unix)]
if let Some(p) = args.symlink_dir {
copy_into(&tsp_path, &p, true)?
}