Как стать автором
Обновить
0
Учи.ру team
интерактивная образовательная платформа

Как мы построили Computer Vision из подручных материалов, чтобы сделать гифки

Время на прочтение 19 мин
Количество просмотров 3.2K

Меня зовут Денис Власов, я Data Scientist в Учи.ру. С помощью моделей машинного обучения из записей онлайн-уроков мы сделали гифки — последовательность из нескольких кадров с наиболее яркими эмоциями учеников. Эти гифки получили их родители в e-mail-рассылке. Вместе с Data Scientist @DariaV Дашей Васюковой расскажем, как без экспертизы в Computer Vision, а только с помощью открытых библиотек и готовых моделей сделать MVP, в основе которого лежат low-res видео. В конце бонус — виджет для быстрой разметки кадров.

Откуда у нас вообще возникла мысль распознавать эмоции? Дело в том, что мы в Учи.ру развиваем онлайн-школу Учи.Дома — сервис персональных видео-уроков для школьников. Но поскольку такой урок — это чисто человеческое взаимодействие, возникла идея «‎прикрутить»‎ к нему немного аналитики. Такие данные могут помочь повысить конверсию, отследить эффективность уроков, замерить вовлеченность учеников и многое другое.

Если у вас, как и у нас, не стоит задача RealTime определения эмоций, можно пойти простым способом: анализировать записи уроков.

Маркеры начала и конца урока

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

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

Разбили видео на кадры

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

Научились детектировать детские улыбки (и не только)

На каждом кадре необходимо обнаружить лицо. Если оно там есть, проверить, родитель это или ученик, а также оценить эмоции на лице. И тут возникло несколько нюансов, которые пришлось учитывать.

Проблема 1. Распознавать лица на картинках низкого качества сложнее

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

Стандартный детектор DNN Face Detector из библиотеки OpenCV, который мы сначала взяли за основу, на наших данных давал неточные результаты. Оказалось, что алгоритм недостаточно хорошо справляется с реальными кадрами из видеочатов: иногда пропускает лица, которые явно есть в кадре, из двух лиц находил только одно или определял лица там, где их нет.

Стандартный детектор DNN Face Detector мог определить как лицо узор на занавеске, игрушечного медведя или даже композицию из картин на стене и стула
Стандартный детектор 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 с веселыми и нейтральными лицами, потом на нашем размеченном вручную датасете с кадрами с камер учеников. Публичные датасеты решили включить, чтобы расширить размер выборки и помочь модели более качественными изображениями, чем кадры видео с планшетов и веб-камер наших учеников.

Размечали с помощью кастомного виджета. Все изображения подвергались одной и той же процедуре предобработки:

  1. Масштабирование кадра до размера 64 на 64 пикселя. В публичных датасетах картинки уже квадратные, поэтому масштабирование не приводит к искажениям пропорций. В собственном датасете мы сначала дополняли детектированную область с лицом до квадрата и потом масштабировали.

  2. Приведение к черно-белой палитре. Визуально черно-белые изображения показались нам «‎чище»‎, кроме того, один из публичных датасетов уже был в черно-белом формате. Ну и интуитивно кажется, что для определения улыбки цвета совсем не нужны, что подтвердилось в экспериментах.

  3. Аугментация. Позволяет в разы увеличить эффективный размер выборки и учесть особенности данных.

  4. Нормализация цветов с помощью CLAHE normalizer из библиотеки OpenCV. По ощущениям, такая нормализация лучше других вытягивает контраст на пересвеченных или темных изображениях.

Дообучаем модель для распознавания улыбок

1. Аугментации

При дообучении мы использовали достаточно жесткие аугментации:

  • Отражали изображение по горизонтали.

  • Поворачивали на случайную величину.

  • Применяли три разных искажения для изменения контраста и яркости.

  • Брали не всю картинку, а квадрат, составляющий не менее 60% от площади исходного изображения.

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

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

Пример аугментаций на одном изображении. Для наглядности аугментации сделаны до масштабирования к разрешению 64х64
Пример аугментаций на одном изображении. Для наглядности аугментации сделаны до масштабирования к разрешению 64х64
Код для аугментаций
# ! pip freeze | grep fastai
# fastai==1.0.44

import fastai

import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib import colors
import seaborn as sns
%matplotlib inline
from pylab import rcParams
plt.style.use('seaborn-talk')
rcParams['figure.figsize'] = 12, 6

path = '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.52

import cv2

import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib import colors
import seaborn as sns
%matplotlib inline
from pylab import rcParams
plt.style.use('seaborn-talk')
rcParams['figure.figsize'] = 12, 6


path = '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 = 4
imsize = 3
fig, 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
Примеры изображений по запросам happy и unhappy

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

  • есть ли в кадре лицо;

  • с какой вероятностью этот человек улыбается;

  • ребенок это или взрослый;

  • есть ли в кадре взрослый, даже если мы не нашли лица.

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

Собрали гифку

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

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

Примеры итоговых GIF с улыбками нашей коллеги и ее детей
Примеры итоговых GIF с улыбками нашей коллеги и ее детей

Что мы в итоге получили?

Исследование и эксперимент показали, что можно быстро и без глубокой экспертизы в Computer Vision научиться различать пользователей и их эмоции по видео (даже если оно плохого качества) на основе открытых библиотек и моделей. 

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

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

Статистика дисконнектов. В этом уроке был единственный дисконнект на стороне ученика
Статистика дисконнектов. В этом уроке был единственный дисконнект на стороне ученика

Другой пример — трекинг настроения ученика на протяжении урока. Он позволяет проанализировать ход занятия и понять, нужно ли что-то менять в его структуре.

Виджеты

Все данные мы размечали сами и делали это довольно быстро (примерно 100 кадров в минуту). В этом нам помогали самописные виджеты:

  1. Виджет для разметки кадров с улыбками.

  2. Виджет для разметки кадров с детьми и взрослыми.

  3. Виджет для разметки кадров с «‎плечом»‎ или «‎локтем»‎ родителя.

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

Этот виджет показывает таймлайн всего урока с кадрами на равных расстояниях. По этим кадрам можно ориентироваться, чтобы находить нужные промежутки видео и отправить лица с этих кадров в разметку.

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

Видео работы виджета

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

Код виджета
import pandas as pd
import numpy as np
import datetime
import random
import os
import ipywidgets as widgets
from IPython.display import display
from pathlib import Path


class 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()
Теги:
Хабы:
+3
Комментарии 4
Комментарии Комментарии 4

Публикации

Информация

Сайт
uchi.ru
Дата регистрации
Дата основания
Численность
501–1 000 человек
Местоположение
Россия

Истории