130 lines
3.6 KiB
Python
Executable File
130 lines
3.6 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Build RAG Knowledge Base from Images
|
|
Creates FAISS index from image directory for retrieval-augmented generation
|
|
"""
|
|
import os
|
|
import sys
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from tqdm import tqdm
|
|
from PIL import Image
|
|
from rag.retriever import CLIPRetriever
|
|
|
|
|
|
def build_knowledge_base(
|
|
images_dir: str,
|
|
output_path: str,
|
|
description_file: str = None,
|
|
batch_size: int = 32
|
|
):
|
|
"""
|
|
Build knowledge base from image directory
|
|
|
|
Args:
|
|
images_dir: Directory containing images
|
|
output_path: Where to save FAISS index
|
|
description_file: Optional JSON file with descriptions per image
|
|
batch_size: Batch size for encoding
|
|
"""
|
|
print(f"\n{'='*60}")
|
|
print("🔨 Building Knowledge Base for RAG")
|
|
print(f"{'='*60}\n")
|
|
|
|
# Find all images
|
|
image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.webp'}
|
|
images_dir = Path(images_dir)
|
|
image_files = [
|
|
f for f in images_dir.iterdir()
|
|
if f.suffix.lower() in image_extensions
|
|
]
|
|
|
|
print(f"📂 Found {len(image_files)} images in {images_dir}")
|
|
|
|
# Load descriptions if provided
|
|
descriptions = {}
|
|
if description_file and os.path.exists(description_file):
|
|
with open(description_file, 'r') as f:
|
|
descriptions = json.load(f)
|
|
print(f"📝 Loaded descriptions for {len(descriptions)} images")
|
|
|
|
# Create metadata
|
|
metadata = []
|
|
image_paths = []
|
|
|
|
for img_file in tqdm(image_files, desc="Preparing metadata"):
|
|
img_name = img_file.name
|
|
|
|
# Create metadata entry
|
|
meta = {
|
|
'id': img_file.stem,
|
|
'filename': img_name,
|
|
'path': str(img_file.absolute()),
|
|
'description': descriptions.get(img_name, f"Image: {img_name}"),
|
|
}
|
|
|
|
metadata.append(meta)
|
|
image_paths.append(str(img_file))
|
|
|
|
# Create retriever
|
|
print("\n🔧 Initializing CLIP retriever...")
|
|
retriever = CLIPRetriever()
|
|
|
|
# Build index
|
|
print(f"\n🚀 Encoding images in batches of {batch_size}...")
|
|
vector_store = retriever.build_index(
|
|
images=image_paths,
|
|
metadata=metadata,
|
|
save_path=output_path,
|
|
batch_size=batch_size
|
|
)
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"✅ Knowledge Base Built Successfully!")
|
|
print(f"{'='*60}")
|
|
print(f" Total images: {len(image_paths)}")
|
|
print(f" Index file: {output_path}")
|
|
print(f" Metadata file: {output_path.replace('.faiss', '.json')}")
|
|
print(f" Embedding dimension: {vector_store.dimension}")
|
|
print(f"{'='*60}\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Build RAG knowledge base from images")
|
|
parser.add_argument(
|
|
"--images_dir",
|
|
type=str,
|
|
required=True,
|
|
help="Directory containing images"
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
type=str,
|
|
required=True,
|
|
help="Output path for FAISS index (.faiss)"
|
|
)
|
|
parser.add_argument(
|
|
"--descriptions",
|
|
type=str,
|
|
default=None,
|
|
help="Optional JSON file with image descriptions"
|
|
)
|
|
parser.add_argument(
|
|
"--batch_size",
|
|
type=int,
|
|
default=32,
|
|
help="Batch size for encoding"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
build_knowledge_base(
|
|
images_dir=args.images_dir,
|
|
output_path=args.output,
|
|
description_file=args.descriptions,
|
|
batch_size=args.batch_size
|
|
)
|