Как научить компьютер различать цвета?

Привет!

Думаю, все заметили, что технологии компьютерного зрения и искусственного интеллекта появились во многих сферах нашей жизни. Аналитика изображений применяется на производстве, в медицине, в системах управления, в географии.

В сфере спорта также не обошлось без компьютерного зрения. Современные технологии используются для анализа игр и стратегий, предоставляя тренерам и спортсменам ценную информацию, которую сложно или невозможно получить вручную. В частности, в футболе компьютерное зрение используется для трекинга, детектирования действий игроков, вычисления их скоростей и решения многих других задач. Аналитика матча вручную может занимать много времени и требует внимания квалифицированных специалистов, а технологии компьютерного зрения позволяют значительно автоматизировать и оптимизировать процессы.

Одна из задач, которая была поставлена в рамках нашего проекта по футбольной аналитике, заключается в определении команды, к которой принадлежит конкретный игрок.

Для человеческого глаза эта задача — проще некуда. Достаточно посмотреть на цвет формы игрока, чтобы понять, в какой он команде. Однако как научить компьютер автоматически понимать цвет? Оказалось, что в общем случае эта задача совсем не тривиальна.

Изначально были попытки обучить на размеченных данных классические детекторы объектов, например, YOLO. Однако оказалось, что модель сильно усложняет задачу — начинает ориентироваться на расстановку игроков, и из-за этого точность становится невысокой. Поэтому возникла идея попробовать более классические и, соответственно, более простые алгоритмы.

Допустим, что решение компьютером также будет приниматься на основе цвета формы игрока. Тогда для каждого игрока на поле нужно получить вектор цвета, определяющий его.

Общий случай

Выделение векторов

Для начала будем считать, что нам известно положение игроков на поле. Детектирование выполняется с помощью модели YOLOv8, дообученной для распознавания игроков, рефери, вратаря и мяча. Для классов людей также можно получить ключевые точки на теле (для этого используется модель PoseNet). Зная, где находятся ключевые точки, можно выбрать области с цветом - взять часть от футболки (тогда размерность вектора-эмбеддинга игрока равна 3), или же взять две части — от футболки и шорт и конкатенировать (здесь размерность 6). Выделенные части показаны на рисунке:

Каждая часть переводится из формата RGB в LAB (так как это позволяет отделить компоненту яркости Lightness от цветовой составляющей). Для дальнейшей разработки выбран именно вариант 6D, так как он несет в себе больше информации и позволяет справляться с такими вариантами формы, когда футболки одинакового цвета, а шорты разные (да, в некоторых чемпионатах это допустимо!). 

Такие векторы могут быть получены как по части видео, так и по полному матчу. Однако для дальнейшего анализа все векторы не нужны, достаточно лишь репрезентативной выборки.

Примеры полученных кластеров из векторов показаны ниже, визуализация сделана через программу Meshlab. Цвет точки — это цвет, кодирующий также местоположение вектора.

Алгоритм разделения команд

Векторы, полученные для игроков одной команды, будут близки в пространстве. Для того, чтобы отличить группы друг от друга, используются методы обучения без учителя, так как на данном этапе меток команды нет. 

В качестве основного метода выбрана кластеризация, основанная на плотности распределения векторов в пространстве. При настройке кластеризатора используются два основных параметра: eps — радиус окрестности, в которой идет подсчет точек; num_samples — количество точек в окрестности данной, необходимое для того, чтобы она принадлежала кластеру. При поиске наилучших параметров можно рассматривать постоянное значение eps и переменное значение num_samples — это позволяет варьировать плотность кластеров, требуемую для кластеризации.

Однако на практике стало понятно, что кластеры могут быть разными по плотности. Например, игроков одной команды лучше видно из-за контрастности относительно поля, их может быть просто больше на имеющемся отрезке видео. Более того, эмбеддинги вратарей и рефери также могут сгруппироваться в отдельные кластеры, но с существенно меньшей плотностью. Стало понятно, что параметр num_samples нужно подбирать для каждого отдельного случая.

