Обнаружение логотипа в прямых трансляциях с использованием YOLOv4

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

Мы приступим к делу. Давайте поговорим о данных (которые были синтетическими):

Я сделал две вещи:

Получите список SVG-файлов интересующих вас логотипов.
Получите много фоновых изображений.
Автоматически аннотировать логотипы ограничивающими рамками (поскольку это проблема обнаружения объектов)

Первая часть не требует особых пояснений. Позвольте мне объяснить, как идти о 2-й части.

Используя Twitch API и оболочку Python FFmpeg, я сделал N скриншотов из M прямых трансляций Twitch. Нам нужны эти кадры, чтобы иметь возможность вставлять наши логотипы и создавать набор данных, поэтому эта часть очень важна.

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

(К сожалению, я не могу поделиться каким-либо кодом для вышеуказанного процесса из-за NDA)

Как только это было сделано, я приступил к обучению и выводам на Colab. Следующий код в основном предоставлен Roboflow с небольшими изменениями:

Настройка cuDNN в Colab для YOLOv4

# Change the number depending on what GPU is listed above, under NVIDIA-SMI > Name.
# Tesla K80: 30
# Tesla P100: 60
# Tesla T4: 75\
%env compute_capability=60

Установка Darknet для YOLOv4 на Colab

%cd /content/
%rm -rf darknet sample_data/

#we clone the fork of darknet maintained by roboflow
#small changes have been made to configure darknet for training
!git clone https://github.com/AlexeyAB/darknet.git

#install environment from the Makefile
%cd darknet/
# compute_30, sm_30 for Tesla K80
# compute_75, sm_75 for Tesla T4
# !sed -i 's/ARCH= -gencode arch=compute_60,code=sm_60/ARCH= -gencode arch=compute_30,code=sm_30/g' Makefile
 
#install environment from the Makefile
#note if you are on Colab Pro this works on a P100 GPU
#if you are on Colab free, you may need to change the Makefile for the K80 GPU
#this goes for any GPU, you need to change the Makefile to inform darknet which GPU you are running on.
!sed -i 's/OPENCV=0/OPENCV=1/g' Makefile
 
#for GPU accelerated training
!sed -i 's/GPU=0/GPU=1/g' Makefile
!sed -i 's/CUDNN=0/CUDNN=1/g' Makefile
!sed -i "s/ARCH= -gencode arch=compute_75,code=sm_75/ARCH= -gencode arch=compute_${compute_capability},code=sm_${compute_capability}/g" Makefile
!make

%cd /content/darknet

# YOLOv4weights
!wget https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.weights
!wget https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.conv.137

Настройте пользовательский набор данных для YOLOv4

Я загрузил все обучающие и тестовые изображения на свой Google Диск. Ячейки ниже собирают набор данных для совместимости с форматом YOLO.

# Mount Drive to Colab
from google.colab import drive
drive.mount('/content/drive')
%cd /content
!mkdir /content/darknet/prediction-images
!mkdir /content/darknet/data/obj

!unzip /content/drive/MyDrive/worlds-lec-dataset.zip -d /content/darknet/logorec/
!unzip /content/drive/MyDrive/youtube_images.zip -d /content/darknet/logorec/

!cp /content/drive/MyDrive/worlds-test-video.mp4 /content/darknet/
!cp /content/drive/MyDrive/lec-test-video.mp4 /content/darknet

%cat /content/darknet/logorec/Worlds-LEC-Dataset/classes.txt | cut -d" " -f2 | cut -d"." -f1 >> /content/darknet/data/obj.names

%cd /content/
!rm -rf /content/weights
# Keeping all imports in a single cell is a good idea so that, 
# if you need to update your imports and add more, you can just
# update one cell without having to re-run other code.

import os
from random import sample
trainshare = 0.8

#Set up training file directories for custom dataset
%cd /content/darknet/

