feat(visualizer): integrate image upload with similarity search

mini-nav/visualizer/app.py

@@ -1,10 +1,15 @@
+import base64
 import datetime
-from typing import List, Optional, Union
+import io
+from typing import List, Optional
 
 import dash_ag_grid as dag
 import dash_mantine_components as dmc
 from dash import Dash, Input, Output, State, callback, dcc, html
 from database import db_manager
+from feature_retrieval import FeatureRetrieval
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModel
 
 
 class APP(Dash):
@@ -12,6 +17,9 @@ class APP(Dash):
 
     _instance: Optional["APP"] = None
 
+    # Feature retrieval singleton
+    _feature_retrieval: FeatureRetrieval
+
     def __new__(cls) -> "APP":
         if cls._instance is None:
             cls._instance = super().__new__(cls)
@@ -20,6 +28,11 @@ class APP(Dash):
     def __init__(self):
         super().__init__(__name__)
 
+        # Initialize FeatureRetrieval
+        processor = AutoImageProcessor.from_pretrained("facebook/dinov2-large")
+        model = AutoModel.from_pretrained("facebook/dinov2-large")
+        APP._feature_retrieval = FeatureRetrieval(processor, model)
+
         df = (
             db_manager.table.search()
             .select(["id", "label", "vector"])
@@ -56,6 +69,7 @@ class APP(Dash):
                         ),
                         html.Div(id="output-image-upload"),
                         dag.AgGrid(
+                            id="ag-grid",
                             rowData=df.to_dicts(),
                             columnDefs=columnDefs,
                         ),
@@ -71,6 +85,7 @@ class APP(Dash):
 
     @callback(
         Output("output-image-upload", "children"),
+        Output("ag-grid", "rowData"),
         Input("upload-image", "contents"),
         State("upload-image", "filename"),
         State("upload-image", "last_modified"),
@@ -80,23 +95,51 @@ class APP(Dash):
         list_of_names: List[str],
         list_of_dates: List[int] | List[float],
     ):
-        def parse_contents(contents: str, filename: str, date: Union[int, float]):
-            return html.Div(
-                [
-                    html.H5(filename),
-                    html.H6(datetime.datetime.fromtimestamp(date)),
-                    # HTML images accept base64 encoded strings in the same format
-                    # that is supplied by the upload
-                    dmc.Image(src=contents),
-                ]
-            )
+        def parse_base64_to_pil(contents: str) -> Image.Image:
+            """Parse base64 string to PIL Image."""
+            # Remove data URI prefix (e.g., "data:image/png;base64,")
+            base64_str = contents.split(",")[1]
+            img_bytes = base64.b64decode(base64_str)
+            return Image.open(io.BytesIO(img_bytes))
 
         if list_of_contents is not None:
+            # Process first uploaded image for similarity search
+            filename = list_of_names[0]
+            uploaddate = list_of_dates[0]
+            imagecontent = list_of_contents[0]
+
+            pil_image = parse_base64_to_pil(imagecontent)
+
+            # Extract feature vector using DINOv2
+            feature_vector = APP._feature_retrieval.extract_single_image_feature(
+                pil_image
+            )
+
+            # Search for similar images in database
+            results_df = (
+                db_manager.table.search(feature_vector)
+                .select(["id", "label", "vector"])
+                .limit(10)
+                .to_polars()
+            )
+
+            # Convert to AgGrid row format
+            row_data = results_df.to_dicts()
+
+            # Display uploaded images
             children = [
-                parse_contents(c, n, d)
-                for c, n, d in zip(list_of_contents, list_of_names, list_of_dates)
+                html.H5(filename),
+                html.H6(str(datetime.datetime.fromtimestamp(uploaddate))),
+                # HTML images accept base64 encoded strings in same format
+                # that is supplied by the upload
+                dmc.Image(src=imagecontent),
+                dmc.Text(f"{feature_vector[:5]}", size="xs"),
             ]
-            return children
+
+            return children, row_data
+
+        # Return empty if no content
+        return [], []
 
 
 app = APP()