239 lines
7.7 KiB
Python
239 lines
7.7 KiB
Python
import os
|
||
import sys
|
||
import json
|
||
import base64
|
||
import time
|
||
from pathlib import Path
|
||
from multiprocessing import Pool
|
||
|
||
# ========== 配置区（直接改这里）==========
|
||
|
||
# 图片目录
|
||
IMAGE_DIR = "./shipin/"
|
||
|
||
# API Keys（填多个就多进程并行）
|
||
API_KEYS = [
|
||
"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",
|
||
"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",
|
||
"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"
|
||
]
|
||
|
||
collection_name = "my_collection1"
|
||
# API 地址
|
||
BASE_URL = "https://integrate.api.nvidia.com/v1"
|
||
|
||
# 模型名称
|
||
MODEL = "nvidia/llama-nemotron-embed-vl-1b-v2"
|
||
|
||
# 输出文件
|
||
OUTPUT = "vectors1.jsonl"
|
||
|
||
# 同一 key 请求间隔（秒），防限流
|
||
DELAY = 0.1
|
||
|
||
# ========== 配置结束 ==========
|
||
|
||
IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.webp'}
|
||
|
||
# ========== 独立函数（供 weixin_bot 等外部调用）==========
|
||
|
||
def get_client(api_key=None, base_url=None):
|
||
"""获取 OpenAI 客户端（单 key 模式，兼容旧接口）"""
|
||
from openai import OpenAI
|
||
return OpenAI(
|
||
api_key=api_key or API_KEYS[0],
|
||
base_url=base_url or BASE_URL
|
||
)
|
||
|
||
|
||
def encode_image_base64(image_path: str) -> str:
|
||
"""将图片编码为 base64"""
|
||
with open(image_path, "rb") as f:
|
||
return base64.b64encode(f.read()).decode("utf-8")
|
||
|
||
|
||
def get_image_vector(image_path: str, client=None) -> list[float]:
|
||
if client is None:
|
||
client = get_client()
|
||
b64, mime = encode_image_base64(image_path) # 确保你的函数能同时返回这些值
|
||
url = f"data:{mime};base64,{b64}"
|
||
input_payload = [
|
||
{
|
||
"content": [
|
||
{"type": "text", "text": "What is in this image?"},
|
||
{"type": "image_url", "image_url": {"url": url}}
|
||
]
|
||
}
|
||
]
|
||
resp = client.embeddings.create(
|
||
input=input_payload, # 使用新格式
|
||
model=MODEL,
|
||
encoding_format="float",
|
||
extra_body={"input_type": "passage","truncate": "NONE", "modality": ["image"]}
|
||
)
|
||
return resp.data[0].embedding
|
||
|
||
|
||
def get_text_vector(text: str, client=None) -> list[float]:
|
||
"""获取文本的向量（外部可直接调用）"""
|
||
if client is None:
|
||
client = get_client()
|
||
resp = client.embeddings.create(
|
||
input=[text],
|
||
model=MODEL,
|
||
encoding_format="float",
|
||
extra_body={"modality": ["text"], "input_type": "query", "truncate": "NONE"}
|
||
)
|
||
return resp.data[0].embedding
|
||
|
||
|
||
def search1(query_text: str, top_k=5, qdrant_host="10.0.0.66", qdrant_port=6333):
|
||
"""在 Qdrant 中搜索（外部可直接调用）"""
|
||
from qdrant_client import QdrantClient
|
||
vector = get_text_vector(query_text)
|
||
qdrant = QdrantClient(host=qdrant_host, port=qdrant_port)
|
||
results = qdrant.query_points(
|
||
collection_name=collection_name,
|
||
query=vector,
|
||
limit=top_k
|
||
)
|
||
return results.points
|
||
|
||
|
||
# ========== 批量向量化（多 key 多进程）==========
|
||
|
||
def scan_images(root_path: str) -> list[str]:
|
||
"""扫描目录下所有图片文件"""
|
||
images = []
|
||
root = Path(root_path)
|
||
if root.is_file() and root.suffix.lower() in IMAGE_EXTENSIONS:
|
||
return [str(root.resolve())]
|
||
for p in sorted(root.rglob("*")):
|
||
if p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS:
|
||
images.append(str(p.resolve()))
|
||
return images
|
||
|
||
|
||
def vectorize(image_path: str, api_key: str, base_url: str, model: str) -> dict:
|
||
from openai import OpenAI
|
||
import mimetypes
|
||
import base64
|
||
try:
|
||
client = OpenAI(api_key=api_key, base_url=base_url)
|
||
|
||
# 获取 base64 和 MIME 类型
|
||
with open(image_path, "rb") as f:
|
||
b64 = base64.b64encode(f.read()).decode("utf-8")
|
||
mime, _ = mimetypes.guess_type(image_path)
|
||
if not mime:
|
||
mime = "image/jpeg"
|
||
|
||
uri = f"data:{mime};base64,{b64}"
|
||
|
||
resp = client.embeddings.create(
|
||
input=[uri],
|
||
model=model,
|
||
encoding_format="float",
|
||
# !!! 关键修改：把 "query" 改为 "passage" !!!
|
||
extra_body={"input_type": "passage", "truncate": "NONE", "modality": ["image"]}
|
||
)
|
||
return {"path": image_path, "vector": resp.data[0].embedding, "status": "ok"}
|
||
except Exception as e:
|
||
return {"path": image_path, "vector": None, "status": "error", "error": str(e)}
|
||
|
||
|
||
def _worker(args) -> dict:
|
||
"""多进程 worker，返回结果同时携带进程信息"""
|
||
image_path, api_key, base_url, model, delay = args
|
||
if delay > 0:
|
||
time.sleep(delay)
|
||
pid = os.getpid()
|
||
# 也可以使用 current_process().name
|
||
result = vectorize(image_path, api_key, base_url, model)
|
||
result["pid"] = pid
|
||
result["api_key_prefix"] = api_key[:10] + "..." # 只显示前几位
|
||
return result
|
||
|
||
|
||
def load_existing(output_path: str) -> dict:
|
||
"""加载已有的向量 JSONL（每行一个 JSON 对象）"""
|
||
result = {}
|
||
if os.path.exists(output_path):
|
||
with open(output_path, "r", encoding="utf-8") as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if line:
|
||
item = json.loads(line)
|
||
result[item["path"]] = item
|
||
return result
|
||
|
||
|
||
def append_vectors(new_items: list[dict], output_path: str):
|
||
"""追加写入新向量到 JSONL 文件（每行一个 JSON 对象）"""
|
||
with open(output_path, "a", encoding="utf-8") as f:
|
||
for item in new_items:
|
||
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||
|
||
|
||
def main():
|
||
images = scan_images(IMAGE_DIR)
|
||
if not images:
|
||
print(f"未在 {IMAGE_DIR} 中找到图片文件")
|
||
sys.exit(1)
|
||
|
||
existing = load_existing(OUTPUT)
|
||
todo = [img for img in images if img not in existing]
|
||
skipped = len(images) - len(todo)
|
||
|
||
if not todo:
|
||
print(f"所有 {len(images)} 张图片已存在于 {OUTPUT}，无需处理")
|
||
return
|
||
|
||
jobs = len(API_KEYS)
|
||
print(f"找到 {len(images)} 张图片，待处理 {len(todo)}，跳过 {skipped}")
|
||
print(f"使用 {jobs} 个 API Key，并行处理")
|
||
|
||
# 构建任务：轮询分配 key，同一 key 的请求错开 delay
|
||
tasks = []
|
||
for i, img in enumerate(todo):
|
||
key = API_KEYS[i % len(API_KEYS)]
|
||
delay = DELAY
|
||
# print(delay)
|
||
tasks.append((img, key, BASE_URL, MODEL, delay))
|
||
|
||
done = 0
|
||
errors = 0
|
||
total = len(existing)
|
||
|
||
with Pool(jobs) as pool:
|
||
for result in pool.imap_unordered(_worker, tasks):
|
||
if result["status"] == "ok":
|
||
done += 1
|
||
append_vectors([{"path": result["path"], "vector": result["vector"]}], OUTPUT)
|
||
total += 1
|
||
# 打印详细信息：进程ID + 文件路径
|
||
print(f"[PID {result['pid']}] 已处理: {os.path.basename(result['path'])}")
|
||
if done % 50 == 0:
|
||
print(f"[进度] 已完成 {done}/{len(todo)}，错误 {errors}，总计 {total}")
|
||
else:
|
||
errors += 1
|
||
print(f"[错误][PID {result.get('pid', '?')}] {os.path.basename(result['path'])}: {result['error']}")
|
||
|
||
print(f"\n处理完毕: 完成 {done}, 跳过 {skipped}, 错误 {errors}")
|
||
print(f"向量已追加到 {OUTPUT}，共 {total} 条记录")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
# from qdrant_client import QdrantClient, models
|
||
# client = QdrantClient(host="10.0.0.66", port=6333)
|
||
# collection_name = "my_collection1" #集合名称
|
||
# client.create_collection(
|
||
# collection_name=collection_name,
|
||
# vectors_config=models.VectorParams(
|
||
# size=2048, # 向量维度，128维
|
||
# distance=models.Distance.COSINE # 距离计算方法，这里用余弦相似度
|
||
# )
|
||
# )
|
||
# client.delete_collection(collection_name="my_collection")
|