Чтобы не увеличивать число допущений и не фиксировать какие-то пороги, была придумана идея каскада кластеризации. Сначала выполняется кластеризация при малой начальной плотности (num_samples=5). Если в результате имеется кластер, который слишком большой, то на нем применяется следующий кластеризатор с увеличенным параметром num_samples (то есть с повышенной начальной плотностью). Кластер подлежит разбиению, если его размер больше 55% от всего объема векторов. Этот параметр появился, исходя из количества игроков на поле, на фрагменте  видео игроков каждой команды в среднем чуть меньше 50%. По итогу получается так, что все итоговые кластеры должны быть не более половины объема всего облака векторов. В алгоритме есть ограничение на количество разбиений, так как значения num_samples предустановлены. 

Безусловно, любая кластеризация получает не только метки кластеров, но и достаточно большое количество шума. Для его устранения также есть решение. 

Технические подробности алгоритмов кластеризации

Во всем алгоритме используется кластеризатор DBSCAN, кроме начального этапа  — на нем кластеризация выполняется алгоритмом HDBSCAN. 

DBSCAN — алгоритм кластеризации, основанный на плотности распределения точек. Для выполнения кластеризации точки делятся на основные точки, достижимые по плотности точки и точки, выпадающие следующим образом:

• Точка p является основной точкой, если по меньшей мере num_samples точек находятся на расстоянии, не превосходящем eps (eps является максимальным радиусом соседства от p), до неё (включая саму точку p). Говорят, что эти точки достижимы прямо из p.

• Точка q прямо достижима из p, если точка q находится на расстоянии, не большем eps, от точки p и p должна быть основной точкой.

• Точка q достижима из p, если имеется путь p_1, …, p_n с p_1=p и p_n=q, где каждая точка p_i+1 достижима из точки p_i (все точки на пути должны быть основными, за исключением q)

• Все точки, не достижимые из основных точек, считаются выбросами.

Теперь, если p является основной точкой, то она формирует кластер вместе со всеми точками (основными или неосновными), достижимыми из неё. Каждый кластер содержит по меньшей мере одну основную точку. Неосновные точки могут быть частью кластера, но они формируют его «край», поскольку не могут быть использованы для достижения других точек.

HDBSCAN — кластеризатор из sklearn-contrib, работает как "надстройка" над исходным DBSCAN, но с переменным параметром epsilon. Это позволяет на одной итерации находить кластеры разных плотностей, а также четче определять границы в тех областях кластера, где его плотность падает. По сравнению с DBSCAN, в качестве начального кластеризатора HDBSCAN получает значительно меньше шума, что исключает ошибку алгоритма устранения шума на этих векторах.

Исправление шума

При кластеризации методами HDBSCAN или DBSCAN появляется много шума, особенно при большой размерности векторов. Для того, чтобы его устранить, используется следующий алгоритм:

1. Для кластеров, чей объем превышает 5%, собирается словарь, состоящий из средних значений векторов этих кластеров. Векторы шума могут быть присвоенными только к этим кластерам. Так кластеры размером в несколько точек не будут стягивать на себя весь шум.

2. Вычисляется матрица расстояний от всех шумовых векторов до медианных (значения словаря, созданного в п.1)

3. В матрице расстояний находится первое и второе по величине расстояния до кластеров (первый и второй минимум). Вычисляется отношение второго минимума к первому.

4. Если данная пропорция больше какого-то порогового значения, то вектор считается отнесенным к ближайшему кластеру, иначе — остается шумом.

Шум устраняется в 2 этапа: сначала для футболок, затем для шорт (при размерности 6D). Порог задается конкретным числом, полученным в результате эксперимента. В среднем на каждом этапе устраняется порядка 60% шума.

Предсказание для новых векторов

Как известно, модели кластеризации не предусматривают функционала для предсказания кластера для новой точки. Обучать же алгоритм кластеризации на всем видео слишком затратно по времени, так как количество векторов будет исчисляться миллионами. Для того, чтобы выполнять тренировку на части векторов из данного видео, а для остальных векторов просто предсказывать кластер, используется классификатор KNN (K ближайших соседей). На последнем этапе он обучается на метках, полученных кластеризатором, и затем используется для предсказания каждого нового эмбеддинга игрока. Это позволяет предсказывать команду быстро и не переобучать алгоритм заново. 

Итоговый результат

Пример получившегося результата представлен на кадрах ниже.

Особенности реализации для неравномерного освещения

На видеосъемке матчей могут быть разные варианты освещения поля. Освещение может быть искусственным, тогда на поле нет теней, т.к. софиты расположены со всех сторон. Также используется естественное освещение, в этом случае на поле появляется тень от стадиона. 

