CH3_04_multi_milvus

1
# Google Colab: run this cell first to install notebook dependencies.
2
import os
3
import subprocess
4
import sys
5

6
REPO_URL = "https://github.com/datawhalechina/all-in-rag.git"
7
REPO_DIR = "/content/all-in-rag"
8
VISUAL_BGE_DIR = os.path.join(REPO_DIR, "code", "C3", "visual_bge")
9
C3_DIR = os.path.join(REPO_DIR, "code", "C3")
10

11

12
def run(cmd: list[str]) -> None:
13
    print("$", " ".join(cmd))
14
    subprocess.check_call(cmd)
15

16

17
# visual_bge ships with this repo; clone it when running on Colab.
18
if not os.path.isdir(VISUAL_BGE_DIR):
19
    run(["git", "clone", "--depth", "1", REPO_URL, REPO_DIR])
20

21
# Third-party packages used in this notebook.
22
# Colab already provides torch, torchvision, numpy, requests, and Pillow.
23
# milvus-lite must match pymilvus 2.5.x, otherwise search() fails with
24
# MilvusException: function_score (version mismatch in the local DB backend).
25
run(
26
    [
27
        sys.executable,
28
        "-m",
29
        "pip",
30
        "install",
31
        "-q",
32
        "pymilvus==3.0.0",
33
        "milvus-lite==3.0",
34
        "opencv-python-headless",
35
        "tqdm",
36
        "transformers>=4.40.0",
37
        "huggingface-hub",
38
        "timm",
39
        "einops",
40
        "ftfy",
41
        "regex",
42
    ]
43
)
44

45
# Install the local visual_bge package.
46
run([sys.executable, "-m", "pip", "install", "-q", "-e", VISUAL_BGE_DIR])
47

48
# Notebook paths (../../models, ../../data) assume cwd is code/C3.
49
os.chdir(C3_DIR)
50

51
print("Dependencies installed.")
52
print("Working directory:", os.getcwd())
53
print("Next: download model weights with `!python download_model.py` if needed.")
54
print("If you previously hit a Milvus search error, restart the runtime, then re-run all cells.")

1
# 2. 定义工具 (编码器和可视化函数)
2
class Encoder:
3
    """编码器类，用于将图像和文本编码为向量。"""
4
    def __init__(self, model_name: str, model_path: str):
5
        self.model = Visualized_BGE(model_name_bge=model_name, model_weight=model_path)
6
        self.model.eval()
7

8
    def encode_query(self, image_path: str, text: str) -> list[float]:
9
        with torch.no_grad():
10
            query_emb = self.model.encode(image=image_path, text=text)
11
        return query_emb.tolist()[0] # 去掉 batch 维度
12

13
    def encode_image(self, image_path: str) -> list[float]:
14
        with torch.no_grad():
15
            query_emb = self.model.encode(image=image_path)
16
        return query_emb.tolist()[0]

1
def visualize_results(query_image_path: str, retrieved_images: list, img_height: int = 300, img_width: int = 300, row_count: int = 3) -> np.ndarray:
2
    """从检索到的图像列表创建一个全景图用于可视化。"""
3
    panoramic_width = img_width * row_count # 创建两张空白画布
4
    panoramic_height = img_height * row_count
5
    panoramic_image = np.full((panoramic_height, panoramic_width, 3), 255, dtype=np.uint8) # 左侧结果区
6
    query_display_area = np.full((panoramic_height, img_width, 3), 255, dtype=np.uint8) # 右侧查询区
7

8
    # 处理查询图像
9
    query_pil = Image.open(query_image_path).convert("RGB")
10
    query_cv = np.array(query_pil)[:, :, ::-1]
11
    resized_query = cv2.resize(query_cv, (img_width, img_height))
12
    bordered_query = cv2.copyMakeBorder(resized_query, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=(255, 0, 0))
13
    query_display_area[img_height * (row_count - 1):, :] = cv2.resize(bordered_query, (img_width, img_height))
14
    cv2.putText(query_display_area, "Query", (10, panoramic_height - 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
15

16
    # 处理检索到的图像
17
    for i, img_path in enumerate(retrieved_images):
18
        row, col = i // row_count, i % row_count
19
        start_row, start_col = row * img_height, col * img_width
20

21
        retrieved_pil = Image.open(img_path).convert("RGB")
22
        retrieved_cv = np.array(retrieved_pil)[:, :, ::-1]
23
        resized_retrieved = cv2.resize(retrieved_cv, (img_width - 4, img_height - 4))
24
        bordered_retrieved = cv2.copyMakeBorder(resized_retrieved, 2, 2, 2, 2, cv2.BORDER_CONSTANT, value=(0, 0, 0))
25
        panoramic_image[start_row:start_row + img_height, start_col:start_col + img_width] = bordered_retrieved
26

27
        # 添加索引号
28
        cv2.putText(panoramic_image, str(i), (start_col + 10, start_row + 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
29

30
    return np.hstack([query_display_area, panoramic_image])

1
# 6. 创建索引
2
print(f"\n--> 正在为 '{COLLECTION_NAME}' 创建索引")
3
index_params = milvus_client.prepare_index_params()
4
index_params.add_index(
5
    field_name="vector",
6
    index_type="HNSW",
7
    metric_type="COSINE",
8
    params={"M": 16, "efConstruction": 256}
9
)
10
milvus_client.create_index(collection_name=COLLECTION_NAME, index_params=index_params)
11
print("成功为向量字段创建 HNSW 索引。")
12
print("索引详情:")
13
print(milvus_client.describe_index(collection_name=COLLECTION_NAME, index_name="vector"))
14
milvus_client.load_collection(collection_name=COLLECTION_NAME)
15
print("已加载 Collection 到内存中。")

1
# 7. 执行多模态检索
2
print(f"\n--> 正在 '{COLLECTION_NAME}' 中执行检索")
3
query_image_path = os.path.join(DATA_DIR, "dragon", "query.png")
4
query_text = "一条龙"
5
query_vector = encoder.encode_query(image_path=query_image_path, text=query_text)
6

7
search_results = milvus_client.search(
8
    collection_name=COLLECTION_NAME,
9
    data=[query_vector],
10
    output_fields=["image_path"],
11
    limit=5,
12
    search_params={"metric_type": "COSINE", "params": {"ef": 128}}
13
)[0]
14

15
retrieved_images = []
16
print("检索结果:")
17
for i, hit in enumerate(search_results):
18
    print(f"  Top {i+1}: ID={hit['id']}, 距离={hit['distance']:.4f}, 路径='{hit['entity']['image_path']}'")
19
    retrieved_images.append(hit['entity']['image_path'])

1
# 8. 可视化与清理
2
print(f"\n--> 正在可视化结果并清理资源")
3
if not retrieved_images:
4
    print("没有检索到任何图像。")
5
else:
6
    panoramic_image = visualize_results(query_image_path, retrieved_images)
7
    combined_image_path = os.path.join(DATA_DIR, "search_result.png")
8
    cv2.imwrite(combined_image_path, panoramic_image)
9
    print(f"结果图像已保存到: {combined_image_path}")
10
    Image.open(combined_image_path).show()
11

12
milvus_client.release_collection(collection_name=COLLECTION_NAME)
13
print(f"已从内存中释放 Collection: '{COLLECTION_NAME}'")
14
milvus_client.drop_collection(COLLECTION_NAME)
15
print(f"已删除 Collection: '{COLLECTION_NAME}'")