wehub-resource-sync
77 lines
2.3 KiB
Python
77 lines
2.3 KiB
Python
import pixeltable as pxt
|
|
import os
|
|
from pixeltable.functions.openai import vision
|
|
from pixeltable.functions.huggingface import sentence_transformer
|
|
|
|
# Set OpenAI API key (replace with your actual key or use an environment variable)
|
|
os.environ['OPENAI_API_KEY'] = 'your-openai-api-key-here'
|
|
|
|
# Base directory for the index
|
|
DIRECTORY = 'image_search'
|
|
TABLE_NAME = f'{DIRECTORY}.images'
|
|
RECREATE = True
|
|
|
|
# Recreate the directory if specified
|
|
if RECREATE:
|
|
pxt.drop_dir(DIRECTORY, force=True)
|
|
|
|
# Check if table exists, create it if not
|
|
if TABLE_NAME not in pxt.list_tables():
|
|
# Create directory and table
|
|
pxt.create_dir(DIRECTORY, if_exists='ignore')
|
|
image_index = pxt.create_table(
|
|
TABLE_NAME,
|
|
{'image_file': pxt.Image},
|
|
if_exists='ignore'
|
|
)
|
|
|
|
# Add GPT-4 Vision analysis
|
|
image_index.add_computed_column(
|
|
image_description=vision(
|
|
prompt="Describe the image. Be specific on the colors you see.",
|
|
image=image_index.image_file,
|
|
model="gpt-4o-mini"
|
|
)
|
|
)
|
|
|
|
# Define the embedding model and create embedding index
|
|
embed_model = sentence_transformer.using(model_id='intfloat/e5-large-v2')
|
|
image_index.add_embedding_index(
|
|
column='image_description',
|
|
string_embed=embed_model,
|
|
if_exists='ignore'
|
|
)
|
|
else:
|
|
image_index = pxt.get_table(TABLE_NAME)
|
|
|
|
# Sample image URLs
|
|
IMAGE_URL = (
|
|
"https://raw.github.com/pixeltable/pixeltable/release/docs/resources/images/"
|
|
)
|
|
image_urls = [
|
|
IMAGE_URL + doc for doc in [
|
|
"000000000030.jpg",
|
|
"000000000034.jpg",
|
|
"000000000042.jpg",
|
|
]
|
|
]
|
|
|
|
# Insert images into the table
|
|
image_index.insert({'image_file': url} for url in image_urls)
|
|
|
|
# Perform a sample query
|
|
query_text = "Show me images containing blue flowers"
|
|
sim = image_index.image_description.similarity(query_text)
|
|
results = (
|
|
image_index.order_by(sim, asc=False)
|
|
.select(image_index.image_file, image_index.image_description, sim=sim)
|
|
.limit(3)
|
|
.collect()
|
|
)
|
|
|
|
# Print results
|
|
print(f"Query Results for '{query_text}' in '{TABLE_NAME}':\n")
|
|
for i, row in enumerate(results.to_pandas().itertuples(), 1):
|
|
print(f"{i}. Score: {row.sim:.4f}")
|
|
print(f" Description: {row.image_description}")
|
|
print(f" Image URL: {row.image_file}\n") |