Для равномерного освещения все понятно — цвет формы игрока может меняться, так как в область, по которой строится эмбеддинг, попадает номер (а он другого цвета), трава или другие игроки. Однако в целом кардинально цвет остается в тех же границах. В том случае, если на поле присутствует тень от крыши стадиона, а остальная часть ярко залита солнцем, один и тот же комплект формы может выглядеть по-разному. Эта проблема решается этим же алгоритмом, с одним небольшим дополнением, описанным далее.

Разделение векторов на группы

Для неравномерного освещения нужно получить два набора векторов — из солнечной и теневой части поля. Значит, нужно построить маску, разделяющую поле на эти части. Можно было бы применить обычную бинаризацию по порогу, однако оказалось, что порог подобрать не так просто. Цвета поля и яркость освещения меняются в зависимости от видео, места и времени съемки. 

Для того, чтобы адаптивно подобрать порог, изображение сначала переводится в черно-белый формат и обрабатывается с помощью преобразования Фурье. В преобразованном изображении выделяются только основные частоты. Это позволяет убрать из рассмотрения игроков, разметку и другие мелкие детали, которые не важны для определения освещенности поля (но при этом могут своим цветом вносить вклад в итоговое распределение пикселей). Затем через обратное преобразование Фурье получаем уже размытое черно-белое изображение. У изображений, имеющих границу света и тени, цвет пикселей различается значительно. Тогда, если построить гистограмму такого изображения, на ней будут отчетливо видны два пика. Таким образом, изображение бимодальное, на нем есть светлая и теневая часть, а значит, с помощью адаптивного алгоритма бинаризации Отсу можно выбрать порог, который корректно разделит изображение на две части. В том случае, если пик на гистограмме единственный, можно сказать, что кадр либо полностью залит солнцем, либо полностью находится в тени. Тогда маска проходит дополнительный этап: исходное изображение переводится в формат LAB, и по среднему значению компоненты L (освещенность) определяется принадлежность к солнечной или теневой части.

На рисунках представлены примеры подобного разделения: первая колонка — исходный черно-белый кадр, вторая — кадр после преобразования Фурье, третья — гистограмма, в четвертой колонке — полученная маска.

Таким образом, векторы разделены на две части — солнечную и теневую. Каждая из групп векторов успешно пройдет кластеризацию и получит свои метки, так же, как и для случая с равномерным освещением. После этого полученные классы нужно совместить. 

Совмещение кластеров

Для совмещения кластеров нужно создать какое-то их представление, которое будет наиболее полно их описывать и при этом позволит вычислить какую-то метрику похожести. Представим кластер как объект в пространстве эмбеддингов, составленный из векторов, являющихся квантилями определенного порядка (от 0.05 до 0.95 включительно с шагом 0.05). Этот объект по сути — набор точек в пространстве эмбеддингов, для общности также выполняется нормализация. Такой подход позволяет представить кластер не как его наиболее частого представителя, а с учетом большей части возможных изменений цвета игрока в данной части поля. 

На самом деле, случай с тенью на поле также может быть разбит на два: когда тень четкая (clear_shadow) и когда тень частичная (partial_shadow). 

1. Для режима clear_shadow характерно то, что тень четкая, отсутствуют различные полутени и места с освещением средней интенсивности. За счет этого форма футболистов выглядит по-разному на разном свете, практически нет такой ситуации, когда футболист стоит в тени, а его освещение при этом такое, как будто он стоит на солнце. 

2. Для режима partial_shadow характерно то, что тень нечеткая (редкий случай). Это значит, что есть фрагменты с полутенью, где преобразование Фурье срабатывает неточно и вместо четкой маски с контуром появятся пятна. Таким образом, в векторы одного освещения могут попадать векторы противоположного, и в процессе кластеризации алгоритм выделит их в отдельный кластер. Поэтому перед тем, как совмещать лейблы на солнце и в тени, необходимо переместить те векторы, которые оказались в неверном кластере, в противоположный.

Примеры показаны на рисунке ниже:

Алгоритм совмещения кластеров, которые попали в неверное освещение

Данный алгоритм используется в режиме partial_shadow.

1. Отдельно для солнечных и теневых векторов создается набор прямых, описанных выше.

2. Если количество таких прямых для солнечной или теневой части меньше или равно 2, то совмещение выполнить невозможно, так как недостаточно крупных кластеров.

