Upload 3 files

image2image.py

@@ -0,0 +1,86 @@
+from functools import partial
+import json
+from multiprocessing.pool import ThreadPool as Pool
+import gradio as gr
+import PIL
+from PIL import Image
+from utils import *
+
+from clip_retrieval.clip_client import ClipClient
+
+def image2text_gr():
+    def clip_api(query_image=None, return_n=8, model_name=clip_base, thumbnail=yes):
+        client = ClipClient(url="http://9.135.121.52:1234//knn-service",
+                            indice_name="ltr_cover_index",
+                            aesthetic_weight=0,
+                            num_images=int(return_n))
+        result = client.query(image=query_image)
+
+        if not result or len(result) == 0:
+            print("no result found")
+            return None
+
+        print(f"get result sucessed, num: {len(result)}")
+
+        cover_urls = [res['cover_url'] for res in result]
+        cover_info = []
+        for res in result:
+            json_info = {"cover_url": res['cover_url'],
+                         "similarity": round(res['similarity'], 6),
+                         "docid": res['docids']}
+            cover_info.append(str(json_info))
+        pool = Pool()
+        new_url2image = partial(url2img, thumbnail=thumbnail)
+        ret_imgs = pool.map(new_url2image, cover_urls)
+        pool.close()
+        pool.join()
+
+        new_ret = []
+        for i in range(len(ret_imgs)):
+            new_ret.append([ret_imgs[i], cover_info[i]])
+        return new_ret
+
+    examples = [
+        ["https://xingchen-data.oss-cn-zhangjiakou.aliyuncs.com/coco/2014/test2014/COCO_test2014_000000000069.jpg", 20,
+         clip_base, "是"],
+        ["https://xingchen-data.oss-cn-zhangjiakou.aliyuncs.com/coco/2014/test2014/COCO_test2014_000000000080.jpg", 20,
+         clip_base, "是"],
+        ["https://xingchen-data.oss-cn-zhangjiakou.aliyuncs.com/coco/2014/train2014/COCO_train2014_000000000009.jpg",
+         20, clip_base, "是"],
+        ["https://xingchen-data.oss-cn-zhangjiakou.aliyuncs.com/coco/2014/train2014/COCO_train2014_000000000308.jpg",
+         20, clip_base, "是"]
+    ]
+
+    title = "<h1 align='center'>CLIP图到图搜索应用</h1>"
+
+    with gr.Blocks() as demo:
+        gr.Markdown(title)
+        gr.Markdown(description)
+        with gr.Row():
+            with gr.Column(scale=1):
+                with gr.Column(scale=2):
+                    img = gr.Textbox(value="https://xingchen-data.oss-cn-zhangjiakou.aliyuncs.com/coco/2014/test2014/COCO_test2014_000000000069.jpg", label="图片地址", elem_id=0, interactive=True)
+                num = gr.components.Slider(minimum=0, maximum=50, step=1, value=8, label="返回图片数（可能被过滤部分）", elem_id=2)
+                model = gr.components.Radio(label="模型选择", choices=[clip_base],
+                                            value=clip_base, elem_id=3)
+                tn = gr.components.Radio(label="是否返回缩略图", choices=[yes, no],
+                                         value=yes, elem_id=4)
+                btn = gr.Button("搜索", )
+            with gr.Column(scale=100):
+                out = gr.Gallery(label="检索结果为：", columns=4, height="auto")
+        inputs = [img, num, model, tn]
+        btn.click(fn=clip_api, inputs=inputs, outputs=out)
+        gr.Examples(examples, inputs=inputs)
+    return demo
+
+
+if __name__ == "__main__":
+    with gr.TabbedInterface(
+            [image2text_gr()],
+            ["图到图搜索"],
+    ) as demo:
+        demo.launch(
+            #enable_queue=True,
+            server_name='127.0.0.1',
+            share=False
+        )

text2image.py

