1

app.py

@@ -0,0 +1,30 @@
+import os.path
+
+from pages.main_page import init_pages
+from vector_db.vector_db_client import VectorDB
+import configparser
+
+CONFIG_MAP = {}
+CONFIG_FILE = "conf/config.ini"
+
+
+def init_config():
+    print(f"init configs {CONFIG_FILE}")
+    if not os.path.exists(CONFIG_FILE):
+        raise FileNotFoundError(f'The config file {CONFIG_FILE} not found.')
+
+    conf = configparser.ConfigParser()
+    conf.read(CONFIG_FILE, encoding="UTF-8")
+
+    for selection in conf.sections():
+        for k, v in conf.items(selection):
+            CONFIG_MAP[f'{selection}.{k}'] = v
+
+
+if __name__ == "__main__":
+    init_config()
+    vdb = VectorDB(CONFIG_MAP)
+
+    server_name = CONFIG_MAP.get('server.name')
+    server_port = int(CONFIG_MAP.get('server.port'))
+    init_pages(vdb, server_name, server_port,CONFIG_MAP)

conf/config.ini

@@ -0,0 +1,18 @@
+
+[vector_db]
+address=http://lb-ozz7dtn0-l7dqtav6xoyir4bm.clb.ap-guangzhou.tencentclb.com:10000
+key=6uwifScRaDLNND2970YKH4uIHe3eZZn37hYu3FN2
+ai_db=test_ai_graph_db
+ai_graph_emb_collection=test_ai_graph_collection
+
+[download_model]
+local_model_path=models/
+model_name=openai/clip-vit-base-patch32
+
+[graph_upload]
+local_graph_path=graphs/
+
+
+[server]
+name=127.0.0.1
+port=8080

pages/chat_search.py

@@ -0,0 +1,94 @@
+import gradio as gr
+from vector_db.vector_db_client import VectorDB
+from PIL import Image
+from transformers import AutoProcessor, CLIPModel
+import os
+import uuid
+from tcvectordb.model.document import SearchParams
+import traceback
+
+LOCAL_MODEL_PATH = "download_model.local_model_path"
+MODEL_NAME = "download_model.model_name"
+LOCAL_GRAPH_PATH = "graph_upload.local_graph_path"
+
+class ChatSearch:
+    def __init__(self, config, vdb: VectorDB):
+        self.vdb = vdb
+        self.model_name = config.get(MODEL_NAME)
+        self.local_model_path = config.get(LOCAL_MODEL_PATH)
+        self.local_graph_path = config.get(LOCAL_GRAPH_PATH)
+        self.model_cache_directory = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), self.local_model_path, self.model_name)
+        self.graph_cache_directory = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), self.local_graph_path)
+
+    def initial_model(self):
+        model = CLIPModel.from_pretrained(self.model_cache_directory)
+        processor = AutoProcessor.from_pretrained(self.model_cache_directory)
+        return model, processor
+    
+    def search_result(self, image):
+        if image is None:
+            return "请先上传图片..."
+        
+        if not os.path.exists(self.model_cache_directory):
+            return f"缓存目录 {self.model_cache_directory} 不存在，无法初始化模型。"
+        
+        model, processor = self.initial_model()
+        try:
+            # 生成唯一的文件名
+            unique_filename = f"{uuid.uuid4().hex}.png"
+            image_path = os.path.join(self.graph_cache_directory, unique_filename)
+            
+            # 保存图片到指定文件夹
+            image.save(image_path)
+            
+            image_vector = self._process_image(image_path, model, processor).squeeze().tolist()  # 转换为一维列表
+            
+            # 假设你的 VectorDB 支持图片搜索
+            collection = self.vdb.get_collection()
+            res = collection.search(
+                vectors=[image_vector],
+                params=SearchParams(ef=200),
+                limit=10,
+                output_fields=['local_graph_path']
+            )
+
+            results = []
+            for i, docs in enumerate(res):
+                for doc in docs:
+                    image_path = doc['local_graph_path']
+                    try:
+                        image = Image.open(image_path)
+                        results.append(image)
+                    except Exception as e:
+                        print(f"无法加载图片 {image_path}: {e}")
+            return results
+        except Exception as e:
+            print(f"问题：{e}\n")
+            error_trace = traceback.format_exc()
+            print(error_trace)
+    
+    def _process_image(self, image_path, emb_model, processor):
+        """
+        处理单个图片文件，将其转换为向量。
+
+        参数:
+        image_path (str): 图片文件的路径。
+
+        返回:
+        torch.Tensor: 图片的向量表示。
+        """
+        image = Image.open(image_path)
+        inputs = processor(images=image, return_tensors="pt")
+        image_features = emb_model.get_image_features(**inputs)
+        return image_features
+
+    def get_chart(self):
+        return gr.Interface(
+            fn=self.search_result,
+            inputs=gr.Image(type="pil", label="上传图片"),
+            outputs=gr.Gallery(label="检索结果"),
+            theme="soft",
+            description="上传图片进行检索",
+            allow_flagging="never"
+        )
+