3. Рассчитывается матрица расстояний между прямыми из солнечной и теневой части. Используется метрика — модифицированное Хаусдорфово расстояние. Оно основано на Хаусдорфовом расстоянии, но берет сумму минимальных расстояний от каждой точки одной кривой до каждой точки другой (вместо максимальной из них).

4. Берутся минимальные значения из данной матрицы.

5. Если минимальное значение меньше, чем 10 (на визуализации это очевидное пересечение и совпадение прямых), то данная пара - один и тот же кластер. Тогда выполняется перемещение того кластера, который меньше в процентном соотношении, к большему. Происходит перемещение векторов, переименование лейблов, пересчет словаря.

6. Возвращаются новые векторы, новые лейблы, новые словари.

Примеры совмещения векторов:

1. Прямые до и после проведенного совмещения:

   

2. Прямые в 3D до объединения (построены первые три компоненты из 6D-вектора):

   

3. Прямые в 3D после объединения (чуть смещены из-за нормализации при построении):

   

Очевидно, совместились кластерами с индексами 12 и 2.

Алгоритм совмещения цветов при разном их освещении

Данный алгоритм используется в режиме clear_shadow после кластеризации в каждом освещении и в режиме partial_shadow после совмещения кластеров, попавших в неверное освещение.

1. Отдельно для солнечных и теневых векторов создается набор прямых, описанных выше.

2. Аналогичным образом рассчитывается матрица расстояний между прямыми из солнечной и теневой части,  метрика — модифицированное Хаусдорфово расстояние.

3. Берутся минимальные значения. 

4. Теперь заменяем индексы на лейблы, которые дали данные минимальные значения.

5. Возвращается массив пар [sun_lbl, shadow_lbl]. В одной паре значения лейблов, первое соответствует кластеру на солнце, второе - кластеру в тени.

На рисунках выше найдены следующие пары: [12,  1], [ 8,  0], [ 2,  3] (после совмещения кластеров из неверного освещения). 

Блок-схемы описанных алгоритмов

Реализация при равномерном освещении

Реализация при затемненном освещении

Оптимизация алгоритма

Как и от любой системы, от алгоритма требуется максимально возможная скорость с минимальными потерями в качестве. Все части алгоритма, для которых существуют написанные библиотеки для работы на GPU, были перенесены на видеокарту. К примеру, для ускорения обработки процесс расчета прямого и обратного преобразования Фурье был перенесен на GPU при помощи библиотеки Cupy. 

Процесс распознавания положений ключевых точек на теле человека занимал много времени даже с учетом инференса на GPU. Нейронная сеть была написана при помощи библиотеки PyTorch, что тоже является неоптимизированным вариантом. Для ускорения вычислений, а также для упрощения интеграции в общую систему, модель была конвертирована в формат ONNX. 

Основная проблема, которая возникла в процессе оптимизации, заключается в бОльшем времени, затрачиваемом на обработку первого кадра. Это связано с неполной инициализацией модели. Для того, чтобы избежать таких задержек, перед запуском используется разогрев модели — инференс на рандомных тензорах нужной размерности. Он занимает порядка 40 секунд, через это время модель полностью готова к использованию, выбросов не наблюдается.

Дальнейшее применение

Алгоритм и его части могут применяться не только для задачи определения команды игрока по цвету его формы. К примеру, при наличии предобученного детектора объектов, такой алгоритм может использоваться для классификации объектов по цветам. В данной ситуации основные правки будут заключаться в функции поиска вектора цвета объекта — нужно более специфично указать область, откуда будут усредняться значения пикселей. Более того, при отсутствии дополнительных моделей,  алгоритм будет работать быстрее, так как модель детекции ключевых точек специфична для задачи классификации игроков по командам.

Алгоритм поиска маски света/тени на кадре также может быть переиспользован. Если объекты достаточно мелкие, но при этом их цвет выделяется на общем фоне, то классические приемы бинаризации будут их выделять. Безусловно, можно применить морфологические операции, которые уберут эти пятна как дефекты маски, но подбор параметров ядра и операций не гарантирует результат в том случае, если размеры объекта на кадре отличаются от тех, для которых были подобраны параметры. Использование преобразования Фурье позволяет размыть мелкие объекты еще до бинаризации, а значит, с учетом адаптивного поиска порога, бинарная маска будет содержать меньше пятен противоположного значения (или не содержать их вовсе).

Спасибо за внимание!