update

README.md

@@ -1,6 +1,6 @@
 ---
 title: Coin CLIP Retrieval
-emoji: 💻
+emoji: 🪙
 colorFrom: yellow
 colorTo: purple
 sdk: gradio

app.py

@@ -0,0 +1,244 @@
+# coding: utf-8
+# Copyright (C) 2023, [Breezedeus](https://github.com/breezedeus).
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import os
+import sys
+import logging
+from typing import List
+
+import yaml
+
+import gradio as gr
+from PIL import Image
+import numpy as np
+from datasets import load_dataset
+import chromadb
+from chromadb import Settings
+
+from coin_clip.utils import resize_img
+from coin_clip.chroma_embedding import ChromaEmbeddingFunction
+from coin_clip.detect import Detector
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+env = os.environ.get('COIN_ENV', 'local')
+if env == 'hf':
+    config_fp = 'hf_config.yaml'
+else:
+    config_fp = 'local_config.yaml'
+logger.info(f'Use config file: {config_fp}')
+
+total_config = yaml.safe_load(open(config_fp))
+DETECTOR = Detector(
+    model_name=total_config['detector']['model_name'],
+    device=total_config['detector']['device'],
+)
+# USE_REMOVE_BG = total_config['use_remove_bg']
+RESIZED_TO_BEFORE_DETECT = total_config['detector'].get('resized_to', 300)
+
+
+def prepare_chromadb():
+    if env == 'local':
+        return
+    from huggingface_hub import snapshot_download
+    snapshot_download(repo_type='model', repo_id='breezedeus/usa-coins-chromadb', local_dir='./')
+
+
+def load_dataset(data_path):
+    logger.info('Load dataset from %s', data_path)
+
+    if env == 'hf':
+        dataset = load_dataset(data_path, split='train')
+    else:
+        dataset = load_dataset("imagefolder", data_dir=data_path, split='train')
+    return dataset
+
+
+def detect(images):
+    outs = []
+    for idx, img in enumerate(images):
+        img = resize_img(img, RESIZED_TO_BEFORE_DETECT)
+        out = DETECTOR.detect(np.array(img))
+        if not out:
+            out = {'position': None, 'scores': 0.0}
+        else:
+            out = out[0]
+            out.pop('label')
+            out['position'] = out.pop('box')
+        out['from_image_idx'] = idx
+        outs.append(out)
+
+    box_images = []
+    for out, img in zip(outs, images):
+        if out['position'] is None:
+            box_images.append(None)
+        else:
+            # box 比例值转化为绝对位置值
+            w, h = img.size
+            box = out['position']
+            box = (int(box[0] * w), int(box[1] * h), int(box[2] * w), int(box[3] * h))
+            box_images.append(img.crop(box))
+
+    return outs, box_images
+
+
+def load_chroma_db(db_dir, collection_name, model_name, device='cpu'):
+    logger.info('Load chroma db from %s', db_dir)
+    client = chromadb.PersistentClient(
+        path=db_dir, settings=Settings(anonymized_telemetry=False)
+    )
+
+    embedding_function = ChromaEmbeddingFunction(model_name, device)
+    collection = client.get_collection(
+        name=collection_name,
+        embedding_function=embedding_function,
+    )
+    return collection
+
+
+def retrieve(query_image: Image.Image, collection, top_k=20) -> List[Image.Image]:
+    query_image = np.array(query_image)
+    retrieved = collection.query(
+        query_images=[query_image], include=['metadatas', 'distances'], n_results=top_k,
+    )
+    logger.info('retrieved ids: %s', retrieved['ids'][0])
+    logger.info('retrieved distances: %s', retrieved['distances'][0])
+    return [ds_dict[id]['image'] for id in retrieved['ids'][0]]
+
+
+dataset = load_dataset(**total_config['dataset'])
+ds_dict = {_d['id']: _d for _d in dataset}
+
+prepare_chromadb()
+cc_collection = load_chroma_db(**total_config['coin_clip_db'])
+clip_collection = load_chroma_db(**total_config['clip_db'])
+
+
+def search(image_file: Image.Image):
+    images = [image_file.convert('RGB')]
+    detected_outs, box_images = detect(images)
+    box_images = [img for img in box_images if img is not None]
+    if len(box_images) == 0:
+        return [
+            gr.update(visible=False),
+            gr.update(visible=True),
+            gr.update(visible=False),
+            gr.update(visible=False),
+        ]
+
+    box_image = box_images[0]
+    # breakpoint()
+    cc_results = retrieve(box_image, cc_collection, top_k=30)
+    clip_results = retrieve(box_image, clip_collection, top_k=30)
+    return [
+        gr.update(value=box_image, visible=True),
+        gr.update(visible=False),
+        gr.update(value=cc_results, visible=True),
+        gr.update(value=clip_results, visible=True),
+    ]
+
+
+def main():
+    title = 'USA Coin Retrieval by'
+    desc = (
+        '<p style="text-align: center">Coin-CLIP: '
+        '<a href="https://huggingface.co/breezedeus/coin-clip-vit-base-patch32" target="_blank">Model</a>, '
+        '<a href="https://github.com/breezedeus/coin-clip" target="_blank">Github</a>; '
+        'Author: <a href="https://www.breezedeus.com" target="_blank">Breezedeus</a> , '
+        '<a href="https://github.com/breezedeus" target="_blank">Github</a> </p>'
+    )
+    examples = [
+        'examples/c2.jpeg',
+        'examples/c20.jpg',
+        'examples/c21.jpg',
+        'examples/c22.png',
+        'examples/c1.jpg',
+        'examples/c11.jpg',
+        'examples/c3.png',
+        'examples/c4.jpg',
+        'examples/c5.jpeg',
+        'examples/c6.jpeg',
+        'examples/c7.jpg',
+        'examples/c8.jpeg',
+    ]
+
+    with gr.Blocks() as demo:
+        gr.Markdown(
+            f'<h1 style="text-align: center; margin-bottom: 1rem;">{title} <a href="https://github.com/breezedeus/coin-clip" target="_blank">Coin-CLIP</a></h1>'
+        )
+        gr.Markdown(desc)
+        with gr.Row(equal_height=False):
+            with gr.Column(variant='compact', scale=1):
+                gr.Markdown('### Image within a coin')
+                image_file = gr.Image(
+                    label='Coin Image to Search',
+                    type="pil",
+                    image_mode='RGB',
+                    height=400,
+                )
+                sub_btn = gr.Button("Submit", variant="primary")
+            with gr.Column(variant='compact', scale=1):
+                gr.Markdown('### Detected Coin')
+                detected_image = gr.Image(
+                    label='Detected Coin',
+                    type="pil",
+                    interactive=False,
+                    image_mode='RGB',
+                    height=400,
+                )
+                no_detect_warn = gr.Markdown(
+                    '**⚠️ Warning**: No coins detected in image', visible=False
+                )
+
+        with gr.Row(equal_height=False):
+            with gr.Column(variant='compact', scale=1):
+                gr.Markdown('### Results from Coin-CLIP')
+                cc_results = gr.Gallery(
+                    label='Coin-CLIP Results', columns=3, height=2200, show_share_button=True, visible=False
+                )
+
+            with gr.Column(variant='compact', scale=1):
+                gr.Markdown('### Results from CLIP')
+                coin_results = gr.Gallery(
+                    label='CLIP Results', columns=3, height=2200, show_share_button=True, visible=False
+                )
+
+            sub_btn.click(
+                search,
+                inputs=[image_file,],
+                outputs=[detected_image, no_detect_warn, cc_results, coin_results],
+            )
+
+        gr.Examples(
+            label='Examples',
+            examples=examples,
+            inputs=image_file,
+            outputs=[detected_image, no_detect_warn, cc_results, coin_results],
+            fn=search,
+            examples_per_page=12,
+            cache_examples=True,
+        )
+
+    demo.queue(max_size=20)
+    demo.launch()
+
+
+if __name__ == '__main__':
+    main()

