Merge pull request JaidedAI#739 from gmuffiness/master

add CRAFT training code

trainer/craft/.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+model/__pycache__/
+wandb/*
+vis_result/*

trainer/craft/README.md

@@ -0,0 +1,105 @@
+# CRAFT-train
+On the official CRAFT github, there are many people who want to train CRAFT models. 
+
+However, the training code is not published in the official CRAFT repository. 
+
+There are other reproduced codes, but there is a gap between their performance and performance reported in the original paper. (https://arxiv.org/pdf/1904.01941.pdf) 
+
+The trained model with this code recorded a level of performance similar to that of the original paper.
+
+```bash
+├── config
+│ ├── syn_train.yaml
+│ └── custom_data_train.yaml
+├── data
+│ ├── pseudo_label
+│ │ ├── make_charbox.py
+│ │ └── watershed.py
+│ ├── boxEnlarge.py
+│ ├── dataset.py
+│ ├── gaussian.py
+│ ├── imgaug.py
+│ └── imgproc.py
+├── loss
+│ └── mseloss.py
+├── metrics
+│ └── eval_det_iou.py
+├── model
+│ ├── craft.py
+│ └── vgg16_bn.py
+├── utils
+│ ├── craft_utils.py
+│ ├── inference_boxes.py
+│ └── utils.py
+├── trainSynth.py
+├── train.py
+├── train_distributed.py
+├── eval.py
+├── data_root_dir (place dataset folder here)
+└── exp (model and experiment result files will saved here)
+```
+
+### Installation
+
+Install using `pip`
+
+``` bash
+pip install -r requirements.txt
+```
+
+
+### Training
+1. Put your training, test data in the following format
+ ```
+ └── data_root_dir (you can change root dir in yaml file)
+ ├── ch4_training_images
+ │ ├── img_1.jpg
+ │ └── img_2.jpg
+ ├── ch4_training_localization_transcription_gt
+ │ ├── gt_img_1.txt
+ │ └── gt_img_2.txt
+ ├── ch4_test_images
+ │ ├── img_1.jpg
+ │ └── img_2.jpg
+ └── ch4_training_localization_transcription_gt
+ ├── gt_img_1.txt
+ └── gt_img_2.txt
+ ```
+ * localization_transcription_gt files format :
+ ```
+ 377,117,463,117,465,130,378,130,Genaxis Theatre
+ 493,115,519,115,519,131,493,131,[06]
+ 374,155,409,155,409,170,374,170,###
+ ```
+2. Write configuration in yaml format (example config files are provided in `config` folder.)
+ * To speed up training time with multi-gpu, set num_worker > 0 
+3. Put the yaml file in the config folder
+4. Run training script like below (If you have multi-gpu, run train_distributed.py)
+5. Then, experiment results will be saved to ```./exp/[yaml]``` by default.
+
+* Step 1 : To train CRAFT with SynthText dataset from scratch
+ * Note : This step is not necessary if you use <a href="https://drive.google.com/file/d/1enVIsgNvBf3YiRsVkxodspOn55PIK-LJ/view?usp=sharing">this pretrain</a> as a checkpoint when start training step 2. You can download and put it in `exp/CRAFT_clr_amp_29500.pth` and change `ckpt_path` in the config file according to your local setup.
+ ```
+ CUDA_VISIBLE_DEVICES=0 python3 trainSynth.py --yaml=syn_train
+ ```
+
+* Step 2 : To train CRAFT with [SynthText + IC15] or custom dataset
+ ```
+ CUDA_VISIBLE_DEVICES=0 python3 train.py --yaml=custom_data_train ## if you run on single GPU
+ CUDA_VISIBLE_DEVICES=0,1 python3 train_distributed.py --yaml=custom_data_train ## if you run on multi GPU
+ ```
+
+### Arguments
+* ```--yaml``` : configuration file name
+
+### Evaluation
+* In the official repository issues, the author mentioned that the first row setting F1-score is around 0.75.
+* In the official paper, it is stated that the result F1-score of the second row setting is 0.87.
+ * If you adjust post-process parameter 'text_threshold' from 0.85 to 0.75, then F1-score reaches to 0.856.
+* It took 14h to train weak-supervision 25k iteration with 8 RTX 3090 Ti.
+ * Half of GPU assigned for training, and half of GPU assigned for supervision setting.
+
+| Training Dataset | Evaluation Dataset | Precision | Recall | F1-score | pretrained model |
+| ------------- |-----|:-----:|:-----:|:-----:|-----:|
+| SynthText | ICDAR2013 | 0.801 | 0.748 | 0.773| <a href="https://drive.google.com/file/d/1enVIsgNvBf3YiRsVkxodspOn55PIK-LJ/view?usp=sharing">download link</a>|
+| SynthText + ICDAR2015 | ICDAR2015 | 0.909 | 0.794 | 0.848| <a href="https://drive.google.com/file/d/1qUeZIDSFCOuGS9yo8o0fi-zYHLEW6lBP/view">download link</a>|

trainer/craft/config/custom_data_train.yaml

@@ -0,0 +1,100 @@
+wandb_opt: False
+
+results_dir: "./exp/"
+vis_test_dir: "./vis_result/"
+
+data_root_dir: "./data_root_dir/"
+score_gt_dir: None # "/data/ICDAR2015_official_supervision"
+mode: "weak_supervision"
+
+
+train:
+ backbone : vgg
+ use_synthtext: False # If you want to combine SynthText in train time as CRAFT did, you can turn on this option
+ synth_data_dir: "/data/SynthText/"
+ synth_ratio: 5
+ real_dataset: custom
+ ckpt_path: "./pretrained_model/CRAFT_clr_amp_29500.pth"
+ eval_interval: 1000
+ batch_size: 5
+ st_iter: 0
+ end_iter: 25000
+ lr: 0.0001
+ lr_decay: 7500
+ gamma: 0.2
+ weight_decay: 0.00001
+ num_workers: 0 # On single gpu, train.py execution only works when num worker = 0 / On multi-gpu, you can set num_worker > 0 to speed up
+ amp: True
+ loss: 2
+ neg_rto: 0.3
+ n_min_neg: 5000
+ data:
+ vis_opt: False
+ pseudo_vis_opt: False
+ output_size: 768
+ do_not_care_label: ['###', '']
+ mean: [0.485, 0.456, 0.406]
+ variance: [0.229, 0.224, 0.225]
+ enlarge_region : [0.5, 0.5] # x axis, y axis
+ enlarge_affinity: [0.5, 0.5]
+ gauss_init_size: 200
+ gauss_sigma: 40
+ watershed:
+ version: "skimage"
+ sure_fg_th: 0.75
+ sure_bg_th: 0.05
+ syn_sample: -1
+ custom_sample: -1
+ syn_aug:
+ random_scale:
+ range: [1.0, 1.5, 2.0]
+ option: False
+ random_rotate:
+ max_angle: 20
+ option: False
+ random_crop:
+ version: "random_resize_crop_synth"
+ option: True
+ random_horizontal_flip:
+ option: False
+ random_colorjitter:
+ brightness: 0.2
+ contrast: 0.2
+ saturation: 0.2
+ hue: 0.2
+ option: True
+ custom_aug:
+ random_scale:
+ range: [ 1.0, 1.5, 2.0 ]
+ option: False
+ random_rotate:
+ max_angle: 20
+ option: True
+ random_crop:
+ version: "random_resize_crop"
+ scale: [0.03, 0.4]
+ ratio: [0.75, 1.33]
+ rnd_threshold: 1.0
+ option: True
+ random_horizontal_flip:
+ option: True
+ random_colorjitter:
+ brightness: 0.2
+ contrast: 0.2
+ saturation: 0.2
+ hue: 0.2
+ option: True
+
+test:
+ trained_model : null
+ custom_data:
+ test_set_size: 500
+ test_data_dir: "./data_root_dir/"
+ text_threshold: 0.75
+ low_text: 0.5
+ link_threshold: 0.2
+ canvas_size: 2240
+ mag_ratio: 1.75
+ poly: False
+ cuda: True
+ vis_opt: False

trainer/craft/config/load_config.py

@@ -0,0 +1,37 @@
+import os
+import yaml
+from functools import reduce
+
+CONFIG_PATH = os.path.dirname(__file__)
+
+def load_yaml(config_name):
+
+ with open(os.path.join(CONFIG_PATH, config_name)+ '.yaml') as file:
+ config = yaml.safe_load(file)
+
+ return config
+
+class DotDict(dict):
+ def __getattr__(self, k):
+ try:
+ v = self[k]
+ except:
+ return super().__getattr__(k)
+ if isinstance(v, dict):
+ return DotDict(v)
+ return v
+
+ def __getitem__(self, k):
+ if isinstance(k, str) and '.' in k:
+ k = k.split('.')
+ if isinstance(k, (list, tuple)):
+ return reduce(lambda d, kk: d[kk], k, self)
+ return super().__getitem__(k)
+
+ def get(self, k, default=None):
+ if isinstance(k, str) and '.' in k:
+ try:
+ return self[k]
+ except KeyError:
+ return default
+ return super().get(k, default=default)

trainer/craft/config/syn_train.yaml

@@ -0,0 +1,68 @@
+wandb_opt: False
+
+results_dir: "./exp/"
+vis_test_dir: "./vis_result/"
+data_dir:
+ synthtext: "/data/SynthText/"
+ synthtext_gt: NULL
+
+train:
+ backbone : vgg
+ dataset: ["synthtext"]
+ ckpt_path: null
+ eval_interval: 1000
+ batch_size: 5
+ st_iter: 0
+ end_iter: 50000
+ lr: 0.0001
+ lr_decay: 15000
+ gamma: 0.2
+ weight_decay: 0.00001
+ num_workers: 4
+ amp: True
+ loss: 3
+ neg_rto: 1
+ n_min_neg: 1000
+ data:
+ vis_opt: False
+ output_size: 768
+ mean: [0.485, 0.456, 0.406]
+ variance: [0.229, 0.224, 0.225]
+ enlarge_region : [0.5, 0.5] # x axis, y axis
+ enlarge_affinity: [0.5, 0.5]
+ gauss_init_size: 200
+ gauss_sigma: 40
+ syn_sample : -1
+ syn_aug:
+ random_scale:
+ range: [1.0, 1.5, 2.0]
+ option: False
+ random_rotate:
+ max_angle: 20
+ option: False
+ random_crop:
+ version: "random_resize_crop_synth"
+ rnd_threshold : 1.0
+ option: True
+ random_horizontal_flip:
+ option: False
+ random_colorjitter:
+ brightness: 0.2
+ contrast: 0.2
+ saturation: 0.2
+ hue: 0.2
+ option: True
+
+test:
+ trained_model: null
+ icdar2013:
+ test_set_size: 233
+ cuda: True
+ vis_opt: True
+ test_data_dir : "/data/ICDAR2013/"
+ text_threshold: 0.85
+ low_text: 0.5
+ link_threshold: 0.2
+ canvas_size: 960
+ mag_ratio: 1.5
+ poly: False

trainer/craft/data/boxEnlarge.py

@@ -0,0 +1,65 @@
+import math
+import numpy as np
+
+
+def pointAngle(Apoint, Bpoint):
+ angle = (Bpoint[1] - Apoint[1]) / ((Bpoint[0] - Apoint[0]) + 10e-8)
+ return angle
+
+def pointDistance(Apoint, Bpoint):
+ return math.sqrt((Bpoint[1] - Apoint[1])**2 + (Bpoint[0] - Apoint[0])**2)
+
+def lineBiasAndK(Apoint, Bpoint):
+
+ K = pointAngle(Apoint, Bpoint)
+ B = Apoint[1] - K*Apoint[0]
+ return K, B
+
+def getX(K, B, Ypoint):
+ return int((Ypoint-B)/K)
+
+def sidePoint(Apoint, Bpoint, h, w, placehold, enlarge_size):
+
+ K, B = lineBiasAndK(Apoint, Bpoint)
+ angle = abs(math.atan(pointAngle(Apoint, Bpoint)))
+ distance = pointDistance(Apoint, Bpoint)
+
+ x_enlarge_size, y_enlarge_size = enlarge_size
+
+ XaxisIncreaseDistance = abs(math.cos(angle) * x_enlarge_size * distance)
+ YaxisIncreaseDistance = abs(math.sin(angle) * y_enlarge_size * distance)
+
+ if placehold == 'leftTop':
+ x1 = max(0, Apoint[0] - XaxisIncreaseDistance)
+ y1 = max(0, Apoint[1] - YaxisIncreaseDistance)
+ elif placehold == 'rightTop':
+ x1 = min(w, Bpoint[0] + XaxisIncreaseDistance)
+ y1 = max(0, Bpoint[1] - YaxisIncreaseDistance)
+ elif placehold == 'rightBottom':
+ x1 = min(w, Bpoint[0] + XaxisIncreaseDistance)
+ y1 = min(h, Bpoint[1] + YaxisIncreaseDistance)
+ elif placehold == 'leftBottom':
+ x1 = max(0, Apoint[0] - XaxisIncreaseDistance)
+ y1 = min(h, Apoint[1] + YaxisIncreaseDistance)
+ return int(x1), int(y1)
+
+def enlargebox(box, h, w, enlarge_size, horizontal_text_bool):
+
+ if not horizontal_text_bool:
+ enlarge_size = (enlarge_size[1], enlarge_size[0])
+
+ box = np.roll(box, -np.argmin(box.sum(axis=1)), axis=0)
+
+ Apoint, Bpoint, Cpoint, Dpoint = box
+ K1, B1 = lineBiasAndK(box[0], box[2])
+ K2, B2 = lineBiasAndK(box[3], box[1])
+ X = (B2 - B1)/(K1 - K2)
+ Y = K1 * X + B1
+ center = [X, Y]
+
+ x1, y1 = sidePoint(Apoint, center, h, w, 'leftTop', enlarge_size)
+ x2, y2 = sidePoint(center, Bpoint, h, w, 'rightTop', enlarge_size)
+ x3, y3 = sidePoint(center, Cpoint, h, w, 'rightBottom', enlarge_size)
+ x4, y4 = sidePoint(Dpoint, center, h, w, 'leftBottom', enlarge_size)
+ newcharbox = np.array([[x1, y1], [x2, y2], [x3, y3], [x4, y4]])
+ return newcharbox