train for SynthText data

Author: backtime92

.gitignore

@@ -2,4 +2,14 @@
 data/test/icdar2013/*.jpg
 data/test/icdar2015/*.jpg
 pretrain/*.pth
-__pycache__
+__init__.py
+craft.py
+craft_ic15_20k.pth
+craft_utils.py
+enlarge.py
+file_utils.py
+requirments.txt
+watershed.py
+basenet/*
+weights/*
+__pycache__

README.md

@@ -1,7 +1,6 @@
 # CRAFT-Reimplementation
 # Note：If you have any problems, please comment. Or you can join us weChat group. The QR code will update in issues #49 .
-# 更新：重构工作已经开始，由于重构全部都是工作之外的时间，所以重构及训练周期相对较长。目前进度新的高斯图生成方式已经实验完成，对比作者论文中的图，结合以前的实验探索和向作者的提问，这次是与作者一样的方式，预计在20号左右release合成数据的训练部分。由于公司连接github较卡，暂时无法上传高斯图生成这部分。
-# 更新2：目前已经在训练合成数据全监督这部分，同时弱监督部分同时在coding及实验。如果合成数据部分训练效果达到预期，预计周五会release这部分训练code。
+# 非常抱歉，一直没有继续维护这个工程，近期看到挺多人关注的，我预计11月底最晚12月初重新维护一下工程，由于当时实习期间整体工程能力不够导致工程相对较乱，这个重新维护会整理清楚。维护周期大概是两周时间，我会重新整理code以及重新训练同时上传训练的pretrain model。同时一些实验的关键和实验思路我会写上注释，欢迎到时关注。
 
 ## Reimplementation：Character Region Awareness for Text Detection Reimplementation based on Pytorch
 
@@ -19,83 +18,6 @@ The full paper is available at: https://arxiv.org/pdf/1904.01941.pdf
 5、4 nvidia GPUs(we use 4 nvidia titanX)                                                                                      
 
 
-## pre-trained model:
-`NOTE: There are old pre-trained models, I will upload the new results pre-trained models' link.`                                                                                
-Syndata:[Syndata for baidu drive](https://pan.baidu.com/s/1MaznjE79JNS9Ld48ZtRefg) ||     [Syndata for google drive](https://drive.google.com/file/d/1FvqfBMZQJeZXGfZLl-840YXoeYK8CNwk/view?usp=sharing)                                                                                                    
-Syndata+IC15:[Syndata+IC15 for baidu drive](https://pan.baidu.com/s/19lJRM6YWZXVkZ_aytsYSiQ) ||      [Syndata+IC15 for google
- drive](https://drive.google.com/file/d/1k17GuBG_omT91tJoIMSlLrorYbLXkq4z/view?usp=sharing)                                   
- Syndata+IC13+IC17:[Syndata+IC13+IC17 for baidu drive](https://pan.baidu.com/s/1PTTzbM9XG0pNe5i-uL6Aag)||      [Syndata+IC13+IC17 for google drive](https://drive.google.com/open?id=1SkJEfaGYIq-eFxfzFVZb-cGdGWR8lPSi) 
-
-
-## Training 
-`Note: When you train the IC15-Data or MLT-Data, please see the annotation in data_loader.py line 92 and line 108-112.`
-
-### Train for Syndata
-- download the Syndata(I will give the link)
-- change the path in basernet/vgg16_bn.py file:
->` (/data/CRAFT-pytorch/vgg16_bn-6c64b313.pth -> /your_path/vgg16_bn-6c64b313.pth).You can download the model here.`[baidu](https://pan.baidu.com/s/1_h5qdwYQAToDi_BB5Eg3vg)||[google](https://drive.google.com/open?id=1ZtvGpFQrbmEisB_GhmZb8UQOtvqY_-tW)                                                                 
-- change the path in trainSyndata.py file:
-> `(1、/data/CRAFT-pytorch/SynthText -> /your_path/SynthText 2、/data/CRAFT-pytorch/synweights/synweights -> /your_path/real_weights)`                                                                      
-- Run **`python trainSyndata.py`**
-
-### Train for IC15 data based on Syndata pre-trained model
-- download the IC15 data, rename the image file and the gt file for  ch4_training_images and ch4_training_localization_transcription_gt,respectively.
-- change the path in basernet/vgg16_bn.py file:                                                                                                                                                              
-> `(/data/CRAFT-pytorch/vgg16_bn-6c64b313.pth -> /your_path/vgg16_bn-6c64b313.pth).You can download the model here.`[baidu](https://pan.baidu.com/s/1_h5qdwYQAToDi_BB5Eg3vg)||[google](https://drive.google.com/open?id=1ZtvGpFQrbmEisB_GhmZb8UQOtvqY_-tW)
-- change the path in trainic15data.py file:                                                                                                                                                                  
->` (1、/data/CRAFT-pytorch/SynthText -> /your_path/SynthText    2、/data/CRAFT-pytorch/real_weights -> /your_path/real_weights)`
-- change the path in trainic15data.py file:                                                                                                                                                                 
-> `(1、/data/CRAFT-pytorch/1-7.pth -> /your_path/your_pre-trained_model_name 2、/data/CRAFT-pytorch/icdar1317 -> /your_ic15data_path/)`
-- Run **`python trainic15data.py`**
-
-### Train for IC13+17 data based on Syndata pre-trained model
-
-- download the MLT data, rename the image file and the gt file,respectively.
-- change the path in basernet/vgg16_bn.py file:                                                                                                                                                              
-> `(/data/CRAFT-pytorch/vgg16_bn-6c64b313.pth -> /your_path/vgg16_bn-6c64b313.pth).You can download the model here.`[baidu](https://pan.baidu.com/s/1_h5qdwYQAToDi_BB5Eg3vg)||[google](https://drive.google.com/open?id=1ZtvGpFQrbmEisB_GhmZb8UQOtvqY_-tW)
-- change the path in trainic-MLT_data.py file:                                                                                                                                                              
->` (1、/data/CRAFT-pytorch/SynthText -> /your_path/SynthText    2、savemodel path-> your savemodel path)`
-- change the path in trainic-MLT_data.py file:                                                                                                                                                         
-> `(1、/data/CRAFT-pytorch/1-7.pth -> /your_path/your_pre-trained_model_name 2、/data/CRAFT-pytorch/icdar1317 -> /your_ic15data_path/)`
-- Run **`python trainic-MLT_data.py`**
-
-### If you want to train for weak supervised use our Syndate pre-trained model:                                                                                                                                
-1、You should first download the pre_trained model trained in the Syndata [baidu](https://pan.baidu.com/s/1MaznjE79JNS9Ld48ZtRefg)||[google](https://drive.google.com/file/d/1FvqfBMZQJeZXGfZLl-840YXoeYK8CNwk/view?usp=sharing).                                                                                                                                                      
-2、change the data path and pre-trained model path.                                                                                                                                                         
-3、run `python trainic15data.py`                                                                                                                                                                           
-
-                                                                                                                    
-**This code supprts for Syndata and icdar2015, and we will release the training code for IC13 and IC17 as soon as possible.**
-
-Methods                                       |dataset      |Recall      |precision      |H-mean
-----------------------------------------------|-------------|------------|---------------|------
-Syndata                                       |ICDAR13      |71.93%      |81.31%         |76.33%                                                                          
-Syndata+IC15                                  |ICDAR15      |76.12%      |84.55%         |80.11%               
-Syndata+MLT(deteval)                          |ICDAR13      |86.81%      |95.28%         |90.85%                                   
-Syndata+MLT(deteval)(new gaussian map method) |ICDAR13      |90.67%      |94.56%         |92.57%                                   
-Syndata+IC15(new gaussian map method)         |ICDAR15      |80.36%      |84.25%         |82.26%
-
-### We have released the latest code with new gaussian map and random crop algorithm. 
-**`Note:new gaussian map method can split the inference gaussian region score map`**                                                                                                                         
-`Sample:`                                                                                           
-<img src="https://github.com/backtime92/CRAFT-Reimplementation/blob/master/image/test3_score.jpg" width="384" height="512" /><img src="https://github.com/backtime92/CRAFT-Reimplementation/blob/master/image/test3_affinity.jpg" width="384" height="256" />                                                                                                                                                      
-
-**`Note:We have solved the problem about detecting big word. Now we are training the model. And any issues or advice are welcome.`**                                                                  
-                                                                                                                                                 
-`Sample:`
-<img src="https://github.com/backtime92/CRAFT-Reimplementation/blob/master/image/test4_score.jpg" width="384" height="512" /><img src="https://github.com/backtime92/CRAFT-Reimplementation/blob/master/image/test4_affinity.jpg" width="384" height="256" />                                                                                                                                    
-
-###weChat QR code
-<img src="https://github.com/backtime92/CRAFT-Reimplementation/blob/master/image/wechatgroup.jpeg" width="150" height="150" />                                                           
-
-
-# Contributing to the project
-`We will release training code as soon as possible， and we have not yet reached the results given in the author's paper. Any pull requests or issues are welcome. We also hope that you could give us some advice for the project.`
-
-# Acknowledgement
-Thanks for Youngmin Baek, Bado Lee, Dongyoon Han, Sangdoo Yun, Hwalsuk Lee excellent work and [code](https://github.com/clovaai/CRAFT-pytorch) for test. In this repo, we use the author repo's basenet and test code.
-
-# License
-For commercial use, please contact us.
-
+# to do list
+Release strong supervision training part in early December 
 

basenet/vgg16_bn.py

@@ -1,43 +1,48 @@
 from collections import namedtuple
 
 import torch
+import torch.nn as nn
+import torch.nn.init as init
 from torchvision import models
 from torchvision.models.vgg import model_urls
-from torchutil import *
-import os
-
-weights_folder = os.path.join(os.path.dirname(__file__) + '/../pretrain')
 
+def init_weights(modules):
+    for m in modules:
+        if isinstance(m, nn.Conv2d):
+            init.xavier_uniform_(m.weight.data)
+            if m.bias is not None:
+                m.bias.data.zero_()
+        elif isinstance(m, nn.BatchNorm2d):
+            m.weight.data.fill_(1)
+            m.bias.data.zero_()
+        elif isinstance(m, nn.Linear):
+            m.weight.data.normal_(0, 0.01)
+            m.bias.data.zero_()
 
 class vgg16_bn(torch.nn.Module):
-    def __init__(self, pretrained=True, freeze=False):
+    def __init__(self, pretrained=True, freeze=True):
         super(vgg16_bn, self).__init__()
         model_urls['vgg16_bn'] = model_urls['vgg16_bn'].replace('https://', 'http://')
-        # vgg_pretrained_features = models.vgg16_bn(pretrained=pretrained).features
-        vgg_pretrained_features = models.vgg16_bn(pretrained=False)
-        if pretrained:
-            vgg_pretrained_features.load_state_dict(
-                copyStateDict(torch.load(os.path.join(weights_folder, '/data/CRAFT-pytorch/vgg16_bn-6c64b313.pth'))))
-        vgg_pretrained_features = vgg_pretrained_features.features
+        vgg_pretrained_features = models.vgg16_bn(pretrained=pretrained).features
         self.slice1 = torch.nn.Sequential()
         self.slice2 = torch.nn.Sequential()
         self.slice3 = torch.nn.Sequential()
         self.slice4 = torch.nn.Sequential()
         self.slice5 = torch.nn.Sequential()
-        for x in range(12):  # conv2_2
+        for x in range(12):         # conv2_2
             self.slice1.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(12, 19):  # conv3_3
+        for x in range(12, 19):         # conv3_3
             self.slice2.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(19, 29):  # conv4_3
+        for x in range(19, 29):         # conv4_3
             self.slice3.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(29, 39):  # conv5_3
+        for x in range(29, 39):         # conv5_3
             self.slice4.add_module(str(x), vgg_pretrained_features[x])
 
         # fc6, fc7 without atrous conv
         self.slice5 = torch.nn.Sequential(
-            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
-            nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6),
-            nn.Conv2d(1024, 1024, kernel_size=1)
+                nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
+                nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6),
+                nn.Conv2d(1024, 1024, kernel_size=1)
         )
 
         if not pretrained:
@@ -46,11 +51,11 @@ def __init__(self, pretrained=True, freeze=False):
             init_weights(self.slice3.modules())
             init_weights(self.slice4.modules())
 
-        init_weights(self.slice5.modules())  # no pretrained model for fc6 and fc7
+        init_weights(self.slice5.modules())        # no pretrained model for fc6 and fc7
 
         if freeze:
-            for param in self.slice1.parameters():  # only first conv
-                param.requires_grad = False
+            for param in self.slice1.parameters():      # only first conv
+                param.requires_grad= False
 
     def forward(self, X):
         h = self.slice1(X)