Upload search_image_by_image/app.py with huggingface_hub

search_image_by_image/app.py

@@ -0,0 +1,158 @@
+import json
+import os
+from collections import defaultdict
+from typing import List, Dict
+
+import faiss
+import gradio as gr
+import numpy as np
+from PIL import Image
+from cheesechaser.datapool import YandeWebpDataPool, ZerochanWebpDataPool, GelbooruWebpDataPool, \
+    KonachanWebpDataPool, AnimePicturesWebpDataPool, DanbooruNewestWebpDataPool, Rule34WebpDataPool
+from hfutils.operate import get_hf_fs, get_hf_client
+from hfutils.utils import TemporaryDirectory
+from imgutils.tagging import wd14
+from imgutils.utils import ts_lru_cache
+
+from pools import quick_webp_pool
+
+_REPO_ID = 'deepghs/anime_sites_indices'
+
+hf_fs = get_hf_fs()
+hf_client = get_hf_client()
+
+_DEFAULT_MODEL_NAME = 'SwinV2_v3_iqdb_10_46796044_8GB'
+_ALL_MODEL_NAMES = [
+    os.path.dirname(os.path.relpath(path, _REPO_ID))
+    for path in hf_fs.glob(f'{_REPO_ID}/*/knn.index')
+]
+
+_SITE_CLS = {
+    'danbooru': DanbooruNewestWebpDataPool,
+    'yandere': YandeWebpDataPool,
+    'zerochan': ZerochanWebpDataPool,
+    'gelbooru': GelbooruWebpDataPool,
+    'konachan': KonachanWebpDataPool,
+    'anime_pictures': AnimePicturesWebpDataPool,
+    'rule34': Rule34WebpDataPool,
+}
+
+
+def _get_from_ids(site_name: str, ids: List[int]) -> Dict[int, Image.Image]:
+    with TemporaryDirectory() as td:
+        site_cls = _SITE_CLS.get(site_name) or quick_webp_pool(site_name, 3)
+        datapool = site_cls()
+        datapool.batch_download_to_directory(
+            resource_ids=ids,
+            dst_dir=td,
+        )
+
+        retval = {}
+        for file in os.listdir(td):
+            id_ = int(os.path.splitext(file)[0])
+            image = Image.open(os.path.join(td, file))
+            image.load()
+            retval[id_] = image
+
+        return retval
+
+
+def _get_from_raw_ids(ids: List[str]) -> Dict[str, Image.Image]:
+    _sites = defaultdict(list)
+    for id_ in ids:
+        site_name, num_id = id_.rsplit('_', maxsplit=1)
+        num_id = int(num_id)
+        _sites[site_name].append(num_id)
+
+    _retval = {}
+    for site_name, site_ids in _sites.items():
+        _retval.update({
+            f'{site_name}_{id_}': image
+            for id_, image in _get_from_ids(site_name, site_ids).items()
+        })
+    return _retval
+
+
+@ts_lru_cache(maxsize=3)
+def _get_index_info(repo_id: str, model_name: str):
+    image_ids = np.load(hf_client.hf_hub_download(
+        repo_id=repo_id,
+        repo_type='model',
+        filename=f'{model_name}/ids.npy',
+    ))
+    knn_index = faiss.read_index(hf_client.hf_hub_download(
+        repo_id=repo_id,
+        repo_type='model',
+        filename=f'{model_name}/knn.index',
+    ))
+
+    config = json.loads(open(hf_client.hf_hub_download(
+        repo_id=repo_id,
+        repo_type='model',
+        filename=f'{model_name}/infos.json',
+    )).read())["index_param"]
+    faiss.ParameterSpace().set_index_parameters(knn_index, config)
+    return image_ids, knn_index
+
+
+def search(model_name: str, img_input, n_neighbours: int):
+    images_ids, knn_index = _get_index_info(_REPO_ID, model_name)
+    embeddings = wd14.get_wd14_tags(
+        img_input,
+        model_name="SwinV2_v3",
+        fmt="embedding",
+    )
+    embeddings = np.expand_dims(embeddings, 0)
+    faiss.normalize_L2(embeddings)
+
+    dists, indexes = knn_index.search(embeddings, k=n_neighbours)
+    neighbours_ids = images_ids[indexes][0]
+
+    captions = []
+    images = []
+    ids_to_images = _get_from_raw_ids(neighbours_ids)
+    for image_id, dist in zip(neighbours_ids, dists[0]):
+        if image_id in ids_to_images:
+            images.append(ids_to_images[image_id])
+            captions.append(f"{image_id}/{dist:.2f}")
+
+    return list(zip(images, captions))
+
+
+if __name__ == "__main__":
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                img_input = gr.Image(type="pil", label="Input")
+
+            with gr.Column():
+                with gr.Row():
+                    n_model = gr.Dropdown(
+                        choices=_ALL_MODEL_NAMES,
+                        value=_DEFAULT_MODEL_NAME,
+                        label='Index to Use',
+                    )
+                with gr.Row():
+                    n_neighbours = gr.Slider(
+                        minimum=1,
+                        maximum=50,
+                        value=20,
+                        step=1,
+                        label="# of images",
+                    )
+                find_btn = gr.Button("Find similar images")
+
+        with gr.Row():
+            similar_images = gr.Gallery(label="Similar images", columns=[5])
+
+        find_btn.click(
+            fn=search,
+            inputs=[
+                n_model,
+                img_input,
+                n_neighbours,
+            ],
+            outputs=[similar_images],
+        )
+
+    demo.queue().launch()