utils/search.py

from utils.setup import load_collection
from utils.get_embeddings import preprocess_image, extract_features, create_resnet18_model

search_params = {
    "metric_type": "L2",
    "params": {"level": 2}
}

collection = load_collection()
model = create_resnet18_model()

def search_images(image_path, topk=6):
    processed_image = preprocess_image(image_path)
    processed_image = extract_features(model, processed_image)
    results = collection.search(
        [processed_image],
        anns_field="image_embedding",
        param=search_params,
        limit=topk,
        guarantee_timestamp=1
    )
    
    return results[0]