add content embedder

1 files changed, 14 insertions, 7 deletions

diff --git a/src/main.rs b/src/main.rs
index 5caa4ce..6524de7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -7,14 +7,18 @@ use sha2::{Sha512_256, Digest};
 use std::{cmp::Ordering, collections::HashMap, fs, io, path::PathBuf};
 
 use embedders::*;
-
+use pure_embedders::*;
+use ai_embedders::*;
 mod embedders;
+mod pure_embedders;
+mod ai_embedders;
 
 #[derive(Debug, Clone, Copy, clap::ValueEnum)]
 enum Embedder {
     Brightness,
     Hue,
     Color,
+    Content,
 }
 
 #[derive(Debug, Parser)]
@@ -33,7 +37,7 @@ struct Config {
 fn get_config() -> Result<Config> {
     let dirs = xdg::BaseDirectories::with_prefix("embeddings-sort")?;
 
-    Ok(Config{base_dirs: dirs})
+    Ok(Config { base_dirs: dirs })
 }
 
 fn get_mst<M>(embeds: &Vec<M>) -> HashMap<usize, Vec<usize>>
@@ -119,8 +123,8 @@ fn hash_file(p: &PathBuf) -> Result<[u8; 32]> {
     Ok(hasher.finalize().into_iter().collect::<Vec<u8>>().try_into().unwrap())
 }
 
-fn process_embedder<E>(mut e: E, args: Args, cfg: Config) -> Result<Vec<PathBuf>>
-    where E: EmbedderT
+fn process_embedder<E>(mut e: E, args: Args, cfg: &Config) -> Result<Vec<PathBuf>>
+    where E: BatchEmbedder
 {
     if args.images.is_empty() {
         return Ok(Vec::new());
@@ -145,6 +149,8 @@ fn process_embedder<E>(mut e: E, args: Args, cfg: Config) -> Result<Vec<PathBuf>
             None => Some(i),
             Some(_) => None,
         }).collect();
+    // TODO only run e.embeds if !missing_embeds_indices.is_empty(); this allows
+    // for optimizations in the ai embedde (move pip to ::embeds() instead of ::new())
     let missing_embeds = e.embeds(&missing_embeds_indices
             .iter()
             .map(|i| args.images[*i].clone())
@@ -168,9 +174,10 @@ fn main() -> Result<()> {
     let args = Args::parse();
 
     let tsp_path = match args.embedder {
-        Embedder::Brightness => process_embedder(BrightnessEmbedder, args, cfg),
-        Embedder::Hue => process_embedder(HueEmbedder, args, cfg),
-        Embedder::Color => process_embedder(ColorEmbedder, args, cfg),
+        Embedder::Brightness => process_embedder(BrightnessEmbedder, args, &cfg),
+        Embedder::Hue => process_embedder(HueEmbedder, args, &cfg),
+        Embedder::Color => process_embedder(ColorEmbedder, args, &cfg),
+        Embedder::Content => process_embedder(ContentEmbedder::new(&cfg), args, &cfg),
     }?;
 
     for p in tsp_path {