File size: 1,608 Bytes
03d82bf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
from pdf2image import convert_from_path
import os
#inget image
import os
import fitz
import chromadb
from chromadb.utils.data_loaders import ImageLoader
from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction # type: ignore
path = "mm_vdb2"
client = chromadb.PersistentClient(path=path)
def extract_and_store_images2(pdf_path,images_dir=r'extracted_images2'):
output_dir = 'extracted_images2' # Directory to save images
# Ensure the output directory exists
os.makedirs(output_dir, exist_ok=True)
# Convert PDF to a list of images (one per page)
pages = convert_from_path(pdf_path, 300) # 300 dpi is a good resolution
# Save each page as an image (screenshot)
for i, page in enumerate(pages):
output_path = os.path.join(output_dir, f'page_{i + 1}.png')
page.save(output_path, 'PNG')
print(f"Saved: {output_path}")
print("Image extraction complete.")
# Step 2: Add extracted images to ChromaDB
image_loader = ImageLoader()
CLIP = OpenCLIPEmbeddingFunction()
image_collection2 = client.get_or_create_collection(name="image2", embedding_function=CLIP, data_loader=image_loader)
ids = []
uris = []
for i, filename in enumerate(sorted(os.listdir(images_dir))):
if filename.endswith('.jpeg') or filename.endswith('.png'):
file_path = os.path.join(images_dir, filename)
ids.append(str(i))
uris.append(file_path)
image_collection2.add(ids=ids, uris=uris)
print("Images added to the database.")
# return image_vdb
return image_collection2
|