#copy image and labels
%cp /content/darknet/logorec/Worlds-LEC-Dataset/*.png data/obj/
%cp /content/darknet/logorec/Worlds-LEC-Dataset/*.txt data/obj/

with open('data/obj.data', 'w') as out:
  out.write(f'classes = 13\n')
  out.write('train = data/train.txt\n')
  out.write('valid = data/valid.txt\n')
  out.write('names = data/obj.names\n')
  out.write('backup = backup/')
 
#write train file (just the image list)
files = [f for f in os.listdir('data/obj/') if f.endswith('png')]
train = sample(files, int(len(files) * trainshare))
validation = set(files) - set(train)

with open('data/train.txt', 'w') as out:
  for img in train:
    out.write('data/obj/' + img + '\n')
 
#write the valid file (just the image list)
with open('data/valid.txt', 'w') as out:
  for img in validation:
    out.write('data/obj/' + img + '\n')

Напишите пользовательскую конфигурацию обучения для YOLOv4

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

#we build config dynamically based on number of classes
#we build iteratively from base config files. This is the same file shape as cfg/yolo-obj.cfg
def file_len(fname):
  with open(fname) as f:
    for i, l in enumerate(f):
      pass
  return i + 1

num_classes = file_len('data/obj.names')
max_batches = num_classes*2000
steps1 = .8 * max_batches
steps2 = .9 * max_batches
steps_str = str(steps1)+','+str(steps2)
num_filters = (num_classes + 5) * 3
 
print("writing config for a custom YOLOv4 detector detecting number of classes: " + str(num_classes))
 
#Instructions from the darknet repo
#change line max_batches to (classes*2000 but not less than number of training images, and not less than 6000), f.e. max_batches=6000 if you train for 3 classes
#change line steps to 80% and 90% of max_batches, f.e. steps=4800,5400
if os.path.exists('./cfg/custom-yolov4-detector.cfg'): 
  os.remove('./cfg/custom-yolov4-detector.cfg')
 
#customize iPython writefile so we can write variables
from IPython.core.magic import register_line_cell_magic
 
@register_line_cell_magic
def writetemplate(line, cell):
    with open(line, 'w') as f:
        f.write(cell.format(**globals()))

Обучить пользовательский детектор YOLOv4

Обучение заключается в вызове файла обучения.

!./darknet detector train data/obj.data cfg/custom-yolov4-detector.cfg yolov4.conv.137 -dont_show -map -clear

Выполнение вывода на тестовых видео и изображениях с сохраненными весами YOLOv4

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

%cd /content

!zip -r /content/drive/MyDrive/worlds-lec-dataset.zip /content/darknet/logorec/Worlds-LEC-Dataset
!zip -r /content/drive/MyDrive/youtube_images.zip /content/darknet/logorec/youtube_images

!cp /content/darknet/worlds-test-video.mp4 /content/drive/MyDrive/ 
!cp /content/darknet/lec-test-video.mp4 /content/drive/MyDrive/ 
#define utility function
def imShow(path):
  import cv2
  import matplotlib.pyplot as plt
  %matplotlib inline
 
  image = cv2.imread(path)
  height, width = image.shape[:2]
  resized_image = cv2.resize(image,(3*width, 3*height), interpolation = cv2.INTER_CUBIC)
 
  fig = plt.gcf()
  fig.set_size_inches(18, 10)
  plt.axis("off")
  plt.imshow(cv2.cvtColor(resized_image, cv2.COLOR_BGR2RGB))
  plt.show()

#check if weigths have saved yet
#backup houses the last weights for our detector
#(file yolo-obj_last.weights will be saved to the build\darknet\x64\backup\ for each 100 iterations)
#(file yolo-obj_xxxx.weights will be saved to the build\darknet\x64\backup\ for each 1000 iterations)
#After training is complete - get result yolo-obj_final.weights from path build\darknet\x64\bac
!ls /content/darknet/backup
#if it is empty you haven't trained for long enough yet, you need to train for at least 100 iterations

#coco.names is hardcoded somewhere in the detector
%cd /content/darknet/
%cp data/obj.names data/coco.names

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

import random, os, shutil

#/test has images that we can test our detector on
test_folder = "/content/darknet/logorec/youtube_images/"
test_images = [f for f in os.listdir(test_folder) if f.endswith('.png')]
img_path = ''
counter = 0

for i in range(0, len(test_images)):
  img_path = test_folder + test_images[i];
  !./darknet detect cfg/custom-yolov4-detector.cfg backup/custom-yolov4-detector_best.weights {img_path} -dont-show
  imShow('predictions.jpg')
  src_dir='/content/darknet/predictions.jpg'
  dst_dir='/content/darknet/prediction-images/prediction-'+str(counter)+".jpg"
  shutil.copy(src_dir,dst_dir)
  counter += 1

%cd /content/darknet/

Сохранение весов, файлов конфигурации, диаграмм и прогнозных изображений на диск

!mkdir /content/darknet/training_results

# Creating the readme file
!touch /content/darknet/training_results/readme.txt
!echo "The training set was trained on YOLOv4's default weights with a distractor set" >> /content/darknet/training_results/readme.txt

# Moving the chart to the results directory
!cp /content/darknet/chart* /content/darknet/training_results

# Moving the config file to the results directory
!cp /content/darknet/cfg/custom-yolov4-detector.cfg /content/darknet/training_results

# Maving the obj.names and obj.data files to the results directory
!cp /content/darknet/data/obj.* /content/darknet/training_results

# Moving the weights to the results directory
!cp /content/darknet/backup/*.weights /content/darknet/training_results

# Moving the predicted images to the results directory  
!cp -r /content/darknet/prediction-images /content/darknet/training_results

# Zip the training result files and copy to the drive
!zip -r march13_results.zip /content/darknet/training_results
!cp march13_results.zip /content/drive/MyDrive/
!./darknet detector demo data/obj.data cfg/custom-yolov4-tiny-detector.cfg backup/custom-yolov4-tiny-detector_best.weights /content/darknet/worlds-test-video.mp4 -dont_show -thresh 0.25 -out_filename /content/RESULT_worlds-test.mp4
!cp /content/RESULT_LEC-test.mp4 /content/drive/MyDrive/

!./darknet detector demo data/obj.data cfg/custom-yolov4-tiny-detector.cfg backup/custom-yolov4-tiny-detector_best.weights /content/darknet/lec-test-video.mp4 -dont_show -thresh 0.25 -out_filename /content/RESULT_LEC-test.mp4
!cp /content/RESULT_worlds-test.mp4 /content/drive/MyDrive/

By:

Онур Андрос Озбек

Монреаль, Квебек

Обнаружение логотипа в прямых трансляциях с использованием YOLOv4 — Онур Озбек