File size: 1,979 Bytes
38ab39c | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | from cora_vision import CoraVision
from cora_memory import CoraMemory
from PIL import Image
import os
def create_dummy_image():
img = Image.new('RGB', (256, 256), color = 'red')
img.save('test_red.png')
return 'test_red.png'
def test_stack():
# 1. Init
print("--- Initializing Stack ---")
vision = CoraVision()
memory = CoraMemory()
# 2. Get Image
img_path = create_dummy_image()
prompt = "A red square test image"
print(f"Created dummy image: {img_path}")
# 3. Vision Analysis
print("--- 1. The Eyes (Vision) ---")
embedding = vision.embed_image(img_path)
print(f"Embedding generated: Vector length {len(embedding)}")
tags = vision.detect_tags(img_path)
print(f"Tags detected: {tags}")
# 4. Memory Storage
print("--- 2. The Memory (Storage) ---")
entry_id = memory.save(img_path, embedding, prompt, tags)
print(f"Stored with ID: {entry_id}")
# 5. Search
print("--- 3. Retrieval (Search) ---")
# Search by the same image embedding (should find itself)
results = memory.search_by_vector(embedding, k=1)
print("Search Results:")
try:
# ChromaDB API structure v0.4+
ids = results['ids'][0]
metadatas = results['metadatas'][0]
distances = results['distances'][0]
for i, uid in enumerate(ids):
print(f"Found: {uid}")
print(f" Metadata: {metadatas[i]}")
print(f" Distance: {distances[i]}")
if entry_id in ids:
print("SUCCESS: Retrieved the saved entry!")
else:
print("FAILURE: Did not retrieve the entry.")
except Exception as e:
print(f"Error parsing results: {e}")
print(f"Raw results: {results}")
# Clean up
if os.path.exists(img_path):
os.remove(img_path)
if __name__ == "__main__":
test_stack()
|