@@ -0,0 +1,78 @@
+from functools import partial
+import json
+from multiprocessing.pool import ThreadPool as Pool
+import gradio as gr
+from utils import *
+
+from clip_retrieval.clip_client import ClipClient
+
+
+def text2image_gr():
+    def clip_api(query_text='', return_n=8, model_name=clip_base, thumbnail="是"):
+        client = ClipClient(url="http://9.135.121.52:1234//knn-service",
+                            indice_name="ltr_cover_index",
+                            aesthetic_weight=0,
+                            num_images=int(return_n))
+        #result = client.query(embedding_input=query_emb)
+        result = client.query(text=query_text)
+        
+        if not result or len(result) == 0:
+            print("no result found")
+            return None
+        
+        print(f"get result sucessed, num: {len(result)}")
+        
+        cover_urls = [res['cover_url'] for res in result]
+        cover_info = []
+        for res in result:
+            json_info = {"cover_url": res['cover_url'],
+                         "similarity": round(res['similarity'], 6),
+                         "docid": res['docids']}
+            cover_info.append(str(json_info))
+        pool = Pool()
+        new_url2image = partial(url2img, thumbnail=thumbnail)
+        ret_imgs = pool.map(new_url2image, cover_urls)
+        pool.close()
+        pool.join()
+
+        new_ret = []
+        for i in range(len(ret_imgs)):
+            new_ret.append([ret_imgs[i], cover_info[i]])
+        return new_ret
+
+    examples = [
+        ["cat", 12, clip_base, "是"],
+        ["dog", 12, clip_base, "是"],
+        ["bag", 12, clip_base, "是"],
+        ["a cat is sit on the table", 12, clip_base, "是"]
+    ]
+
+    title = "<h1 align='center'>CLIP文到图搜索应用</h1>"
+
+    with gr.Blocks() as demo:
+        gr.Markdown(title)
+        gr.Markdown(description)
+        with gr.Row():
+            with gr.Column(scale=1):
+                with gr.Column(scale=2):
+                    text = gr.Textbox(value="cat", label="请填写文本", elem_id=0, interactive=True)
+                num = gr.components.Slider(minimum=0, maximum=50, step=1, value=8, label="返回图片数（可能被过滤部分）", elem_id=2)
+                model = gr.components.Radio(label="模型选择", choices=[clip_base],
+                                            value=clip_base, elem_id=3)
+                thumbnail = gr.components.Radio(label="是否返回缩略图", choices=[yes, no],
+                                                value=yes, elem_id=4)
+                btn = gr.Button("搜索", )
+            with gr.Column(scale=100):
+                out = gr.Gallery(label="检索结果为：", columns=4, height="auto") #.style(grid=4, height=200)
+        inputs = [text, num, model, thumbnail]
+        btn.click(fn=clip_api, inputs=inputs, outputs=out)
+        gr.Examples(examples, inputs=inputs)
+    return demo
+
+if __name__ == "__main__":
+    gr.close_all()
+    with gr.TabbedInterface(
+            [text2image_gr()],
+            ["文到图搜索"],
+    ) as demo:
+        demo.launch(server_name='127.0.0.1', share=False)

utils.py

@@ -0,0 +1,43 @@
+import os
+from PIL import Image
+from PIL import ImageFile
+import requests
+import base64
+from io import BytesIO
+
+clip_base = "CLIP(Base)"
+description = "本项目为CLIP模型的DEMO，可用于图文检索和图像、文本的表征提取，应用于搜索、推荐等应用场景。"
+
+yes = "是"
+no = "否"
+
+server_ip = os.environ.get("CLIP_SERVER_IP", "9.135.121.52")
+
+clip_service_url_d = {
+    clip_base: f'http://{server_ip}/knn-service',
+}
+
+
+def pil_base64(image, img_format="JPEG"):
+    Image.MAX_IMAGE_PIXELS = 1000000000
+    ImageFile.LOAD_TRUNCATED_IMAGES = True
+    img_buffer = BytesIO()
+    image.save(img_buffer, format=img_format)
+    byte_data = img_buffer.getvalue()
+    base64_str = base64.b64encode(byte_data)
+    return base64_str.decode("utf-8")
+
+
+def url2img(img_url, thumbnail=yes):
+    try:
+        #print(img_url, thumbnail)
+        #image = Image.open(requests.get(img_url, stream=True).raw)
+        path = img_url.split("9.22.26.31")[1]
+        image = Image.open(path).convert("RGB")
+        max_ = max(image.size)
+        if max_ > 224 and thumbnail == yes:
+            ratio = max_ // 224
+            image.thumbnail(size=(image.width // ratio, image.height // ratio))
+        return image
+    except Exception as e:
+        print(e)