Введение
В современных задачах компьютерного зрения часто требуется не только анализировать изображения, но и эффективно работать с большими объемами визуальной информации. Для этого необходимо создать векторное представление изображений, которое можно использовать для поиска, классификации и других задач. В данной статье рассматривается процесс создания растрово-векторной базы данных с использованием предобученной нейронной сети ResNet50 для извлечения признаков изображений и библиотеки FAISS для организации быстрого поиска по векторным представлениям.
Шаг 1: Загрузка и предобработка изображений
Основой для создания векторной базы данных является извлечение признаков из изображений. Для этого мы используем предобученную модель ResNet50, которая хорошо справляется с задачами классификации изображений и извлечения признаков.
load_and_preprocess_image(img_path, target_size=(224, 224)):
img = image.load_img(img_path, target_size=target_size)
img_array = image.img_to_array(img)
img_array = np.expand_dims(img_array, axis=0)
img_array = preprocess_input(img_array)
return img_array
Функция load_and_preprocess_image загружает изображение, изменяет его размер до требуемого (224x224 пикселя) и выполняет стандартную предобработку для модели ResNet50. Предобработка включает нормализацию значений пикселей, что важно для корректной работы модели.
Шаг 2: Создание эмбеддингов изображений
После того как изображение подготовлено, необходимо извлечь из него признаки. Мы используем предобученную модель ResNet50 без верхней части сети (классификатора), чтобы извлечь только векторные представления изображений. Для этого модель будет использоваться как средство получения "эмбеддингов" (векторных представлений), которые содержат информацию о содержимом изображения.
def create_embeddings(image_folder, model):
embeddings = []
image_paths = []for img_name in os.listdir(image_folder):
img_path = os.path.join(image_folder, img_name)
if os.path.isfile(img_path) and img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
img_array = load_and_preprocess_image(img_path)
embedding = model.predict(img_array)
embeddings.append(embedding.flatten()) # Преобразуем в 1D вектор image_paths.append(img_path)return np.array(embeddings), image_paths
Функция create_embeddings проходит по всем изображениям в указанной папке, для каждого изображения создает его эмбеддинг с помощью модели ResNet50, а затем сохраняет эти эмбеддинги в список.
Шаг 3: Создание FAISS индекса
Теперь, когда мы имеем эмбеддинги для всех изображений, нужно организовать поиск по этим векторным представлениям. Для этого используется библиотека FAISS (Facebook AI Similarity Search), которая позволяет эффективно строить и использовать индексы для поиска похожих векторов.
def create_faiss_index(embeddings, index_file='faiss_index.idx'):
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension) # Используем L2 расстояние (евклидово) index.add(embeddings) # Добавляем эмбеддинги в индекс faiss.write_index(index, index_file) # Сохраняем индекс в файл
return index
FAISS позволяет создать индекс для векторов с использованием различных типов метрик. В данном случае используется евклидово расстояние (L2), что подходит для сравнения изображений по их признакам.
Шаг 4: Главная функция
Главная функция собирает все шаги вместе: она загружает предобученную модель, извлекает эмбеддинги из изображений и сохраняет результаты в FAISS индекс и файл с путями к изображениям.
def main(image_folder, index_file='faiss_index.idx'):
# Загружаем предобученную модель ResNet50 без классификатора model = ResNet50(weights='imagenet', include_top=False, pooling='avg')# Создаем эмбеддинги для всех изображений в папке embeddings, image_paths = create_embeddings(image_folder, model)
# Создаем и сохраняем FAISS индекс create_faiss_index(embeddings, index_file)
# Сохраняем пути к изображениям для дальнейшего использования with open('image_paths.txt', 'w') as f:
for path in image_paths:
f.write(f"{path}n")print(f"Векторная база данных создана и сохранена в {index_file}")
print(f"Пути к изображениям сохранены в image_paths.txt")
В этой функции:
-
Загружается модель ResNet50 без классификатора.
-
Для каждого изображения в указанной папке извлекаются эмбеддинги.
-
Эти эмбеддинги добавляются в FAISS индекс.
-
Сохраняются пути к изображениям в текстовый файл для дальнейшего использования.
Шаг 5: Запуск
Для того чтобы запустить весь процесс, достаточно указать папку с изображениями и вызвать функцию main.
if name == "__main__":
image_folder = "C:/book/ras/" # Укажите путь к папке с изображениями main(image_folder)
Заключение
Создание растрово-векторной базы данных с использованием предобученной модели ResNet50 и FAISS — это мощный инструмент для организации поиска по изображениям. Векторное представление изображений позволяет легко сравнивать их между собой, а использование FAISS значительно ускоряет процесс поиска схожих изображений в больших базах данных.
Этот подход может быть использован для множества задач, таких как поиск похожих изображений, кластеризация изображений, распознавание объектов и другие приложения в области компьютерного зрения.
P.S. Если интересна данная тема могу добавить статью про поиск в такой базе данных и обучение которое повышает качество поиска.
Автор: stilrambler
