Меня зовут Денис Власов, я Data Scientist в Учи.ру. С помощью
моделей машинного обучения из записей онлайн-уроков мы сделали
гифки последовательность из нескольких кадров с наиболее яркими
эмоциями учеников. Эти гифки получили их родители в
e-mail-рассылке. Вместе с Data Scientist @DariaV Дашей
Васюковой расскажем, как без экспертизы в Computer Vision, а только
с помощью открытых библиотек и готовых моделей сделать MVP, в
основе которого лежат low-res видео. В конце бонус виджет для
быстрой разметки кадров.
Откуда у нас вообще возникла мысль распознавать эмоции? Дело в
том, что мы в Учи.ру развиваем онлайн-школу Учи.Дома сервис
персональных видео-уроков для школьников. Но поскольку такой урок
это чисто человеческое взаимодействие, возникла идея прикрутить к
нему немного аналитики. Такие данные могут помочь повысить
конверсию, отследить эффективность уроков, замерить вовлеченность
учеников и многое другое.
Если у вас, как и у нас, не стоит задача RealTime определения
эмоций, можно пойти простым способом: анализировать записи
уроков.
Маркеры начала и конца урока
Как правило, продолжительность записи с камеры ученика не равна
фактической длине урока. Ученики часто подключаются с опозданием, а
во время урока могут быть дисконнекты и повторные подключения.
Поэтому для начала мы определили, что именно будем считать
уроком.
Для этого мы соотнесли записи с камер учеников и учителей.
Учительские видео помогли обозначить период урока: он начинается,
когда включены обе камеры одновременно, и заканчивается, когда хотя
бы одна камера отключается совсем.
Разбили видео на кадры
Для упрощения анализа нарезали получившиеся отрезки видео
учеников на картинки. Нам хватило одного кадра в секунду: если
ребенок проявил какую-то эмоцию, она будет присутствовать на лице
несколько секунд. Большая степень гранулярности усложнила бы
разметку, но существенно не повлияла на результат.
Научились детектировать детские улыбки (и не только)
На каждом кадре необходимо обнаружить лицо. Если оно там есть,
проверить, родитель это или ученик, а также оценить эмоции на лице.
И тут возникло несколько нюансов, которые пришлось учитывать.
Проблема 1. Распознавать лица на картинках низкого качества
сложнее
Видео пользователей часто бывает низкого качества даже без учета
компрессии видео. Например, ученик может заниматься в темной
комнате, в кадре может быть включенная настольная лампа или люстра,
за спиной ученика может быть яркое окно, лицо может быть в кадре не
полностью.
Стандартный детектор DNN Face Detector из библиотеки OpenCV,
который мы сначала взяли за основу, на наших данных давал неточные
результаты. Оказалось, что алгоритм недостаточно хорошо справляется
с реальными кадрами из видеочатов: иногда пропускает лица, которые
явно есть в кадре, из двух лиц находил только одно или определял
лица там, где их нет.
Стандартный детектор DNN Face Detector
мог определить как лицо узор на занавеске, игрушечного медведя или
даже композицию из картин на стене и стула
Поэтому мы решили попробовать обучить свой детектор. Для этого
взяли реализацию RetinaNet-модели на PyTorch. В качестве данных для
обучения подали результаты работы стандартного детектора и
убедились, что новая модель учится находить лица. Затем подготовили
обучающую и валидационную выборку, просматривая и при необходимости
исправляя результаты работы детектора на новых кадрах: исправлять
разметку работающей модели получается быстрее, чем отмечать лица на
кадре с нуля.
Размечали итеративно: после добавления новой порции размеченных
кадров мы заново обучали модель. А после проверки ее работы
сохраняли разметку для новых кадров, наращивая обучающую выборку.
Всего мы разметили 2624 кадра из 388 видеозаписей, на которых в
сумме было 3325 лиц.
Таким образом удалось обучить более чувствительный в наших
условиях детектор. В валидационной выборке из 140 кадров старый
детектор нашел 150 лиц, а пропустил 38. Новый же пропустил только
5, а 183 обнаружил верно.
Проблема 2. В кадре присутствует не только ребенок
Поскольку на видео-уроках часто присутствуют не только дети, но
и родители, важно научить модель отличать одних от других. В нашем
случае это дает уверенность, что на гифке родитель увидит своего
ребенка, а не себя. Также данные о присутствии родителя на уроке
могут помочь проанализировать продуктовые метрики.
Мы обучили две отдельные модели. На момент эксперимента не было
нужных публичных датасетов, поэтому данные для обучения мы
разметили сами.
Первая модель должна определять, кому принадлежит лицо в кадре:
родителю или ученику. Кажется, что с разметкой не должно было
возникнуть никаких проблем, ведь отличить взрослого от ребенка
просто. Это действительно так, если перед нами целое видео. Но
когда мы имеем дело с отдельными кадрами, то оказывается, что:
-
возраст людей на кадре с низким разрешением становится
неочевидным;
-
дети присутствуют в кадре практически весь урок, а взрослые
минуты.
В процессе разметки мы заметили, что очень часто родитель
присутствует на уроке плечом или локтем. Так мы назвали тип кадров,
когда камера направлена на ученика, но видно, что рядом сидит
родитель. Обычно на таких уроках видно плечо сидящего рядом
взрослого или только локоть.
На всех трех кадрах родитель
присутствует, но по отдельному кадру найти его бывает непросто
Вторая модель должна была находить именно такие родительские
плечи. Очевидно, что в этой задаче детектор лиц не применим,
поэтому надо обучаться на кадрах целиком. Конечно, таких датасетов
мы не нашли в публичном доступе и разметили около 250 000 кадров,
на которых есть часть родителя, и кадры без них. Разметки на
порядок больше, чем в других задачах, потому что размечать гораздо
легче: можно смотреть не отдельные кадры, а отрезки видео и в
несколько кликов отмечать, например, что вот эти 15 минут (900
кадров!) родитель присутствовал.
На дашборде урока с аналитикой доступны графики присутствия
родителями по мнению обеих моделей. Они помогают понять, когда
родитель просто интересуется процессом урока, а когда скорее
общается с преподавателем.
На верхнем графике вероятность
присутствия родителя хотя бы плечом, а на нижнем вероятность того,
что родитель смотрит в камеру, например, общается с преподавателем
Проблема 3. Дети улыбаются по-разному
На практике оказалось, что не так уж просто понять, улыбается
ребенок или нет. И если с улыбчивыми ребятами проблем нет, то
детекция сдержанных улыбок оказывается нетривиальной задачей даже
для человека.
За основу классификатора настроения мы взяли предобученную
модель ResNet34 из библиотеки fast.ai. Эту же библиотеку
использовали для дообучения модели в два этапа: сначала на
публичных датасетах
facial_expressions и
SMILEsmileD с веселыми и нейтральными лицами, потом на
нашем размеченном вручную датасете с кадрами с камер учеников.
Публичные датасеты решили включить, чтобы расширить размер выборки
и помочь модели более качественными изображениями, чем кадры видео
с планшетов и веб-камер наших учеников.
Размечали с помощью кастомного виджета. Все изображения
подвергались одной и той же процедуре предобработки:
-
Масштабирование кадра до размера 64 на 64 пикселя. В публичных
датасетах картинки уже квадратные, поэтому масштабирование не
приводит к искажениям пропорций. В собственном датасете мы сначала
дополняли детектированную область с лицом до квадрата и потом
масштабировали.
-
Приведение к черно-белой палитре. Визуально черно-белые
изображения показались нам чище, кроме того, один из публичных
датасетов уже был в черно-белом формате. Ну и интуитивно кажется,
что для определения улыбки цвета совсем не нужны, что подтвердилось
в экспериментах.
-
Аугментация. Позволяет в разы увеличить эффективный размер
выборки и учесть особенности данных.
-
Нормализация цветов с помощью CLAHE normalizer из библиотеки
OpenCV. По ощущениям, такая нормализация лучше других вытягивает
контраст на пересвеченных или темных изображениях.
Дообучаем модель для распознавания улыбок
1. Аугментации
При дообучении мы использовали достаточно жесткие
аугментации:
-
Отражали изображение по горизонтали.
-
Поворачивали на случайную величину.
-
Применяли три разных искажения для изменения контраста и
яркости.
-
Брали не всю картинку, а квадрат, составляющий не менее 60% от
площади исходного изображения.
-
Обрезали с одной из четырех сторон, вставляя черный
прямоугольник на место обрезанной части.
Первое преобразование нужно исключительно для увеличения размера
выборки. Остальные дополнительно позволяют приблизить публичные
датасеты к нашей задаче. Особенно полезной оказалась последняя
самописная аугментация. Она имитирует ученика, камера которого
смотрит слегка в сторону, и в результате его лицо в кадре
оказывается обрезанным. При детектировании лица и дополнении до
квадрата, обрезанная часть превращается в черную область. Без
аугментаций таких изображений было не достаточно, чтобы модель
научилась понимать, что это, но достаточно, чтобы испортить
качество в среднем. Кроме того, эти ошибки были очевидны для
человека.
Пример аугментаций на одном изображении.
Для наглядности аугментации сделаны до масштабирования к разрешению
64х64Код для аугментаций
# ! pip freeze | grep fastai# fastai==1.0.44import fastaiimport matplotlib.pyplot as pltfrom matplotlib import cmfrom matplotlib import colorsimport seaborn as sns%matplotlib inlinefrom pylab import rcParamsplt.style.use('seaborn-talk')rcParams['figure.figsize'] = 12, 6path = 'facial_expressions/images/'def _side_cutoff( x, cutoff_prob=0.25, cutoff_intensity=(0.1, 0.25)): if np.random.uniform() > cutoff_prob: return x # height and width h, w = x.shape[1:] h_cutoff = np.random.randint( int(cutoff_intensity[0]*h), int(cutoff_intensity[1]*h) ) w_cutoff = np.random.randint( int(cutoff_intensity[0]*w), int(cutoff_intensity[1]*w) ) cutoff_side = np.random.choice( range(4), p=[.34, .34, .16, .16] ) # top, bottom, left, right. if cutoff_side == 0: x[:, :h_cutoff, :] = 0 elif cutoff_side == 1: x[:, h-h_cutoff:, :] = 0 elif cutoff_side == 2: x[:, :, :w_cutoff] = 0 elif cutoff_side == 3: x[:, :, w-w_cutoff:] = 0 return x# side cutoff goes frist.side_cutoff = fastai.vision.TfmPixel(_side_cutoff, order=99)augmentations = fastai.vision.get_transforms( do_flip=True, flip_vert=False, max_rotate=25.0, max_zoom=1.25, max_lighting=0.5, max_warp=0.0, p_affine=0.5, p_lighting=0.5, xtra_tfms = [side_cutoff()])def get_example(): return fastai.vision.open_image( path+'George_W_Bush_0016.jpg', )def plots_f(rows, cols, width, height, **kwargs): [ get_example() .apply_tfms( augmentations[0], **kwargs ).show(ax=ax) for i,ax in enumerate( plt.subplots( rows, cols, figsize=(width,height) )[1].flatten()) ]plots_f(3, 5, 15, 9, size=size)
2. Нормализация цвета
Мы попробовали несколько вариантов предобработки и остановились
на CLAHE нормализации. Этим способом яркость и гамма выравниваются
не по всему изображению, а по его частям. Результат будет
приемлемый, даже если на одном изображении есть и затемненные
участки, и засвеченные.
Пример нормализации цвета на изображениях
из публичного датасетаКод для нормализации цвета
# pip freeze | grep opencv# > opencv-python==4.5.2.52import cv2import matplotlib.pyplot as pltfrom matplotlib import cmfrom matplotlib import colorsimport seaborn as sns%matplotlib inlinefrom pylab import rcParamsplt.style.use('seaborn-talk')rcParams['figure.figsize'] = 12, 6path = 'facial_expressions/images/'imgs = [ 'Guillermo_Coria_0021.jpg', 'Roger_Federer_0012.jpg',]imgs = list( map( lambda x: path+x, imgs ))clahe = cv2.createCLAHE( clipLimit=2.0, tileGridSize=(4, 4))rows_cnt = len(imgs)cols_cnt = 4imsize = 3fig, ax = plt.subplots( rows_cnt, cols_cnt, figsize=(cols_cnt*imsize, rows_cnt*imsize))for row_num, f in enumerate(imgs): img = cv2.imread(f) col_num = 0 img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) ax[row_num, col_num].imshow(img, cmap='gray') ax[row_num, col_num].set_title('bw', fontsize=14) col_num += 1 img_normed = cv2.normalize( img, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F ) ax[row_num, col_num].imshow(img_normed, cmap='gray') ax[row_num, col_num].set_title('bw normalize', fontsize=14) col_num += 1 img_hist_normed = cv2.equalizeHist(img) ax[row_num, col_num].imshow(img_hist_normed, cmap='gray') ax[row_num, col_num].set_title('bw equalizeHist', fontsize=14) col_num += 1 img_clahe = clahe.apply(img) ax[row_num, col_num].imshow(img_clahe, cmap='gray') ax[row_num, col_num].set_title('bw clahe_norm', fontsize=14) col_num += 1 for col in ax[row_num]: col.set_xticks([]) col.set_yticks([])plt.show()
В итоге мы получили модель, способную отличить улыбку от
нейтрального выражения лица в с качеством 0.93 по метрике ROC AUC.
Иными словами, если взять из выборки по случайному кадру с улыбкой
и без, то с вероятностью 93% модель присвоит большую вероятность
улыбки кадру с улыбающимся лицом. Этот показатель мы использовали
для сравнения разных вариантов дообучения и пайплайнов. Но
интуитивно кажется, что это достаточно высокий уровень точности:
даже человек не всегда может определить эмоцию на лице другого
человека. К тому же в реальности существует гораздо больше
выражений лиц кроме однозначной радости и однозначной печали.
Во-первых, улыбка может восприниматься по-разному разными
людьми, размечающими выборку. Например, многое зависит от того,
какие лица видел разметчик ранее: после серии кадров с улыбчивыми
лицами бывает непросто признать улыбающимся хмурого ребенка,
который чуть приподнял уголки губ.
Во-вторых, когда мы размечали данные, в команде тоже возникали
споры: улыбается ребенок на этом кадре или нет. В целом это не
должно мешать модели: легкий шум в разметке помогает бороться с
переобучением.
3. Увеличение объема выборки
На этом этапе у нас было около 5400 размеченных кадров, нам
хотелось понять: достаточно ли такого объема для обучения. Мы
разделили кадры на две подвыборки: половина для оценки качества
(валидация), другая для обучения (трейн). Каждая группа состояла из
кадров с разными учениками: если бы одни и те же лица попали в
разные выборки, результаты оценки качества были бы завышены.
Мы несколько раз дообучили модель на публичных датасетах и
подвыборках разного размера из трейна и проверили качество на
валидационной выборке. На графике видно, что качество растет по
мере увеличения объема подвыборки, поэтому мы разметили
дополнительно еще 3 тыс. кадров: финальную модель обучали на
выборке из 8500 кадров. Скорее всего и при таком объеме данных
качество еще не начало выходить на плато, но перепроверять это мы
не стали.
Такое упражнение называется построением learning curve, и оно в
очередной раз подтверждает тезис: объем данных самое важное в
модели машинного обучения.
Качество на отложенной выборке растет по
мере увеличения выборки для дообучения
4. Картинки Google для обогащения выборки
Мы пробовали спарсить первые 1000 результатов картинок по
запросам в духе happy, unhappy, smiling, neutral и т. д. Не ожидали
получить данные высокого качества, поэтому планировали потом
просмотреть их глазами и удалить совсем неподходящие. В итоге мы
быстро поняли, что никакая фильтрация эти картинки не спасет,
поэтому отказались от этой идеи совсем.
Примеры изображений по запросам happy и unhappy
В итоге мы получили четыре модели, которые с высокой точностью
могли показать:
-
есть ли в кадре лицо;
-
с какой вероятностью этот человек улыбается;
-
ребенок это или взрослый;
-
есть ли в кадре взрослый, даже если мы не нашли лица.
Для этих данных продуктовые аналитики могут придумать множество
способов применения, один из них мы попробовали реализовать.
Собрали гифку
С помощью моделей мы для каждого кадра видео получили
вероятности присутствия родителя или ребенка и вероятности улыбки
на найденных лицах. Из этих кадров мы выбирали по 9 кадров с
улыбками ребенка, которые склеивались в гифку без участия
человека.
Разработчики также настроили автоматическую вставку GIF в письмо
для рассылки. Для этого в шаблоне письма был предусмотрен
дополнительный скрипт, который проверяет, есть ли в базе данных GIF
по конкретному уроку.
Примеры итоговых GIF с улыбками нашей
коллеги и ее детей
Что мы в итоге получили?
Исследование и эксперимент показали, что можно быстро и без
глубокой экспертизы в Computer Vision научиться различать
пользователей и их эмоции по видео (даже если оно плохого качества)
на основе открытых библиотек и моделей.
Возможно, впоследствии мы расширим эту практику, если возникнет
новая идея или обнаружится дополнительная потребность. Но уже
сейчас можно сказать, что этот опыт был интересным и полезным, а
дополнительные данные об эмоциях на уроках уже могут использовать
наши аналитики для построения своих дашбордов и графиков.
Например, можно посмотреть количество отключений, которые
происходили в процессе урока. Подобную информацию можно
использовать, чтобы рекомендовать учителю или ученику более
стабильное подключение.
Статистика дисконнектов. В этом уроке
был единственный дисконнект на стороне ученика
Другой пример трекинг настроения ученика на протяжении урока. Он
позволяет проанализировать ход занятия и понять, нужно ли что-то
менять в его структуре.
Виджеты
Все данные мы размечали сами и делали это довольно быстро
(примерно 100 кадров в минуту). В этом нам помогали самописные
виджеты:
-
Виджет для разметки кадров с улыбками.
-
Виджет для разметки кадров с детьми и взрослыми.
-
Виджет для разметки кадров с плечом или локтем родителя.
Мы хотим поделиться кодом второго виджета как наиболее полного.
Скорее всего, вы не сможете заменить в нем путь к файлам и
использовать для своей задачи, потому что он слишком специфичен. Но
если вам понадобиться написать свой велосипед для разметки, можете
почерпнуть что-то полезное.
Этот виджет показывает таймлайн всего урока с кадрами на равных
расстояниях. По этим кадрам можно ориентироваться, чтобы находить
нужные промежутки видео и отправить лица с этих кадров в
разметку.
Выбрав промежуток видео, где присутствует только ученик или
только родитель, можно размечать лица целыми десятками в один клик.
Когда в кадре присутствуют одновременно взрослый и ребенок, в
разметке помогает обученная модель. Она сортирует показанные лица
по взрослости и назначает предварительные метки. Остается только
исправить ошибки модели в неочевидных случаях.
Видео работы виджета
Таким образом можно быстро набирать размеченные данные,
одновременно отслеживая, какие лица вызывают у модели затруднения.
Например, наличие очков модель поначалу считала явным признаком
взрослого человека. Пришлось отдельно искать кадры с детьми в
очках, чтобы исправить это заблуждение.
Код виджета
import pandas as pdimport numpy as npimport datetimeimport randomimport osimport ipywidgets as widgetsfrom IPython.display import displayfrom pathlib import Pathclass BulkLabeler(): def __init__(self, frames_path, annotations_path, labels = ['0', '1'], predict_fn = None, frame_width=120, num_frames = 27, face_width = 120, num_faces = 27, myname = '?', ): self.predict_fn = predict_fn self.labels = labels self.frames_path = frames_path self.frame_width = frame_width self.num_frames = num_frames self.face_width = face_width self.num_faces = num_faces self.myname = myname self.faces_batch = [] # get annotations self.annotations_path = annotations_path processed_videos = [] if annotations_path.exists(): annotations = pd.read_csv(annotations_path) processed_videos = annotations.file.str.split('/').str[-3].unique() else: with open(self.annotations_path, 'w') as f: f.write('file,label,by,created_at\n') # get list of videos self.video_ids = [x for x in os.listdir(frames_path) if x not in processed_videos] random.shuffle(self.video_ids) self.video_ind = -1 self._make_video_widgets_row() self._make_frames_row() self._make_range_slider() self._make_buttons_row() self._make_faces_row() self._make_video_stats_row() display(widgets.VBox([self.w_video_row, self.w_frames_row, self.w_slider_row, self.w_buttons_row, self.w_faces_row, self.w_faces_label, self.w_video_stats])) self._on_next_video_click(0) ### Video name and next video button def _make_video_widgets_row(self): # widgets for current video name and "Next video" button self.w_current_video = widgets.Text( value='', description='Current video:', disabled=False, layout = widgets.Layout(width='500px') ) self.w_next_video_button = widgets.Button( description='Next video', button_style='info', # 'success', 'info', 'warning', 'danger' or '' tooltip='Go to the next video', icon='right-arrow' ) self.w_video_row = widgets.HBox([self.w_current_video, self.w_next_video_button]) self.w_current_video.observe(self._on_video_change, names='value') self.w_next_video_button.on_click(self._on_next_video_click) def _on_next_video_click(self, _): while True: self.video_ind += 1 current_video = self.video_ids[self.video_ind] if next(os.scandir(self.frames_path/current_video/'student_faces'), None) is not None: break self.w_current_video.value = current_video def _on_video_change(self, change): self.video_id = change['new'] self.frame_nums_all = sorted(int(f.replace('.jpg','')) for f in os.listdir(self.frames_path/self.video_id/'student_src')) start, stop = min(self.frame_nums_all), max(self.frame_nums_all) self.w_range_slider.min = start self.w_range_slider.max = stop step = self.frame_nums_all[1] - self.frame_nums_all[0] if len(self.frame_nums_all)>1 else 1 self.w_range_start.step = step self.w_range_stop.step = step # change to slider value will cause frames to be redrawn self.w_range_slider.value = [start, stop] # reset faces self.faces_df = None self._reset_faces_row() self.w_video_stats.value = f'Video {self.video_id} no annotations yet.' def _close_video_widgets_row(self): self.w_current_video.close() self.w_next_video_button.close() self.w_video_row.close() ### Video frames box def _make_frames_row(self): frame_boxes = [] self.w_back_buttons = {} self.w_forward_buttons = {} for i in range(self.num_frames): back_button = widgets.Button(description='<',layout=widgets.Layout(width='20px',height='20px')) self.w_back_buttons[back_button] = i back_button.on_click(self._on_frames_back_click) label = widgets.Label(str(i+1), layout = widgets.Layout(width=f'{self.frame_width-50}px')) forward_button = widgets.Button(description='>',layout=widgets.Layout(width='20px',height='20px')) self.w_forward_buttons[forward_button] = i forward_button.on_click(self._on_frames_forward_click) image = widgets.Image(width=f'{self.frame_width}px') frame_boxes.append(widgets.VBox([widgets.HBox([back_button, label, forward_button]), image])) self.w_frames_row = widgets.GridBox(frame_boxes, layout = widgets.Layout(width='100%', display='flex', flex_flow='row wrap')) def _on_frames_back_click(self, button): frame_ind = self.w_back_buttons[button] frame = int(self.w_frames_row.children[frame_ind].children[0].children[1].value) start, stop = self.w_range_slider.value self.w_range_slider.value = [frame, stop] def _on_frames_forward_click(self, button): frame_ind = self.w_forward_buttons[button] frame = int(self.w_frames_row.children[frame_ind].children[0].children[1].value) start, stop = self.w_range_slider.value self.w_range_slider.value = [start, frame] def _close_frames_row(self): for box in self.w_frames_row.children: label_row, image = box.children back, label, forward = label_row.children image.close() back.close() label.close() forward.close() box.close() self.w_frames_row.close() ### Frames range slider def _make_range_slider(self): self.w_range_start = widgets.BoundedIntText( value=0, min=0, max=30000, step=1, description='Frames from:', disabled=False, layout = widgets.Layout(width='240px') ) self.w_range_stop = widgets.BoundedIntText( value=30000, min=0, max=30000, step=1, description='to:', disabled=False, layout = widgets.Layout(width='240px') ) self.w_range_slider = widgets.IntRangeSlider( value=[0, 30000], min=0, max=30000, step=1, description='', disabled=False, continuous_update=False, orientation='horizontal', readout=True, readout_format='d', layout=widgets.Layout(width='500px') ) self.w_range_flip = widgets.Button(description='Flip range', button_style='', # 'success', 'info', 'warning', 'danger' or '' tooltip='Invert frames selection', layout = widgets.Layout(width=f'{self.frame_width}px'), icon='retweet' ) self.w_range_slider.observe(self._on_slider_change, names='value') self.w_range_start.observe(self._on_range_start_change, names='value') self.w_range_stop.observe(self._on_range_stop_change, names='value') self.w_range_flip.on_click(self._on_range_flip) self.w_slider_row = widgets.HBox([self.w_range_start, self.w_range_slider, self.w_range_stop, self.w_range_flip]) def _close_range_slider(self): self.w_range_start.close() self.w_range_stop.close() self.w_range_slider.close() self.w_range_flip.close() self.w_slider_row.close() def _on_range_flip(self, _): start, stop = self.w_range_slider.value left, right = self.w_range_slider.min, self.w_range_slider.max if start==left and right==stop: pass elif start - left > right - stop: self.w_range_slider.value=[left, start] else: self.w_range_slider.value=[stop, right] def _on_range_start_change(self, change): new_start = change['new'] start, stop = self.w_range_slider.value self.w_range_slider.value = [new_start, stop] def _on_range_stop_change(self, change): new_stop = change['new'] start, stop = self.w_range_slider.value self.w_range_slider.value = [start, new_stop] def _on_slider_change(self, change): start, stop = change['new'] # update text controls self.w_range_start.max = stop self.w_range_start.value = start self.w_range_stop.min = start self.w_range_stop.max = self.w_range_slider.max self.w_range_stop.value = stop # show frames that fit current selection frame_nums = [i for i in self.frame_nums_all if i>=start and i<=stop] N = len(frame_nums) n = self.num_frames inds = [int(((N-1)/(n-1))*i) for i in range(n)] # load new images into image widgets for ind, box in zip(inds, self.w_frames_row.children): frame_num = frame_nums[ind] filename = self.frames_path/self.video_id/'student_src'/f'{frame_num}.jpg' with open(filename, "rb") as image: f = image.read() label, image = box.children label.children[1].value = str(frame_num) image.value = f ### Buttons row def _make_buttons_row(self): labels = list(self.labels) if self.predict_fn is not None: labels.append('model') self.w_default_label = widgets.ToggleButtons(options=labels, value=self.labels[0], description='Default label:') self.w_next_batch_button = widgets.Button(description='New batch', button_style='info', # 'success', 'info', 'warning', 'danger' or '' tooltip='Show next batch of faces from current frame range', icon='arrow-right' ) self.w_save_button = widgets.Button(description='Save labels', button_style='success', # 'success', 'info', 'warning', 'danger' or '' tooltip='Save current labels', icon='check' ) self.w_buttons_row = widgets.HBox([self.w_default_label, self.w_next_batch_button, self.w_save_button]) self.w_next_batch_button.on_click(self._on_next_batch_click) self.w_save_button.on_click(self._on_save_labels_click) def _close_buttons_row(self): self.w_default_label.close() self.w_next_batch_button.close() self.w_save_button.close() self.w_buttons_row.close() def _on_next_batch_click(self, _): if self.faces_df is None: self._create_faces_df() # select a sample from faces_df start, stop = self.w_range_slider.value subdf = self.faces_df.loc[lambda df: df.frame_num.ge(start)& df.frame_num.le(stop)& df.label.eq('')] num_faces = min(len(subdf), self.num_faces) if num_faces == 0: self.faces_batch = [] self.w_faces_label.value = 'No more unlabeled images in this frames range' self.w_faces_label.layout.visibility = 'visible' for box in self.w_faces_row.children: box.layout.visibility = 'hidden' else: self.w_faces_label.layout.visibility = 'hidden' self.faces_batch = subdf.sample(num_faces).index # if we have a model then we use it to sort images if self.predict_fn is not None: probs, labels = self._predict() # sort faces according to probability ind = sorted(range(len(probs)), key=probs.__getitem__) self.faces_batch = [self.faces_batch[i] for i in ind] labels = [labels[i] for i in ind] # create labels for each face if self.w_default_label.value != 'model': labels = [self.w_default_label.value]*len(self.faces_batch) # update faces UI for facefile, label, box in zip(self.faces_batch, labels, self.w_faces_row.children): image, buttons = box.children with open(self.frames_path/facefile, "rb") as im: image.value = im.read() buttons.value = label box.layout.visibility = 'visible' if len(self.faces_batch) < len(self.w_faces_row.children): for box in self.w_faces_row.children[len(self.faces_batch):]: box.layout.visibility = 'hidden' def _predict(self): probs = [] labels = [] for facefile in self.faces_batch: prob, label = self.predict_fn(self.frames_path/facefile) probs.append(prob) labels.append(label) self.faces_df.loc[self.faces_batch, 'prob'] = probs return probs, labels def _on_save_labels_click(self, _): self.w_save_button.description='Saving...' with open(self.annotations_path, 'a') as f: for file, box in zip(self.faces_batch, self.w_faces_row.children): label = box.children[1].value self.faces_df.at[file,'label'] = label print(file, label, self.myname, str(datetime.datetime.now()),sep=',', file=f) # update current video statistics stats = self.faces_df.loc[self.faces_df.label.ne(''),'label'].value_counts().sort_index() stats_str = ', '.join(f'{label}: {count}' for label, count in stats.items()) self.w_video_stats.value = f'Video {self.video_id} {stats_str}.' self.w_save_button.description = 'Save labels' # ask for next batch self._on_next_batch_click(0) ### Faces row def _make_faces_row(self): face_boxes = [] for i in range(self.num_faces): image = widgets.Image(width=f'{self.face_width}px') n = len(self.labels) toggle_buttons_width = int(((self.face_width-5*(n-1))/n)) toggle_buttons = widgets.ToggleButtons(options=self.labels, value=self.w_default_label.value, style=widgets.ToggleButtonsStyle(button_width=f'{toggle_buttons_width}px')) face_boxes.append(widgets.VBox([image, toggle_buttons])) self.w_faces_row = widgets.GridBox(face_boxes, layout = widgets.Layout(width='100%', display='flex', flex_flow='row wrap')) self.w_faces_label = widgets.Label() self._reset_faces_row() def _close_faces_row(self): for box in self.w_faces_row.children: image, buttons = box.children for w in [image, buttons, box]: w.close() self.w_faces_row.close() self.w_faces_label.close() def _reset_faces_row(self): for box in self.w_faces_row.children: box.layout.visibility = 'hidden' self.w_faces_label.layout.visibility = 'visible' self.w_faces_label.value = 'Press "New batch" button to see a new batch of faces' self.faces_batch = [] ### Video statistics row def _make_video_stats_row(self): self.w_video_stats = widgets.Label('No video currently selected') def _close_video_stats_row(self): self.w_video_stats.close() def _create_faces_df(self): folder = Path(self.video_id,'student_faces') df = pd.DataFrame({'file':[folder/f for f in os.listdir(self.frames_path/folder)]}) df['frame_num'] = df.file.apply(lambda x: int(x.stem.split('_')[0])) df['label'] = '' #TODO maybe existing annotations? df['prob'] = np.nan df = df.sort_values(by='frame_num').set_index('file') self.faces_df = df def close(self): self._close_video_widgets_row() self._close_frames_row() self._close_range_slider() self._close_buttons_row() self._close_faces_row() self._close_video_stats_row()