app.py

# 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()