File size: 2,879 Bytes
ccdd4a4 |
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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import chromadb
import json
COLLECTION_NAME = 'video_analysis_data'
# ... (rest of imports/constants)
def generate_text_summary(record):
"""
Converts a structured detection record into a natural language text description
by summarizing all detected objects clearly.
"""
video_id = record['video_id']
timestamp = record['timestamp_sec']
detections = record['detections']
if not detections:
return f"Analysis of video '{video_id}' at {timestamp} seconds: No objects were detected in this frame."
# Group detections by label for a complete object count summary
object_counts = {}
for det in detections:
label = det['label']
object_counts[label] = object_counts.get(label, 0) + 1
summary_parts = []
if object_counts:
# Format: N instances of 'label', M instances of 'other_label', etc.
object_descriptions = [
f"{count} instances of '{label}'"
for label, count in object_counts.items()
]
summary_parts.append("Detected objects include: " + ", ".join(object_descriptions) + ".")
summary_doc = f"Analysis of video '{video_id}' at {timestamp} seconds: {' '.join(summary_parts)}"
return summary_doc
def index_analysis_data(json_file='raw_analysis.json', collection_name='video_analysis_data'):
"""
Loads raw analysis, generates documents, and indexes them in ChromaDB.
"""
try:
with open(json_file, 'r') as f:
raw_data = json.load(f)
except FileNotFoundError:
print(f"Error: {json_file} not found. Run 'video_analyzer.py' first.")
return
# Initialize ChromaDB client
client = chromadb.PersistentClient(path="./chroma_db") # Stores data locally
# Changed collection name to be more generic
collection = client.get_or_create_collection(name=collection_name)
documents = []
metadatas = []
ids = []
print(f"Indexing {len(raw_data)} analysis records...")
for i, record in enumerate(raw_data):
doc_text = generate_text_summary(record)
if doc_text:
documents.append(doc_text)
# Metadata is crucial for filtering and context
metadatas.append({
'video_id': record['video_id'],
'timestamp_sec': record['timestamp_sec'],
'frame_id': record['frame_id']
})
ids.append(f"doc_{i}")
# ChromaDB automatically handles embedding and storage
if documents:
collection.add(
documents=documents,
metadatas=metadatas,
ids=ids
)
print(f"Successfully indexed {len(documents)} documents into ChromaDB collection '{collection_name}'.")
else:
print("No valid documents generated for indexing.")
if __name__ == '__main__':
index_analysis_data() |