/
rnekrasov
/
pgvector-python
Обзор
Документация
Войти
/
rnekrasov
/
pgvector-python
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
examples/hash_image_search.py
43 строки
1 KB
Andrew Kane
Added hash image search example [skip ci]
24 мар 2024, 21:04
24 мар 2024, 21:04
751dc1d
Код
Авторство
О чём код?
from datasets import load_dataset import matplotlib.pyplot as plt import psycopg from imagehash import phash def hash_image(img): return ''.join(['1' if v else '0' for v in phash(img).hash.flatten()]) conn = psycopg.connect(dbname='pgvector_example', autocommit=True) conn.execute('DROP TABLE IF EXISTS images') conn.execute('CREATE TABLE images (id bigserial PRIMARY KEY, hash bit(64))') print('Loading dataset') dataset = load_dataset('mnist') print('Generating hashes') images = [{'hash': hash_image(row['image'])} for row in dataset['train']] print('Storing hashes') cur = conn.cursor() with cur.copy('COPY images (hash) FROM STDIN') as copy: for image in images: copy.write_row([image['hash']]) print('Querying hashes') results = [] for i in range(5): image = dataset['test'][i]['image'] result = conn.execute('SELECT id FROM images ORDER BY bit_count(hash # %s) LIMIT 5', (hash_image(image),)).fetchall() nearest_images = [dataset['train'][row[0] - 1]['image'] for row in result] results.append([image] + nearest_images) print('Showing results (first column is query image)') fig, axs = plt.subplots(len(results), len(results[0])) for i, result in enumerate(results): for j, image in enumerate(result): ax = axs[i, j] ax.imshow(image) ax.set_axis_off() plt.show(block=True)