bubble detector using YOLOv4

Note : It's not the final version code. I will the refine and update the code.

Overview

Models detection speech bubble in webtoons or cartoons. I have referenced and implemented pytorch-YOLOv4 to detect speech bubble. The key point for improving performance is data analysis. In the case of speech bubbles, there are various forms. Therefore, I define the form of speech bubbles and present the results of training by considering the distribution of data.

Definition of Speech Bubble

Various speech bubble forms of real webtoons

• In fact, there are various colors and various shapes of speech bubbles in webtoons.

New Definition

Key standard for Data Definition: Shape, Color, Form

standard

• shape : Ellipse(tawon), Thorn(gasi), Sea_urchin(seonggye), Rectangle(sagak), Cloud(gurm)

• Color : Black/white(bw), Colorful(color), Transparency(tran), Gradation

• Form : Basic, Double Speech bubble, Multi-External, Scatter-type

• example image

• In this project, two categories are applied, shape and color, and form and Gradation are classified as ect.
  

classes

This class is not about detection, but about speech bubble data distribution.

Install dependencies

• Pytorch Version

  • Pytorch 1.4.0 for TensorRT 7.0 and higher
  • Pytorch 1.5.0 and 1.6.0 for TensorRT 7.1.2 and higher

• Install Dependencies Code

  pip install onnxruntime numpy torch tensorboardX scikit_image tqdm easydict Pillow skimage opencv_python pycocotools

  or

  pip install -r requirements.txt

  
  

Pretrained model

Train

• 1. Download weight

• 2. Train

  python train.py -g gpu_id -classes number of classes  -dir 'data_dir' -pretrained 'pretrained_model.pth'

  or

  Train.sh 

• 3. Config setting

  • cfg.py

    • class = 1
    • learning_rate = 0.001
    • max_batches = 2000 (class * 2000)
    • steps = [1600, 1800], (max_batches 0.8 , max_batches 0.9)
    • train_dir = your dataset root
      • root tree
        
        The image folder contains .jpg or .png image files. The XML folder contains .XML files(label).

  • cfg/yolov4.cfg

    • class 1
    • filter 18 (4 + 1 + class) * 3 (line: 961, 1049, 1137)

If you want to train custom dataset, use the information above.

Demo

• 1. Download weight
• 2. Demo

  python demp.py -cfgfile cfgfile -weightfile pretrained_model.pth -imgfile image_dir 

  • defualt cfgfile is ./cfg/yolov4.cfg

Metric

• 1. validation dataset

• The above distribution is based on speech bubbles, not cuts.
• The distribution is not constant because there are a number of speech bubbles inside a single cut. In addition, for some classes, examples are difficult to find, resulting in an unbalanced distribution as shown above.