First Release

.gitignore

@@ -0,0 +1,3 @@
+.venv
+.env
+__pycache__

app.py

@@ -0,0 +1,53 @@
+from PIL import Image
+import gradio as gr
+
+from config import (
+    description,
+    article,
+)
+from encode import get_embeddings
+from database import search_vector, format_search_results
+
+def search_images(values):
+    image = Image.new("RGBA", values["composite"].size, (255, 255, 255, 255))
+    image.paste(values["composite"], mask=values["composite"])
+    embedding = get_embeddings([image])[0]
+    results = search_vector(embedding, limit=100)
+    formatted = format_search_results(results)
+
+    _deduplicated = '\t'.join(dict.fromkeys(result.kanji for result in formatted))
+    # TODO Format the results better
+    # Huge boxes using the right font for each of them?
+
+    return f"Results: {_deduplicated}"
+
+# TODO FIND OUT HOW TO CHANGE THE DEFAULT EDITOR TAB?
+input_image = gr.ImageEditor(
+    label="Write the Kanji you want to search for",
+    show_label=False,
+    type="pil",
+    brush=gr.Brush(
+        default_size=3,
+        color_mode="fixed",
+        colors=["#000000", "#ffffff"],
+    ),
+)
+
+
+output_box = gr.Textbox()
+
+
+demo = gr.Interface(
+    fn=search_images,
+    inputs=[input_image],
+    outputs=output_box,
+    title="Kanji Lookup",
+    description=description,
+    article=article,
+    examples="examples",
+    # cache_examples=False,
+    # live=True,
+)
+
+if __name__ == "__main__":
+    demo.launch()

config.py

@@ -0,0 +1,16 @@
+import os
+
+qdrant_location = os.getenv('QDRANT_URL', "localhost")
+qdrant_api_key = os.getenv('QDRANT_API_KEY')
+
+description = """This is a Kanji image search demo. Draw or upload an image of an individual Kanji character."""
+
+article = """
+### About this project
+
+You can find the source code as well as more information in https://github.com/etrotta/kanji_lookup
+
+It uses the "kha-white/manga-ocr-base" ViT Encoder model to create embeddings, then uses a vector database (qdrant) to find similar characters.
+
+The vector database has been populated with over 10k characters from [The KANJIDIC project](https://www.edrdg.org/wiki/index.php/KANJIDIC_Project), each rendered in multiple fonts downloaded from Google Fonts
+"""

database.py

@@ -0,0 +1,34 @@
+import dataclasses
+import torch
+from qdrant_client import QdrantClient, models
+
+from config import qdrant_location, qdrant_api_key
+
+qdrant = QdrantClient(qdrant_location, api_key=qdrant_api_key)
+
+def search_vector(query_vector: torch.Tensor, limit: int=20) -> list[models.ScoredPoint]:
+    hits = qdrant.search(
+        collection_name="kanji",
+        # query_vector=query_vector,
+        query_vector=query_vector.numpy(),
+        limit=limit,
+        with_payload=True,
+    )
+    return hits
+
+@dataclasses.dataclass
+class SearchResult:
+    kanji: str
+    font: str
+    score: float
+
+def format_search_results(hits: list[models.ScoredPoint]) -> list[SearchResult]:
+    formatted = []
+    for point in hits:
+        kanji, font = point.payload["kanji"], point.payload["font"]
+        formatted.append(SearchResult(
+            kanji = kanji,
+            font = font,
+            score = point.score,
+        ))
+    return formatted

encode.py

@@ -0,0 +1,25 @@
+from PIL import Image
+import torch
+from transformers import (
+    VisionEncoderDecoderModel,
+    ViTImageProcessor,  # Load extractor
+    ViTModel,  # Load ViT encoder
+)
+MODEL = "kha-white/manga-ocr-base"
+
+print("Loading models")
+feature_extractor: ViTImageProcessor = ViTImageProcessor.from_pretrained(MODEL, requires_grad=False)
+encoder: ViTModel = VisionEncoderDecoderModel.from_pretrained(MODEL).encoder
+
+if torch.cuda.is_available():
+    print('Using CUDA')
+    encoder.cuda()
+else:
+    print('Using CPU')
+
+def get_embeddings(images: list[Image.Image]) -> torch.Tensor:
+    """Processes the images and returns their Embeddings"""
+    images_rgb = [image.convert("RGB") for image in images]
+    with torch.inference_mode():
+        pixel_values: torch.Tensor = feature_extractor(images_rgb, return_tensors="pt")["pixel_values"]
+        return encoder(pixel_values.to(encoder.device))["pooler_output"].cpu()