pages/initial_and_upload.py

@@ -0,0 +1,232 @@
+import time
+import gradio as gr
+from huggingface_hub import snapshot_download
+import os
+import zipfile
+from PIL import Image, UnidentifiedImageError
+from transformers import AutoProcessor, CLIPModel
+from vector_db.vector_db_client import VectorDB
+from tcvectordb.model.document import Document
+import uuid
+import traceback
+import numpy as np 
+# 生成随机的 UUID
+LOCAL_MODEL_PATH = "download_model.local_model_path"
+MODEL_NAME = "download_model.model_name"
+LOCAL_GRAPH_PATH="graph_upload.local_graph_path"
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+init_css="""
+            <style>
+                .equal-height-row {
+                    display: flex;
+                }
+                .equal-height-column {
+                    flex: 1;
+                    display: flex;
+                    flex-direction: column;
+                }
+                .equal-height-column > * {
+                    flex: 1;
+                }
+            </style>
+            """
+class Initial_and_Upload:
+
+    def __init__(self, config,vdb: VectorDB):
+        self.vdb = vdb
+        self.model_name = config.get(MODEL_NAME)
+        self.local_model_path = config.get(LOCAL_MODEL_PATH)
+        self.local_graph_path=config.get(LOCAL_GRAPH_PATH)
+        self.model_cache_directory = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), self.local_model_path, self.model_name)
+        self.graph_cache_directory = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), self.local_graph_path)
+
+    def initial_model(self):
+        model = CLIPModel.from_pretrained(self.model_cache_directory)
+        processor = AutoProcessor.from_pretrained(self.model_cache_directory)
+        return model,processor
+    
+    def _download_model(self, model_name, progress=gr.Progress()):
+        """
+        下载指定的Hugging Face模型并保存在指定位置。
+
+        参数:
+        model_name (str): 模型在Hugging Face上的名字。
+        save_directory (str): 模型文件保存的位置。
+        """
+ 
+
+        os.environ['TRANSFORMERS_CACHE'] = self.model_cache_directory
+        
+        # 创建保存目录（如果不存在）
+        if not os.path.exists(self.model_cache_directory):
+            os.makedirs(self.model_cache_directory)
+        text = f"[正在尝试下载] 模型 {model_name}，因为涉及到模型相关的多个文件下载，进度仅在后台显示。\n"
+        progress(0.5, desc=text)
+        try:
+            # 下载模型
+            snapshot_download(
+                repo_id=model_name,
+                local_dir=self.model_cache_directory,
+                local_dir_use_symlinks=False,
+            )
+            
+            progress(1, f"模型 {model_name} 已下载并保存在 {self.model_cache_directory}")
+            text += f"模型 {model_name} 已下载并保存在 {self.model_cache_directory}"
+            
+            time.sleep(0.3)
+            return text
+        except Exception as e:
+            text += f"[下载失败] 失败原因：{e}"
+            return text
+
+    def _process_image(self, image_path,emb_model,processor):
+        """
+        处理单个图片文件，将其转换为向量。
+
+        参数:
+        image_path (str): 图片文件的路径。
+
+        返回:
+        torch.Tensor: 图片的向量表示。
+        """
+        
+        image = Image.open(image_path)
+        # image.verify()  # 验证图片是否有效
+        inputs = processor(images=image, return_tensors="pt")
+        image_features = emb_model.get_image_features(**inputs)
+        return image_features
+
+    def _handle_upload(self, file, progress=gr.Progress()):
+        """
+        处理上传的文件，识别是图片还是ZIP压缩包，并将图片转换为向量。
+
+        参数:
+        file (file): 上传的文件。
+
+        返回:
+        str: 文件类型和处理结果。
+        """
+        output_text = ""
+        image_vectors = []
+        if not os.path.exists(self.model_cache_directory):
+            output_text += f"缓存目录 {self.model_cache_directory} 不存在，无法初始化模型。"
+        else:
+            model, processor = self.initial_model()
+            collection = self.vdb.get_collection()
+            
+            if zipfile.is_zipfile(file.name):
+                with zipfile.ZipFile(file.name, 'r') as zip_ref:
+                    zip_ref.extractall(self.local_graph_path)
+                    image_files = [file_name for file_name in zip_ref.namelist() if file_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')) and not file_name.startswith('__MACOSX') and not file_name.startswith('._')]
+                    
+                    total_files = len(image_files)
+                    for i, file_name in enumerate(image_files):
+                        image_path = os.path.join(self.local_graph_path, file_name)
+                        try:
+                            image_vector = self._process_image(image_path, model, processor).squeeze().tolist()  # 转换为一维列表
+                            random_uuid = str(uuid.uuid4())  # 转换为字符串
+                            collection.upsert(documents=[Document(id=random_uuid, vector=image_vector, local_graph_path=image_path)], build_index=True)
+                            output_text += f"处理图片: {file_name}\n"
+                        except UnidentifiedImageError:
+                            output_text += f"无法识别图片文件: {file_name}\n"
+                        
+                        # 更新进度
+                        progress((i + 1) / total_files)
+                    
+                output_text += "上传的是ZIP压缩包，已解压缩并处理所有图片。"
+            else:
+                try:
+                    # 保存单张图片到指定文件夹
+                    image_path = os.path.join(self.graph_cache_directory, os.path.basename(file.name))
+                    with open(file.name, "rb") as f_src:
+                        with open(image_path, "wb") as f_dst:
+                            f_dst.write(f_src.read())
+                    
+                    image_vector = self._process_image(image_path, model, processor).squeeze().tolist()  # 转换为一维列表
+                    random_uuid = str(uuid.uuid4())  # 转换为字符串
+                    collection.upsert(documents=[Document(id=random_uuid, vector=image_vector, local_graph_path=image_path)], build_index=True)
+                    output_text += "上传的是图片文件，并已处理。\n"
+                    
+                    # 更新进度
+                    progress(1.0)
+                except (IOError, SyntaxError) as e:
+                    output_text += f"无法识别文件类型：{e}\n"
+
+        # 返回处理结果和图片向量
+        return output_text, image_vectors
+    def _initialize_vector_db(self, progress=gr.Progress()):
+        """
+        初始化向量数据库。
+
+        返回:
+        str: 初始化结果。
+        """
+        output_text = f"[正在尝试连接] VectorDB {self.vdb.address}\n"
+        progress(0, desc=output_text)
+        try:
+            client = self.vdb.create_client()
+            client.list_databases()
+            progress(0.05, f"[连接成功] VectorDB {self.vdb.address}\n")
+            output_text += f"[连接成功] VectorDB {self.vdb.address}\n"
+            client.close()
+
+            progress(0.1, f"[正在初始化] ai database '{self.vdb.db_name}'\n")
+            output_text += f"[正在初始化] ai database '{self.vdb.db_name}'\n"
+            self.vdb.init_database()
+            progress(0.3, f"[初始化完成] ai database '{self.vdb.db_name}'\n")
+            output_text += f"[初始化完成] ai database '{self.vdb.db_name}'\n"
+
+            progress(0.5, f"[正在初始化] ai collection '{self.vdb.ai_graph_emb_collection}'\n")
+            output_text += f"[正在初始化] ai collection '{self.vdb.ai_graph_emb_collection}'\n"
+            self.vdb.init_graph_collection()
+            progress(0.9, f"[初始化完成] ai collection '{self.vdb.ai_graph_emb_collection}'\n")
+            output_text += f"[初始化完成] ai collection '{self.vdb.ai_graph_emb_collection}'\n"
+
+            progress(1, f"您可以去图片上传栏目上传图片或ZIP压缩包，然后进一步进行[图片搜索]")
+            output_text += f"您可以去图片上传栏目上传图片或ZIP压缩包，然后进一步进行[图片搜索]"
+
+            time.sleep(0.3)
+        except Exception as e:
+            output_text += f"[数据库访问失败] 失败原因：{e}"
+            error_trace = traceback.format_exc()
+            print(error_trace)
+        return output_text
+
+    def get_init_panel(self):
+        with gr.Blocks() as demo:
+            gr.HTML(init_css)
+            with gr.Row():
+                
+                with gr.Column():
+                    model_name_input = gr.Textbox(lines=1, label="模型名称", placeholder="请输入Hugging Face模型名称...", value=self.model_name)
+                    output = gr.Textbox(lines=10, label="下载进度", placeholder="下载进度将在这里显示...")
+                    init_button = gr.Button("开始下载模型")
+
+                    init_button.click(
+                        fn=self._download_model,
+                        inputs=[model_name_input],
+                        outputs=output
+                    )
+                with gr.Column():
+                    db_init_output = gr.Textbox(lines=14.5, label="数据库初始化结果", placeholder="数据库初始化结果将在这里显示...")
+                    db_init_button = gr.Button("初始化向量数据库")
+
+                    db_init_button.click(
+                        fn=self._initialize_vector_db,
+                        inputs=[],
+                        outputs=db_init_output
+                    )
+            with gr.Row():
+                    upload_file = gr.File(label="上传图片或ZIP压缩包")
+            with gr.Row():
+                    upload_output = gr.Textbox(lines=10, label="上传结果", placeholder="上传结果将在这里显示...")
+            with gr.Row():
+                    upload_button = gr.Button("上传文件")
+
+                    upload_button.click(
+                        fn=self._handle_upload,
+                        inputs=[upload_file],
+                        outputs=[upload_output, gr.State()]
+                    )
+
+        return demo

pages/main_page.py

@@ -0,0 +1,41 @@
+import gradio as gr
+from pages.initial_and_upload import Initial_and_Upload
+from pages.chat_search import ChatSearch
+from vector_db.vector_db_client import VectorDB
+
+main_css = """
+.secondary.svelte-cmf5ev {
+   background: #6366f1;
+   color: #ffffff;
+}
+.main-text {
+   color: #6366f1;
+   padding-top: 20px;
+   text-align: center;
+}
+.main-title {
+   font-size: 28px;
+   font-weight: bold;
+   padding-top: 10px;
+   padding-bottom: 15px;
+   color: #6366f1;
+   text-align: center;
+}
+"""
+
+
+def init_pages(vdb: VectorDB, server_name: str, server_port,config):
+    initial_and_upload = Initial_and_Upload(config,vdb)
+    chat_search=ChatSearch(config,vdb)
+
+    with gr.Blocks(title="Tencent VectorDB", theme="soft", css=main_css) as demo:
+        with gr.Row():
+            gr.HTML("<div class='main-title'>Tencent VectorDB AI Demo -- Graph search</div>")
+
+        with gr.Tab("初始化页面"):
+            initial_and_upload.get_init_panel()
+        with gr.Tab("图搜图界面"):
+            chat_search.get_chart()
+
+    
+    demo.launch(server_name=server_name, server_port=server_port)

requirements.txt

@@ -0,0 +1,9 @@
+tencentcloud-sdk-python-common
+pydantic
+gradio==4.39.0
+transformers==4.44.2
+torch==2.4.1
+torchvision==0.19.1
+huggingface_hub
+tcvectordb
+fastapi==0.111.1

vector_db/vector_db_client.py

@@ -0,0 +1,63 @@
+import tcvectordb
+from tcvectordb.model.database import Database
+from tcvectordb.model.collection import Collection
+from tcvectordb.model.index import Index, VectorIndex, FilterIndex, HNSWParams
+from tcvectordb.model.enum import FieldType, IndexType, MetricType
+VDB_ADDRESS = "vector_db.address"
+VDB_KEY = "vector_db.key"
+AI_DB_NAME = "vector_db.ai_db"
+AI_COLLECTION_NAME = "vector_db.ai_graph_emb_collection"
+
+
+class VectorDB:
+    def __init__(self, config):
+        self.address = config.get(VDB_ADDRESS)
+        self.key = config.get(VDB_KEY)
+        self.db_name = config.get(AI_DB_NAME)
+        self.ai_graph_emb_collection = config.get(AI_COLLECTION_NAME)
+
+        print(f"Try to connect vector db {self.address}")
+        self.client = self.create_client()
+        self._test_simple()
+
+    def create_client(self):
+        return tcvectordb.RPCVectorDBClient(
+            url=self.address,
+            username='root',
+            key=self.key,
+            timeout=30
+        )
+
+    def _test_simple(self):
+        self.client.list_databases()
+
+    def init_database(self):
+        try:
+            self.client.create_database(self.db_name)
+        except tcvectordb.exceptions.VectorDBException:
+            self.client.drop_database(self.db_name)
+            self.client.create_database(self.db_name)
+
+    def init_graph_collection(self):
+        index = Index(
+                    FilterIndex(name='id', field_type=FieldType.String, index_type=IndexType.PRIMARY_KEY),
+                    FilterIndex(name='local_graph_path', field_type=FieldType.String, index_type=IndexType.FILTER),
+                    VectorIndex(name='vector', dimension=512, index_type=IndexType.HNSW,
+                                metric_type=MetricType.COSINE, params=HNSWParams(m=16, efconstruction=200))
+                )    
+
+        database: Database = self.client.database(self.db_name)
+        try:
+            database.create_collection(name=self.ai_graph_emb_collection ,shard=1,replicas=2,index=index,
+            description='this is a collection of graph embedding'
+
+)
+        except tcvectordb.exceptions.VectorDBException:
+            database.drop_collection(self.ai_graph_emb_collection)
+            database.create_collection(name=self.ai_graph_emb_collection ,shard=1,replicas=2,index=index,
+            description='this is a collection of graph embedding'
+
+)
+    def get_collection(self) -> Collection:
+        database: Database = self.client.database(self.db_name)
+        return database.collection(self.ai_graph_emb_collection)