Related image search (#62)

go.mod

@@ -4,6 +4,7 @@ go 1.19
 
 require (
 	github.com/EdlinOrg/prominentcolor v1.0.0
+	github.com/alecthomas/participle/v2 v2.0.0
 	github.com/deepmap/oapi-codegen v1.8.2
 	github.com/dgraph-io/ristretto v0.0.2
 	github.com/docker/go-units v0.4.0

go.sum

@@ -62,6 +62,10 @@ github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAE
 github.com/ajstarks/svgo v0.0.0-20180226025133-644b8db467af/go.mod h1:K08gAheRH3/J6wwsYMMT4xOr94bZjxIelGM0+d/wbFw=
 github.com/ajstarks/svgo v0.0.0-20200320125537-f189e35d30ca h1:kWzLcty5V2rzOqJM7Tp/MfSX0RMSI1x4IOLApEefYxA=
 github.com/ajstarks/svgo v0.0.0-20200320125537-f189e35d30ca/go.mod h1:K08gAheRH3/J6wwsYMMT4xOr94bZjxIelGM0+d/wbFw=
+github.com/alecthomas/assert/v2 v2.2.2 h1:Z/iVC0xZfWTaFNE6bA3z07T86hd45Xe2eLt6WVy2bbk=
+github.com/alecthomas/participle/v2 v2.0.0 h1:Fgrq+MbuSsJwIkw3fEj9h75vDP0Er5JzepJ0/HNHv0g=
+github.com/alecthomas/participle/v2 v2.0.0/go.mod h1:rAKZdJldHu8084ojcWevWAL8KmEU+AT+Olodb+WoN2Y=
+github.com/alecthomas/repr v0.2.0 h1:HAzS41CIzNW5syS8Mf9UwXhNH1J9aix/BvDRf1Ml2Yk=
 github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
@@ -346,6 +350,7 @@ github.com/hashicorp/go-multierror v1.1.0/go.mod h1:spPvp8C1qA32ftKqdAHm4hHTbPw+
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/hcl v0.0.0-20170914154624-68e816d1c783/go.mod h1:oZtUIOe8dh44I2q6ScRibXws4Ajl+d+nod3AaR9vL5w=
+github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
 github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
 github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=

internal/image/database.go

@@ -1215,6 +1215,36 @@ func (source *Database) List(dirs []string, options ListOptions) <-chan InfoList
 	return out
 }
 
+func (source *Database) GetImageEmbedding(id ImageId) (clip.Embedding, error) {
+	conn := source.pool.Get(nil)
+	defer source.pool.Put(conn)
+
+	stmt := conn.Prep(`
+		SELECT inv_norm, embedding
+		FROM clip_emb
+		WHERE file_id = ?;`)
+	defer stmt.Reset()
+
+	stmt.BindInt64(1, int64(id))
+
+	if exists, err := stmt.Step(); err != nil {
+		return nil, err
+	} else if !exists {
+		return nil, nil
+	}
+
+	invnorm := uint16(clip.InvNormMean + stmt.ColumnInt64(0))
+
+	size := stmt.ColumnLen(1)
+	bytes := make([]byte, size)
+	read := stmt.ColumnBytes(1, bytes)
+	if read != size {
+		return nil, fmt.Errorf("error reading embedding: buffer underrun, expected %d actual %d bytes", size, read)
+	}
+
+	return clip.FromRaw(bytes, invnorm), nil
+}
+
 func (source *Database) ListEmbeddings(dirs []string, options ListOptions) <-chan EmbeddingsResult {
 	out := make(chan EmbeddingsResult, 100)
 	go func() {

internal/image/source.go

@@ -411,6 +411,10 @@ func (source *Source) GetImagePath(id ImageId) (string, error) {
 	return path, nil
 }
 
+func (source *Source) GetImageEmbedding(id ImageId) (clip.Embedding, error) {
+	return source.database.GetImageEmbedding(id)
+}
+
 func (source *Source) IndexFiles(dir string, max int, counter chan<- int) {
 	dir = filepath.FromSlash(dir)
 	indexed := make(map[string]struct{})

internal/scene/sceneSource.go

@@ -15,6 +15,7 @@ import (
 	"photofield/internal/layout"
 	"photofield/internal/metrics"
 	"photofield/internal/render"
+	"photofield/search"
 )
 
 type SceneSource struct {
@@ -82,12 +83,26 @@ func (source *SceneSource) loadScene(config SceneConfig, imageSource *image.Sour
 
 		if scene.Search != "" {
 			searchDone := metrics.Elapsed("search embed")
-			embedding, err := imageSource.Clip.EmbedText(scene.Search)
-			if err != nil {
-				log.Println("search embed failed")
-				scene.Error = fmt.Sprintf("Search failed: %s", err.Error())
+			q, err := search.Parse(scene.Search)
+			if err == nil {
+				similar, err := q.QualifierInt("img")
+				if err == nil {
+					embedding, err := imageSource.GetImageEmbedding(image.ImageId(similar))
+					if err != nil {
+						log.Println("search get similar failed")
+						scene.Error = fmt.Sprintf("Search failed: %s", err.Error())
+					}
+					scene.SearchEmbedding = embedding
+				}
+			}
+			if scene.SearchEmbedding == nil && scene.Error == "" {
+				embedding, err := imageSource.Clip.EmbedText(scene.Search)
+				if err != nil {
+					log.Println("search embed failed")
+					scene.Error = fmt.Sprintf("Search failed: %s", err.Error())
+				}
+				scene.SearchEmbedding = embedding
 			}
-			scene.SearchEmbedding = embedding
 			searchDone()
 		}
 

search/query.go

@@ -0,0 +1,72 @@
+package search
+
+import (
+	"fmt"
+	"strconv"
+
+	"github.com/alecthomas/participle/v2"
+	"github.com/alecthomas/participle/v2/lexer"
+)
+
+type Query struct {
+	Terms []*Term `parser:"@@*" json:"terms"`
+}
+
+type Term struct {
+	String    *string        `parser:"@String" json:"string,omitempty"`
+	Qualifier *Qualifier     `parser:"| @@" json:"qualifier,omitempty"`
+	Word      *string        `parser:"| @Word" json:"word,omitempty"`
+	Pos       lexer.Position `parser:"" json:"start"`
+	EndPos    lexer.Position `parser:"" json:"end"`
+}
+
+type Qualifier struct {
+	Key   string `parser:"@Word ':'"`
+	Value string `parser:"@Word (@':' @Word)*"`
+}
+
+var lex *lexer.StatefulDefinition
+var par *participle.Parser[Query]
+
+func init() {
+	lex = lexer.MustSimple([]lexer.SimpleRule{
+		{Name: "Whitespace", Pattern: `[ \t]+`},
+		{Name: "Word", Pattern: `\w+`},
+		{Name: "String", Pattern: `"(\\"|[^"])*"`},
+		{Name: "Colon", Pattern: `:`},
+	})
+
+	par = participle.MustBuild[Query](
+		participle.Lexer(lex),
+		participle.Elide("Whitespace"),
+		participle.Unquote("String"),
+	)
+}
+
+func Parse(str string) (*Query, error) {
+	return par.ParseString("", str)
+}
+
+func (q *Query) QualifierInt(key string) (int, error) {
+	if q == nil {
+		return 0, fmt.Errorf("nil query")
+	}
+
+	if len(q.Terms) == 0 {
+		return 0, fmt.Errorf("empty query")
+	}
+
+	if len(q.Terms) > 1 {
+		return 0, fmt.Errorf("too many terms")
+	}
+
+	if q.Terms[0].Qualifier == nil {
+		return 0, fmt.Errorf("no qualifier")
+	}
+
+	if q.Terms[0].Qualifier.Key != key {
+		return 0, fmt.Errorf(`qualifier not %s`, key)
+	}
+
+	return strconv.Atoi(q.Terms[0].Qualifier.Value)
+}

ui/src/components/CollectionView.vue

@@ -16,6 +16,7 @@
       @selectTagId="onSelectTagId"
       @region="onScrollRegion"
       @scene="scrollScene = $event"
+      @search="onSearch"
     >
     </scroll-viewer>
 
@@ -33,6 +34,7 @@
       :fullpage="true"
       @region="onStripRegion"
       @scene="stripScene = $event"
+      @search="onSearch"
     >
     </strip-viewer>
 
@@ -137,6 +139,15 @@ const search = computed(() => {
   return route.query.search;
 })
 
+const onSearch = (search) => {
+  router.push({
+    query: {
+      ...route.query,
+      search,
+    }
+  });
+}
+
 const debug = computed(() => {
   const v = {};
   for (const key in route.query) {

ui/src/components/RegionMenu.vue

@@ -39,6 +39,9 @@
         <ui-item @click="copyImageLink()">
           Copy Image Link
         </ui-item>
+        <ui-item @click="findSimilar()">
+          Find Similar Images
+        </ui-item>
       </ui-nav>
       <div v-if="expanded" class="thumbnails">
         <a
@@ -71,7 +74,7 @@ import { useViewport } from '../use';
 
 export default {
   props: ["region", "scene", "flipX", "flipY", "tileSize"],
-  emits: ["close"],
+  emits: ["close", "search"],
   components: { TileViewer, ExpandButton },
   data() {
     return {
@@ -122,6 +125,12 @@ export default {
       await navigator.clipboard.writeText(this.fileUrl);
       this.$emit("close");
     },
+    findSimilar() {
+      const id = this.region?.data?.id;
+      if (!id) return;
+      this.$emit("search", `img:${id}`);
+      this.$emit("close");
+    },
   }
 };
 </script>

ui/src/components/ScrollViewer.vue

@@ -59,6 +59,7 @@
         :flipY="contextFlip.y"
         :tileSize="512"
         @close="closeContextMenu()"
+        @search="emit('search', $event)"
       ></RegionMenu>
     </ContextMenu>
   </div>
@@ -97,6 +98,7 @@ const emit = defineEmits({
   reindex: null,
   region: null,
   selectTagId: null,
+  search: null,
 })
 
 const {

ui/src/components/StripViewer.vue

@@ -65,6 +65,7 @@
         :flipY="contextFlip.y"
         :tileSize="512"
         @close="closeContextMenu()"
+        @search="onSearch($event)"
       ></RegionMenu>
     </ContextMenu>
   </div>
@@ -100,6 +101,7 @@ const emit = defineEmits({
   scene: null,
   reindex: null,
   region: null,
+  search: null,
 })
 
 const {
@@ -147,6 +149,11 @@ const { region, navigate, exit, mutate: updateRegion } = useSeekableRegion({
   regionId,
 });
 
+const onSearch = async (term) => {
+  await exit();
+  emit("search", term);
+}
+
 const fileId = computed(() => region.value?.data?.id);
 
 const favorite = async (tag) => {

ui/src/use.js

@@ -240,9 +240,9 @@ export function useSeekableRegion({ scene, collectionId, regionId }) {
     id: index,
   });
 
-  const exit = () => {
+  const exit = async () => {
     applyTask.cancelAll();
-    router.push({
+    await router.push({
       name: "collection",
       params: {
         collectionId: collectionId.value,