hf_config.yaml

@@ -0,0 +1,20 @@
+dataset:
+  data_path: "breezedeus/usa-coins"
+
+detector:
+  model_name: "google/owlvit-base-patch32"
+  device: "cpu"
+  resized_to: 300
+
+coin_clip_db:
+  db_dir: "data/coin_clip_chroma.db"
+  model_name: "breezedeus/coin-clip-vit-base-patch32"
+  collection_name: "coin_clip_collection"
+  device: "cpu"
+
+clip_db:
+  db_dir: "data/clip_chroma.db"
+  model_name: "openai/clip-vit-base-patch32"
+  collection_name: "clip_collection"
+  device: "cpu"
+  device: "cpu"

local_config.yaml

@@ -0,0 +1,20 @@
+dataset:
+  data_path: "data/coin_usa"
+
+detector:
+  model_name: "google/owlvit-base-patch32"
+  device: "cpu"
+  resized_to: 300
+
+coin_clip_db:
+  db_dir: "data/coin_clip_chroma.db"
+  model_name: "../coin-clip-vit-base-patch32"
+  collection_name: "coin_clip_collection"
+  device: "cpu"
+
+clip_db:
+  db_dir: "data/clip_chroma.db"
+  model_name: "openai/clip-vit-base-patch32"
+  collection_name: "clip_collection"
+  device: "cpu"
+  device: "cpu"

requirements.txt

@@ -0,0 +1,8 @@
+--extra-index-url https://pypi.org/simple
+
+coin_clip==0.1
+huggingface_hub
+matplotlib
+chromadb
+datasets
+numpy