Введение
Этим летом я взялся разгребать свою изрядно захламленную папку с сохранёнными отовсюду картинками, и естественно это сорочье гнездо хотелось как-то упорядочить, потому что я привык общаться в чатах и на форумах, прилагая аттач с соответствующей реакцией. Когда картинок накопилось достаточно много, копаться в архиве стало уже довольно сложно, так что я начал описывать свою реакцию уже выбирая случайную картинку. Поэтому я решил исследовать, какие вообще сейчас представлены способы индексации картинок.
Самый простой способ - обращаться к внешнему банку данных типа text2image, например к Google Images. Конечно, из минусов - это удаление приглянувшегося файла с сайта, где он хранился.
Отсортировать картинки по папкам. Это привязывает к иерархической структуре хранения и заставляет сортировать все картинки вручную, что при большой коллекции сделать переход довольно болезненным. Кроме того, некоторые картинки можно отнести к нескольким категориям.
Маркировать картинки тегами.
Теги можно хранить непосредственно внутри файла в EXIF-секции, и для них даже будет доступен поиск Проводником. Некоторое время я маркировал коллекцию таким образом, затем упёрся в предел её эффективности, т.к. не все виды файлов поддерживают EXIF.
Теги можно найти во внешних хранилищах (если картинка там есть). Арт-галереи, как правило, дают авторам прикреплять к своим рисункам некоторое количество тегов; движки обратного поиска изображений типа IQDB позволяют добраться до страницы рисунка при имеющемся превью. Минус, конечно, в пропускной способности - разметка всей коллекции при миграции на новую систему займёт некоторое время (у меня заняло три дня с учётом того, что я выставил кулдауны обращения к этим сервисам, чтобы не нагружать их слишком сильно и не нарушать их policy). При массивных коллекциях на десятки тысяч изображений маркировка тегами может растянуться на недели и месяцы.
Взять теги из расшаренной коллекции, которая блуждает по сети пользователей HydrusNetwork. Возможно, для меня это было бы правильным решением, т.к. сейчас это наиболее массивная любительская экосистема для любителей каталогизации картинок, которую я нашёл.
Можно найти нейросеть-классификатор, которая проставит теги без обращения к другим источникам. Это решение ограничено только вычислительными мощностями, которыми располагает хозяин коллекции.
Искать по автоматически создаваемым индексам типа perceptive hash. Они хороши для построение поисковика картинок, но внутри немаркированной коллекции с их помощью можно разве что найти дубликаты для последующего удаления.
Распознавать текст на картинке и искать по нему. Актуально для демотиваторов и всего в этом духе.
Взять для индекса текст от описательные нейросети, строящие фразу на естественном языке, описывающую происходящее на картинке. Пример - BLIP. Пока это довольно любопытно, но пока я не придумал, как это полноценно использовать для поиска.
Итак, наиболее любопытным оказался способ 3.4 - разметка коллекции нейросетью.
Выбор нейросети
При поиске я руководствовался простыми критериями: ленью и любопытством. Ленивая предпосылка означала, что я не стану собирать терабайты картинок (или хотя бы грузить их с Kaggle) и тренировать сетку неделями, чтобы потом размечать коллекцию на несколько гигабайтов. Кроме того, избыточное погружение в не очень хорошо знакомую мне тему нейронных сетей могло погасить мой интерес к пет-проекту. Любопытная предпосылка диктовала подыскать наиболее зрелищное решение из существующих. Поэтому я искал варианты, где кто-то щедрый уже не только разметил датасет, но и выложил pretrained модель.
Итак, я нашёл как минимум два крупных проекта разметки изображений: более академический ImageNet и более хаотический Danbooru (который заблокирован в РФ по известным причинам). В одном больше фотографий, в другом больше рисованных картинок - в зависимости от характера коллекции мог быть ближе тот или иной датасет. На их основе существуют нейросети-классификаторы, у которых можно без проблем найти файлы весов... Вот только не все они выложены в формате, который подходит для проекта.
Изначально я писал быстрый прототип на Electron. Внутрь постепенно утаптывались различные модули типа распознавания текста, ui-фреймворка, так что добавление нейросетки в бандл десктопного приложения было в уже ограниченных условиях, и варианты вырисовывались примерно такие:
Можно эксплуатировать оригинальное решение на популярном фреймворке PyTorch, завернуть его в docker-модуль или в веб-сервер, чтобы оно функционировало в своей естественной среде обитания и время от времени глубокомысленно сообщало наружу, что обнаружило на картинке с лисой кошку, инфа 67%! Это не требовало переписывать на Javascript ничего, но добавляло лишние сущности, которые мне было лень обслуживать.
Можно взять решение на более универсальном фреймворке Tensorflow, построить вокруг него обвязку на TensorflowJS (адаптации под Javascript), и этот вариант даже будет способен читать точку сохранения модели из protobuf-файла, куда она сгружена.
Немного поискав, я обнаружил отличный на первый взгляд вариант:
взять *.pth файл от модели на PyTorch
прогреть его тестовой картинкой
сконвертировать в файл ONNX - специальный промежуточный формат конвертации между разными фреймворками нейросетей.
сконвертировать в protobuf формат, который есть у Tensorflow.
написать обвязку, которая примерно воспроизводит обработку картинки до и после скармливания оригинальной нейросети.
Самым сложным пунктом оказался последний.
Препроцессинг изображений в PyTorch
Итак, если опустить детали, есть код под PyTorch и FastAI в духе
from fastbook import *
from pandas import DataFrame, read_csv
from fastai.callback.progress import ProgressCallback
from os import listdir
from os.path import isfile, join
model_path = "models/model.pth" # тут лежит интересующий файл весов
data_path = "test/tags.csv.gz"
tags_path = "data/tags.json"
df = read_csv(data_path)
vocab = json.load(open(tags_path))
threshold = 0.01
limit = 50
bs = 64
dirpath = '../dataset-for-tagging/original2/'
dblock = DataBlock(
blocks=(ImageBlock, MultiCategoryBlock(vocab=vocab)),
get_x=lambda df: Path("test") / df["filename"],
get_y=lambda df: df["tags"].split(" "),
item_tfms=Resize(224, method=ResizeMethod.Squish),
batch_tfms=[RandomErasing()]
)
dls = dblock.dataloaders(df)
learn = vision_learner(dls, "resnet152", pretrained=False)
model_file = open(model_path, "rb")
learn.load(model_file, with_opt=False)
learn.remove_cb(ProgressCallback)
filepaths = [dirpath+f for f in listdir(dirpath) if isfile(join(dirpath, f))]
tags = {}
for filepath in filepaths:
dl = learn.dls.test_dl([PILImage.create(filepath)], bs=bs)
batch, _ = learn.get_preds(dl=dl) # тут нейросеть делает магию предсказаний
for scores in batch:
df = DataFrame({"tag": learn.dls.vocab, "score": scores})
df = df[df.score >= threshold].sort_values(
"score", ascending=False).head(limit)
tags[filepath] = dict(zip(df.tag, df.score))
print(tags)Из него можно сделать выводы, что нейросеть имеет архитектуру ResNet152, что картинка и что значимая информация - это порядок значений выходных узлов, которые соответствуют тегам. Классификатор берёт самые значимые узлы выходного слоя нейронной сети, сортирует их по значениям и берёт первые 50.
{
"800px-Wikipe-tan_dressed_in_a_Halloween_costume.png": {
"1girl": 0.9686011672019958,
"solo": 0.9571052193641663,
...
"blue_hair": 0.6504012942314148,
"short_hair": 0.5599637627601624,
"halloween": 0.472868412733078,
"blue_eyes": 0.4222845435142517,
"dress": 0.3580203950405121,
"open_mouth": 0.3514541685581207,
"jack-o'-lantern": 0.35024315118789673,
"smile": 0.26754775643348694,
"pumpkin": 0.24310718476772308,
"long_sleeves": 0.2105046510696411,
"simple_background": 0.2009338140487671,
"looking_at_viewer": 0.17024405300617218,
"original": 0.1686396598815918,
"blush": 0.12935996055603027,
"hat": 0.11987852305173874,
...
"black_background": 0.11153995990753174,
"touhou": 0.10229469835758209,
"bow": 0.09583188593387604,
":d": 0.09576313942670822,
...
"standing": 0.02223421074450016
},
...Не идеально, но и не совсем плохо. Поскольку в исходном датасете было меньше 1000 изображений с этим персонажем - её тег не попал в выборку; а некоторые теги включены в выдачу по ошибке. Но в остальном качество меня устраивало до некоторой степени, поэтому после некоторого размышления я решил включить нейросеть внутрь основного проекта как ассистента: не доверять ей расстановку всех тегов (чтобы не замусорить разметку), а подсказывать по востребованию список тегов, чтобы выбрать из них наиболее подходящие. К сожалению, для этого подхода надо уже знать, что и кто изображены на картинке.
Сконвертировав файл весов нейросети и немного потыкав веточкой в него с разных сторон, опытным путём можно выяснить, что он находится в режиме serve, что на вход input.1 он принимает тензор с размерностью [1, 3, 224, 224], выход у него называется ret.11, а внутри у него чёрный ящик. Возможно, при экспорте можно обозвать входы и выходы как-нибудь покрасивее взамен сгенерированных автоматически, но в данном случае это особой роли не играет.
Проделаем аналогичный препроцессинг картинки внутри Tensorflow.
PyTorch традиционно принимает на вход цветовые каналы в другом порядке, чем TensorFlow, поэтому мне понадобилось преобразовать тензор, иначе код в принципе не запускался без ошибки. Почему так? Я слышал, PyTorch больше ориентирован на вычисления на видеокарте по сравнению с Tensorflow, а такая конфигурация входного тензора экономит копеечку скорости.
Также исходная нейросеть была натренирована на пакетный приём изображений, поэтому понадобится увеличить размерность тензора, положив таким образом картинку в пакет, где она будет одна. В TensorflowJS для этого есть, например, метод expandDims.
tensor = tensor.expandDims().transpose([0, 3, 1, 2]);Итого:
const tf = require('@tensorflow/tfjs-node');
const fs = require('fs');
const path = require('path');
const tags = require('../pytorch-autotagger/data/tags.json');
const LIMIT = 50;
const PREPARED_IMAGE_SIZE = 224;
const MAX_COLOR_VALUE = 255;
async function getSortedTags(filepath) {
let model = await tf.node.loadSavedModel(
'./model/danbooru/',
['serve'],
'serving_default'
);
return tf.tidy(() => {
let tensor = tf.node
.decodeImage(fs.readFileSync(filepath), 3)
.resizeBilinear([PREPARED_IMAGE_SIZE, PREPARED_IMAGE_SIZE]);
tensor = tensor.expandDims().transpose([0, 3, 1, 2]); // move color channel to 2nd place
let scores = model.predict({ 'input.1': tensor })['ret.11'];
let scoredTags = [];
scores = scores.dataSync();
for (let i in scores) {
scoredTags.push({ score: scores[i], tag: tags[i] });
}
let sortedScoredTags = scoredTags.sort((a, b) => a.score - b.score);
return sortedScoredTags
.slice(sortedScoredTags.length - LIMIT, sortedScoredTags.length)
.reverse();
});
}
/**
* @param {string} dirPath
* @returns {string[]}
*/
function flattyReadDir(dirPath) {
return fs
.readdirSync(dirPath, { withFileTypes: true })
.filter((f) => f.isFile())
.map((f) => path.join(dirPath, f.name));
}
(async () => {
let result = {};
let filepaths = flattyReadDir('../dataset-for-tagging/original/');
for (let i = 0; i < filepaths.length; i++) {
let filepath = filepaths[i];
result[filepath] = (await getSortedTags(filepath)).reduce(
(accum, current) => {
accum[current['tag']] = current['score'];
return accum;
},
{}
);
}
console.log(result);
})();
И получим пшик. Вместо вероятностей тегов от нуля до единицы на выходе плавают трёхзначные числа. Порядок был каким угодно, но определённо не тем, который был в оригинале.
{
"800px-Wikipe-tan_dressed_in_a_Halloween_costume.png": {
"open_mouth": 277.4236755371094,
":d": 159.86825561523438,
"long_hair": 135.67080688476562,
"brown_eyes": 123.95932006835938,
"blush": 117.69956970214844,
"1girl": 114.89971923828125,
...
"fang": 11.402185440063477,
"red_eyes": 4.072291374206543,
"short_hair": -1.1306633949279785,
"large_breasts": -3.761528968811035,
"purple_hair": -11.944032669067383,
"collarbone": -12.208208084106445,
"hair_ornament": -15.728532791137695,
"shiny": -18.439069747924805,
},Немного поискав и поспрашивав, я выяснил что нужно поделить значения цветовых каналов на 255 (чтобы они укладывались в промежуток [0, 1] и вычисления несильно раздували выходные значения).
const MAX_COLOR_VALUE = 255;
// ...
tensor = tensor.div(MAX_COLOR_VALUE);Для сравнения результатов с образцовым результатом от PyTorch я написал простой скрипт, который брал два выходных json-файла после эксперимента и считал количество тегов, которые совпадали у файлов, а затем считал итоговый процент совпадений на всей тестовой мини-выборке картинок. Функция расстояния, таким образом, выглядела примерно таким образом:
// ...
function hammingDistanceBetweenSets(list1, list2) {
let count = 0;
for (let i = 0; i < list1.length; i++) {
if (!list2.includes(list1[i])) {
count = count + 1;
}
}
return count;
}
// ...Итого:
$ node compare-json-files.js tensorflow-1 pytorch-autotagger
82% 800px-Mural_in_Heidelberg.jpg
72% 800px-Neko_Wikipe-tan.png
66% 800px-Wikipe-tan_dressed_in_a_Halloween_costume.png
72% 800px-Wikipe-tan_mopping.png
86% 800px-Пасущаяся_лошадь_и_грозовое_небо.jpg
80% Dierentegel,_vliegende_vogel,_hoekmotief_ossenkop,_objectnr_17872.jpg
84% Tiere_an_einem_Hang_des_Fahrentalsgrabens_2.jpg
92% Utzenaich_(Gemeindeamt).jpg
64% Wiki-sisters.png
72% Wikipe-tan_donations.png
76% Wikipe-tan_the_Library_of_Babel.png
----
76.9090909090909% total
$ node compare-json-files.js tensorflow-2 pytorch-autotagger
18% 800px-Mural_in_Heidelberg.jpg
8% 800px-Neko_Wikipe-tan.png
18% 800px-Wikipe-tan_dressed_in_a_Halloween_costume.png
24% 800px-Wikipe-tan_mopping.png
6% 800px-Пасущаяся_лошадь_и_грозовое_небо.jpg
14% Dierentegel,_vliegende_vogel,_hoekmotief_ossenkop,_objectnr_17872.jpg
22% Tiere_an_einem_Hang_des_Fahrentalsgrabens_2.jpg
10% Utzenaich_(Gemeindeamt).jpg
8% Wiki-sisters.png
26% Wikipe-tan_donations.png
22% Wikipe-tan_the_Library_of_Babel.png
----
16% totalТаким образом, есть метрика, благодаря которой исследование будет не окончательно слепым.
Первый подход хуже второго, но и второй даёт искажение результатов.
Переворошив гугл-выдачу, я построил следующие гипотезы:
1) Нужно не только нормализовать значения цветовых каналов как в PyTorch, нужно сместить цветовое пространство, как иногда оно используется в PyTorch. Это предлагается не во всех моделях обучения.
tensor = tensor
.div(MAX_COLOR_VALUE)
.sub([0.485, 0.456, 0.406])
.div([0.229, 0.224, 0.225]);Результат: оценка отклонений ухудшилась до 69.1%.
2) Надо добавить для выравнивания на вход softmax, уж это точно не испортит дело!
tensor = tensor
.div(MAX_COLOR_VALUE)
.softmax();Результат: оценка отклонений ухудшилась до 44.7%.
3) Я перепутал местами цветовые каналы и они лежат как-то иначе:
tensor = tensor.expandDims().transpose([0, 3, 2, 1]);Оценка отклонений ухудшена до 30.5%
4) Искажений вносит ресайз картинки до квадрата 224х224, который по-разному устроен в PyTorch и в Tensorflow. Я подготовил вспомогательный датасет из картинок, заранее сжатых до 224х224, после чего опробовал на них лучшие найденные алгоритмы из обоих фреймворков нейросетей.
$ node compare-json-files.js tensorflow-2-squared-224 pytorch-autotagger-squared224
2% 800px-Mural_in_Heidelberg.jpg
...
0% Wikipe-tan_the_Library_of_Babel.png
----
0.54% totalИ это очень хороший результат!
К сожалению, он так же хорош, как и бесполезен.
Я визуализировал процесс подготовки картинки в обоих фреймворках.
for filename in filenames:
image = PILImage.create(
dirpath_for_original + filename)
rsz = Resize(224, method=ResizeMethod.Squish)
rsz(image, split_idx=0).save(dirpath_for_squared + filename)
let tensor = tf.node
.decodeImage(fs.readFileSync(filepath), 3)
.resizeBilinear([PREPARED_IMAGE_SIZE, PREPARED_IMAGE_SIZE]);
tf.node.encodeJpeg(tensor, 'rgb').then((image) => {
fs.writeFileSync(
newFilepath,
Buffer.from(image)
);
У Tensorflow обнаружились некоторые проблемы с антиалиасингом, поэтому я попробовал использовать библиотеку sharp. Как ни странно, в исходном коде на PyTorch используется по умолчанию билинейная интерполяция (по крайней мере я так предположил из чтения исходников), однако она дала лучшее сглаживание краёв, чем Tensorflow.
const PREPARED_IMAGE_SIZE = 224
// ...
let sharpedImage = sharp(filepath);
sharpedImage = sharpedImage.resize(PREPARED_IMAGE_SIZE, PREPARED_IMAGE_SIZE, {
fit: sharp.fit.cover,
kernel: sharp.kernel.mitchell,
});
let buffer = await sharpedImage.toBuffer();
let tensor = tf.node.decodeImage(buffer, 3);
tf.node.encodeJpeg(tensor, 'rgb').then((image) => {
fs.writeFileSync(
newFilepath,
Buffer.from(image)
);
});
Библиотека Sharp действительно решала проблемы с зубчатыми краями, но добавляла другую проблему: по какой-то причине она иначе интерпретировала альфа-канал прозрачности, и это вносило собственные искажения. Эта библиотека предлагает 5 способов интерполяции при изменении размеров, и перебором мне удалось снизить процент ошибок с 16% до 13% на первоначальной тестовой выборке (kernel: sharp.kernel.mitchell).
Впрочем, на более крупном датасете один из самых первых алгоритмов препроцессинга картинки для TensorflowJS оказался и самым близким по точности к тому, что выдаёт нейросеть на PyTorch.
Постпроцессинг
Поскольку я хотел сохранить прежнее впечатление от результата выполнения "Это charactername, инфа сотка!", то результирующие значения приглаживаются: нормализуются внутрь диапазона [0, 1] и затем при рендеринге результата переводятся в проценты. Это не влияет на порядок выходных узлов (т. к. преобразование линейно и его множитель положителен), поэтому выдача тегов не искажается этим действием.
При использовании решения на PyTorch можно было бы обойтись без этого шага.
let scores = model.predict({ 'input.1': tensor })['ret.11'];
scores = scores
.sub(tf.min(scores).dataSync()[0])
.div(tf.max(scores).dataSync()[0] - tf.min(scores).dataSync()[0]);
Выводы
Несмотря на различия между фреймворками нейросетей, вы действительно можете конвертировать файлы весов между ними благодаря проекту-посреднику ONNX, чтобы подобрать архитектуру на одном фреймворке, а выложить в продакшн на другом. Вы даже сможете сменить язык программирования для работы с этой нейросетью. Но к сожалению, различие инструментов исказит точность первоначального результата, если вы не сможете точно воспроизвести алгоритм препроцессинга входных данных.
Ссылки
https://github.com/danbooru/autotagger - найденный репозиторий с исходной моделью под PyTorch.
https://github.com/tauraloke/tegownitsa - менеджер тегов с адаптированной моделью под TensorflowJS.
https://github.com/tauraloke/autotagger-paper-listings - листинги экспериментов, которые были показаны в статье.
