Initial commit of EasyOCR model

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

MANIFEST.in

@@ -0,0 +1,8 @@
+include LICENSE.txt
+include README.md
+
+include easyocr/model/*
+include easyocr/character/*
+include easyocr/dict/*
+include easyocr/scripts/compile_dbnet_dcn.py
+recursive-include easyocr/DBNet *

README.md

@@ -0,0 +1,178 @@
+# EasyOCR
+
+[![PyPI Status](https://badge.fury.io/py/easyocr.svg)](https://badge.fury.io/py/easyocr)
+[![license](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/JaidedAI/EasyOCR/blob/master/LICENSE)
+[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.to/easyocr)
+[![Tweet](https://img.shields.io/twitter/url/https/github.com/JaidedAI/EasyOCR.svg?style=social)](https://twitter.com/intent/tweet?text=Check%20out%20this%20awesome%20library:%20EasyOCR%20https://github.com/JaidedAI/EasyOCR)
+[![Twitter](https://img.shields.io/badge/twitter-@JaidedAI-blue.svg?style=flat)](https://twitter.com/JaidedAI)
+
+Ready-to-use OCR with 80+ [supported languages](https://www.jaided.ai/easyocr) and all popular writing scripts including: Latin, Chinese, Arabic, Devanagari, Cyrillic, etc.
+
+[Try Demo on our website](https://www.jaided.ai/easyocr)
+
+Integrated into [Huggingface Spaces 🤗](https://huggingface.co/spaces) using [Gradio](https://github.com/gradio-app/gradio). Try out the Web Demo: [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/tomofi/EasyOCR)
+
+
+## What's new
+- 24 September 2024 - Version 1.7.2
+    - Fix several compatibilities
+
+- [Read all release notes](https://github.com/JaidedAI/EasyOCR/blob/master/releasenotes.md)
+
+## What's coming next
+- Handwritten text support
+
+## Examples
+
+![example](examples/example.png)
+
+![example2](examples/example2.png)
+
+![example3](examples/example3.png)
+
+
+## Installation
+
+Install using `pip`
+
+For the latest stable release:
+
+``` bash
+pip install easyocr
+```
+
+For the latest development release:
+
+``` bash
+pip install git+https://github.com/JaidedAI/EasyOCR.git
+```
+
+Note 1: For Windows, please install torch and torchvision first by following the official instructions here https://pytorch.org. On the pytorch website, be sure to select the right CUDA version you have. If you intend to run on CPU mode only, select `CUDA = None`.
+
+Note 2: We also provide a Dockerfile [here](https://github.com/JaidedAI/EasyOCR/blob/master/Dockerfile).
+
+## Usage
+
+``` python
+import easyocr
+reader = easyocr.Reader(['ch_sim','en']) # this needs to run only once to load the model into memory
+result = reader.readtext('chinese.jpg')
+```
+
+The output will be in a list format, each item represents a bounding box, the text detected and confident level, respectively.
+
+``` bash
+[([[189, 75], [469, 75], [469, 165], [189, 165]], '愚园路', 0.3754989504814148),
+ ([[86, 80], [134, 80], [134, 128], [86, 128]], '西', 0.40452659130096436),
+ ([[517, 81], [565, 81], [565, 123], [517, 123]], '东', 0.9989598989486694),
+ ([[78, 126], [136, 126], [136, 156], [78, 156]], '315', 0.8125889301300049),
+ ([[514, 126], [574, 126], [574, 156], [514, 156]], '309', 0.4971577227115631),
+ ([[226, 170], [414, 170], [414, 220], [226, 220]], 'Yuyuan Rd.', 0.8261902332305908),
+ ([[79, 173], [125, 173], [125, 213], [79, 213]], 'W', 0.9848111271858215),
+ ([[529, 173], [569, 173], [569, 213], [529, 213]], 'E', 0.8405593633651733)]
+```
+Note 1: `['ch_sim','en']` is the list of languages you want to read. You can pass
+several languages at once but not all languages can be used together.
+English is compatible with every language and languages that share common characters are usually compatible with each other.
+
+Note 2: Instead of the filepath `chinese.jpg`, you can also pass an OpenCV image object (numpy array) or an image file as bytes. A URL to a raw image is also acceptable.
+
+Note 3: The line `reader = easyocr.Reader(['ch_sim','en'])` is for loading a model into memory. It takes some time but it needs to be run only once.
+
+You can also set `detail=0` for simpler output.
+
+``` python
+reader.readtext('chinese.jpg', detail = 0)
+```
+Result:
+``` bash
+['愚园路', '西', '东', '315', '309', 'Yuyuan Rd.', 'W', 'E']
+```
+
+Model weights for the chosen language will be automatically downloaded or you can
+download them manually from the [model hub](https://www.jaided.ai/easyocr/modelhub) and put them in the '~/.EasyOCR/model' folder
+
+In case you do not have a GPU, or your GPU has low memory, you can run the model in CPU-only mode by adding `gpu=False`.
+
+``` python
+reader = easyocr.Reader(['ch_sim','en'], gpu=False)
+```
+
+For more information, read the [tutorial](https://www.jaided.ai/easyocr/tutorial) and [API Documentation](https://www.jaided.ai/easyocr/documentation).
+
+#### Run on command line
+
+```shell
+$ easyocr -l ch_sim en -f chinese.jpg --detail=1 --gpu=True
+```
+
+## Train/use your own model
+
+For recognition model, [Read here](https://github.com/JaidedAI/EasyOCR/blob/master/custom_model.md).
+
+For detection model (CRAFT), [Read here](https://github.com/JaidedAI/EasyOCR/blob/master/trainer/craft/README.md).
+
+## Implementation Roadmap
+
+- Handwritten support
+- Restructure code to support swappable detection and recognition algorithms
+The api should be as easy as
+``` python
+reader = easyocr.Reader(['en'], detection='DB', recognition = 'Transformer')
+```
+The idea is to be able to plug in any state-of-the-art model into EasyOCR. There are a lot of geniuses trying to make better detection/recognition models, but we are not trying to be geniuses here. We just want to make their works quickly accessible to the public ... for free. (well, we believe most geniuses want their work to create a positive impact as fast/big as possible) The pipeline should be something like the below diagram. Grey slots are placeholders for changeable light blue modules.
+
+![plan](examples/easyocr_framework.jpeg)
+
+## Acknowledgement and References
+
+This project is based on research and code from several papers and open-source repositories.
+
+All deep learning execution is based on [Pytorch](https://pytorch.org). :heart:
+
+Detection execution uses the CRAFT algorithm from this [official repository](https://github.com/clovaai/CRAFT-pytorch) and their [paper](https://arxiv.org/abs/1904.01941) (Thanks @YoungminBaek from [@clovaai](https://github.com/clovaai)). We also use their pretrained model. Training script is provided by [@gmuffiness](https://github.com/gmuffiness).
+
+The recognition model is a CRNN ([paper](https://arxiv.org/abs/1507.05717)). It is composed of 3 main components: feature extraction (we are currently using [Resnet](https://arxiv.org/abs/1512.03385)) and VGG, sequence labeling ([LSTM](https://www.bioinf.jku.at/publications/older/2604.pdf)) and decoding ([CTC](https://www.cs.toronto.edu/~graves/icml_2006.pdf)). The training pipeline for recognition execution is a modified version of the [deep-text-recognition-benchmark](https://github.com/clovaai/deep-text-recognition-benchmark) framework. (Thanks [@ku21fan](https://github.com/ku21fan) from [@clovaai](https://github.com/clovaai)) This repository is a gem that deserves more recognition.
+
+Beam search code is based on this [repository](https://github.com/githubharald/CTCDecoder) and his [blog](https://towardsdatascience.com/beam-search-decoding-in-ctc-trained-neural-networks-5a889a3d85a7). (Thanks [@githubharald](https://github.com/githubharald))
+
+Data synthesis is based on [TextRecognitionDataGenerator](https://github.com/Belval/TextRecognitionDataGenerator). (Thanks [@Belval](https://github.com/Belval))
+
+And a good read about CTC from distill.pub [here](https://distill.pub/2017/ctc/).
+
+## Want To Contribute?
+
+Let's advance humanity together by making AI available to everyone!
+
+3 ways to contribute:
+
+**Coder:** Please send a PR for small bugs/improvements. For bigger ones, discuss with us by opening an issue first. There is a list of possible bug/improvement issues tagged with ['PR WELCOME'](https://github.com/JaidedAI/EasyOCR/issues?q=is%3Aissue+is%3Aopen+label%3A%22PR+WELCOME%22).
+
+**User:** Tell us how EasyOCR benefits you/your organization to encourage further development. Also post failure cases in [Issue  Section](https://github.com/JaidedAI/EasyOCR/issues) to help improve future models.
+
+**Tech leader/Guru:** If you found this library useful, please spread the word! (See [Yann Lecun's post](https://www.facebook.com/yann.lecun/posts/10157018122787143) about EasyOCR)
+
+## Guideline for new language request
+
+To request a new language, we need you to send a PR with the 2 following files:
+
+1. In folder [easyocr/character](https://github.com/JaidedAI/EasyOCR/tree/master/easyocr/character),
+we need 'yourlanguagecode_char.txt' that contains list of all characters. Please see format examples from other files in that folder.
+2. In folder [easyocr/dict](https://github.com/JaidedAI/EasyOCR/tree/master/easyocr/dict),
+we need 'yourlanguagecode.txt' that contains list of words in your language.
+On average, we have ~30000 words per language with more than 50000 words for more popular ones.
+More is better in this file.
+
+If your language has unique elements (such as 1. Arabic: characters change form when attached to each other + write from right to left 2. Thai: Some characters need to be above the line and some below), please educate us to the best of your ability and/or give useful links. It is important to take care of the detail to achieve a system that really works.
+
+Lastly, please understand that our priority will have to go to popular languages or sets of languages that share large portions of their characters with each other (also tell us if this is the case for your language). It takes us at least a week to develop a new model, so you may have to wait a while for the new model to be released.
+
+See [List of languages in development](https://github.com/JaidedAI/EasyOCR/issues/91)
+
+## Github Issues
+
+Due to limited resources, an issue older than 6 months will be automatically closed. Please open an issue again if it is critical.
+
+## Business Inquiries
+
+For Enterprise Support, [Jaided AI](https://www.jaided.ai/) offers full service for custom OCR/AI systems from implementation, training/finetuning and deployment. Click [here](https://www.jaided.ai/contactus?ref=github) to contact us.

custom_model.md

@@ -0,0 +1,24 @@
+# Custom recognition models
+
+## How to train your custom model
+
+You can use your own data or generate your own dataset. To generate your own data, we recommend using
+[TextRecognitionDataGenerator](https://github.com/Belval/TextRecognitionDataGenerator). We provide an example of a dataset [here](https://jaided.ai/easyocr/modelhub/).
+After you have a dataset, you can train your own model by following this repository
+[deep-text-recognition-benchmark](https://github.com/clovaai/deep-text-recognition-benchmark).
+The network needs to be fully convolutional in order to predict flexible text length. Our current network is 'None-VGG-BiLSTM-CTC'.
+Once you have your trained model (a `.pth` file), you need 2 additional files describing recognition network architecture and model configuration.
+An example is provided in `custom_example.zip` file [here](https://jaided.ai/easyocr/modelhub/).
+
+Please do not create an issue about data generation and model training in this repository. If you have any question regarding data generation and model training, please ask in the respective repositories.
+
+Note: We also provide our version of a training script [here](https://github.com/JaidedAI/EasyOCR/tree/master/trainer). It is a modified version from [deep-text-recognition-benchmark](https://github.com/clovaai/deep-text-recognition-benchmark).
+
+## How to use your custom model
+
+To use your own recognition model, you need the three files as explained above. These three files have to share the same name (i.e. `yourmodel.pth`, `yourmodel.yaml`, `yourmodel.py`) that you will then use to call your model with EasyOCR API.
+
+We provide [custom_example.zip](https://jaided.ai/easyocr/modelhub/)
+as an example. Please download, extract and place `custom_example.py`, `custom_example.yaml` in the `user_network_directory` (default = `~/.EasyOCR/user_network`) and place `custom_example.pth` in model directory (default = `~/.EasyOCR/model`)
+Once you place all 3 files in their respective places, you can use `custom_example` by
+specifying `recog_network` like this `reader = easyocr.Reader(['en'], recog_network='custom_example')`.

model.py

@@ -0,0 +1,24 @@
+import easyocr
+from typing import List
+
+class EasyOCRModel:
+    def __init__(self):
+        self.reader = easyocr.Reader(['en'])  # Initialize with English; add languages if needed.
+
+    def predict(self, image_path: str) -> List[str]:
+        """
+        Perform OCR on the given image.
+
+        Args:
+            image_path (str): Path to the input image.
+
+        Returns:
+            List[str]: Extracted text from the image.
+        """
+        return self.reader.readtext(image_path, detail=0)
+
+# Test the model locally
+if __name__ == "__main__":
+    model = EasyOCRModel()
+    result = model.predict("sample_image.jpg")
+    print(result)

requirements.txt

@@ -0,0 +1,12 @@
+torch
+torchvision>=0.5
+opencv-python-headless
+scipy
+numpy
+Pillow
+scikit-image
+python-bidi
+PyYAML
+Shapely
+pyclipper
+ninja

scripts/.gitignore

@@ -0,0 +1,2 @@
+JMnedict*
+JMdict*

scripts/generate-ja.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+require 'json'
+require 'nokogiri'
+require 'parallel'
+require 'ruby-progressbar'
+
+JMDICT_XML = 'JMdict_e'
+JMNEDICT_XML = 'JMnedict.xml'
+PUNC = '【】《》〈〉｟｠｛｝［］〔〕（）『』「」、；：・？〜＝。！⁉︎‥…〜※＊〽♪♫♬♩〇〒〶〠〄ⓍⓁⓎ→'.chars
+
+def download_dict(xml)
+  return if File.exist?(File.expand_path(xml, __dir__))
+
+  archive = "#{xml}.gz"
+  url = "http://ftp.monash.edu/pub/nihongo/#{archive}"
+  `cd #{File.dirname(__FILE__)} && wget #{url} && gunzip #{archive}`
+end
+
+def read_word(word)
+  word.css('k_ele keb').map(&:text) + word.css('r_ele reb').map(&:text)
+end
+
+def read_dict(filename, root)
+  xml = Nokogiri::XML(File.open(File.expand_path(filename, __dir__)))
+  words = xml.css("#{root} > entry")
+  Parallel.flat_map(words, in_threads: 16, progress: root) do |word|
+    read_word(word)
+  end
+end
+
+def write_files(words)
+  src_dir = File.expand_path('../easyocr', __dir__)
+  ja_dict = File.join(src_dir, 'dict', 'ja.txt')
+  ja_char = File.join(src_dir, 'character', 'ja_char2.txt')
+  ja_char_old = File.join(src_dir, 'character', 'ja_char.txt')
+  ja_punc = File.join(src_dir, 'character', 'ja_punc.txt')
+
+  words -= PUNC
+  chars = words.join.chars.uniq
+  chars_old = IO.read(ja_char_old).split("\n")
+
+  puts "new characters: #{(chars - chars_old).size}"
+  puts "missing characters: #{(chars_old - chars).size}"
+  puts chars_old - chars
+
+  IO.write(ja_dict, words.join("\n"))
+  IO.write(ja_char, chars.join("\n"))
+  IO.write(ja_punc, PUNC.join("\n"))
+end
+
+download_dict(JMDICT_XML)
+download_dict(JMNEDICT_XML)
+words = read_dict(JMDICT_XML, 'JMdict') + read_dict(JMNEDICT_XML, 'JMnedict')
+write_files(words)

setup.cfg

@@ -0,0 +1,2 @@
+[metadata]
+description_file = README.md

setup.py

@@ -0,0 +1,35 @@
+"""
+End-to-End Multi-Lingual Optical Character Recognition (OCR) Solution
+"""
+from io import open
+from setuptools import setup
+
+with open('requirements.txt', encoding="utf-8-sig") as f:
+    requirements = f.readlines()
+
+def readme():
+    with open('README.md', encoding="utf-8-sig") as f:
+        README = f.read()
+    return README
+
+setup(
+    name='easyocr',
+    packages=['easyocr'],
+    include_package_data=True,
+    version='1.7.2',
+    install_requires=requirements,
+    entry_points={"console_scripts": ["easyocr= easyocr.cli:main"]},
+    license='Apache License 2.0',
+    description='End-to-End Multi-Lingual Optical Character Recognition (OCR) Solution',
+    long_description=readme(),
+    long_description_content_type="text/markdown",
+    author='Rakpong Kittinaradorn',
+    author_email='r.kittinaradorn@gmail.com',
+    url='https://github.com/jaidedai/easyocr',
+    download_url='https://github.com/jaidedai/easyocr.git',
+    keywords=['ocr optical character recognition deep learning neural network'],
+    classifiers=[
+        'Development Status :: 5 - Production/Stable'
+    ],
+
+)

trainer/README.md

@@ -0,0 +1,3 @@
+# EasyOCR trainer
+
+use `trainer.ipynb` with yaml config in `config_files` folder

trainer/all_data/folder.txt

@@ -0,0 +1 @@
+place dataset folder here

trainer/config_files/en_filtered_config.yaml

@@ -0,0 +1,45 @@
+number: '0123456789'
+symbol: "!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~ €"
+lang_char: 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
+experiment_name: 'en_filtered'
+train_data: 'all_data'
+valid_data: 'all_data/en_val'
+manualSeed: 1111
+workers: 6
+batch_size: 32 #32
+num_iter: 300000
+valInterval: 20000
+saved_model: '' #'saved_models/en_filtered/iter_300000.pth'
+FT: False
+optim: False # default is Adadelta
+lr: 1.
+beta1: 0.9
+rho: 0.95
+eps: 0.00000001
+grad_clip: 5
+#Data processing
+select_data: 'en_train_filtered' # this is dataset folder in train_data
+batch_ratio: '1' 
+total_data_usage_ratio: 1.0
+batch_max_length: 34 
+imgH: 64
+imgW: 600
+rgb: False
+contrast_adjust: False
+sensitive: True
+PAD: True
+contrast_adjust: 0.0
+data_filtering_off: False
+# Model Architecture
+Transformation: 'None'
+FeatureExtraction: 'VGG'
+SequenceModeling: 'BiLSTM'
+Prediction: 'CTC'
+num_fiducial: 20
+input_channel: 1
+output_channel: 256
+hidden_size: 256
+decode: 'greedy'
+new_prediction: False
+freeze_FeatureFxtraction: False
+freeze_SequenceModeling: False

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

trainer/craft/data/dataset.py

@@ -0,0 +1,542 @@
+import os
+import re
+import itertools
+import random
+
+import numpy as np
+import scipy.io as scio
+from PIL import Image
+import cv2
+from torch.utils.data import Dataset
+import torchvision.transforms as transforms
+
+from data import imgproc
+from data.gaussian import GaussianBuilder
+from data.imgaug import (
+    rescale,
+    random_resize_crop_synth,
+    random_resize_crop,
+    random_horizontal_flip,
+    random_rotate,
+    random_scale,
+    random_crop,
+)
+from data.pseudo_label.make_charbox import PseudoCharBoxBuilder
+from utils.util import saveInput, saveImage
+
+
+class CraftBaseDataset(Dataset):
+    def __init__(
+        self,
+        output_size,
+        data_dir,
+        saved_gt_dir,
+        mean,
+        variance,
+        gauss_init_size,
+        gauss_sigma,
+        enlarge_region,
+        enlarge_affinity,
+        aug,
+        vis_test_dir,
+        vis_opt,
+        sample,
+    ):
+        self.output_size = output_size
+        self.data_dir = data_dir
+        self.saved_gt_dir = saved_gt_dir
+        self.mean, self.variance = mean, variance
+        self.gaussian_builder = GaussianBuilder(
+            gauss_init_size, gauss_sigma, enlarge_region, enlarge_affinity
+        )
+        self.aug = aug
+        self.vis_test_dir = vis_test_dir
+        self.vis_opt = vis_opt
+        self.sample = sample
+        if self.sample != -1:
+            random.seed(0)
+            self.idx = random.sample(range(0, len(self.img_names)), self.sample)
+
+        self.pre_crop_area = []
+
+    def augment_image(
+        self, image, region_score, affinity_score, confidence_mask, word_level_char_bbox
+    ):
+        augment_targets = [image, region_score, affinity_score, confidence_mask]
+
+        if self.aug.random_scale.option:
+            augment_targets, word_level_char_bbox = random_scale(
+                augment_targets, word_level_char_bbox, self.aug.random_scale.range
+            )
+
+        if self.aug.random_rotate.option:
+            augment_targets = random_rotate(
+                augment_targets, self.aug.random_rotate.max_angle
+            )
+
+        if self.aug.random_crop.option:
+            if self.aug.random_crop.version == "random_crop_with_bbox":
+                augment_targets = random_crop_with_bbox(
+                    augment_targets, word_level_char_bbox, self.output_size
+                )
+            elif self.aug.random_crop.version == "random_resize_crop_synth":
+                augment_targets = random_resize_crop_synth(
+                    augment_targets, self.output_size
+                )
+            elif self.aug.random_crop.version == "random_resize_crop":
+
+                if len(self.pre_crop_area) > 0:
+                    pre_crop_area = self.pre_crop_area
+                else:
+                    pre_crop_area = None
+
+                augment_targets = random_resize_crop(
+                    augment_targets,
+                    self.aug.random_crop.scale,
+                    self.aug.random_crop.ratio,
+                    self.output_size,
+                    self.aug.random_crop.rnd_threshold,
+                    pre_crop_area,
+                )
+
+            elif self.aug.random_crop.version == "random_crop":
+                augment_targets = random_crop(augment_targets, self.output_size,)
+
+            else:
+                assert "Undefined RandomCrop version"
+
+        if self.aug.random_horizontal_flip.option:
+            augment_targets = random_horizontal_flip(augment_targets)
+
+        if self.aug.random_colorjitter.option:
+            image, region_score, affinity_score, confidence_mask = augment_targets
+            image = Image.fromarray(image)
+            image = transforms.ColorJitter(
+                brightness=self.aug.random_colorjitter.brightness,
+                contrast=self.aug.random_colorjitter.contrast,
+                saturation=self.aug.random_colorjitter.saturation,
+                hue=self.aug.random_colorjitter.hue,
+            )(image)
+        else:
+            image, region_score, affinity_score, confidence_mask = augment_targets
+
+        return np.array(image), region_score, affinity_score, confidence_mask
+
+    def resize_to_half(self, ground_truth, interpolation):
+        return cv2.resize(
+            ground_truth,
+            (self.output_size // 2, self.output_size // 2),
+            interpolation=interpolation,
+        )
+
+    def __len__(self):
+        if self.sample != -1:
+            return len(self.idx)
+        else:
+            return len(self.img_names)
+
+    def __getitem__(self, index):
+        if self.sample != -1:
+            index = self.idx[index]
+        if self.saved_gt_dir is None:
+            (
+                image,
+                region_score,
+                affinity_score,
+                confidence_mask,
+                word_level_char_bbox,
+                all_affinity_bbox,
+                words,
+            ) = self.make_gt_score(index)
+        else:
+            (
+                image,
+                region_score,
+                affinity_score,
+                confidence_mask,
+                word_level_char_bbox,
+                words,
+            ) = self.load_saved_gt_score(index)
+            all_affinity_bbox = []
+
+        if self.vis_opt:
+            saveImage(
+                self.img_names[index],
+                self.vis_test_dir,
+                image.copy(),
+                word_level_char_bbox.copy(),
+                all_affinity_bbox.copy(),
+                region_score.copy(),
+                affinity_score.copy(),
+                confidence_mask.copy(),
+            )
+
+        image, region_score, affinity_score, confidence_mask = self.augment_image(
+            image, region_score, affinity_score, confidence_mask, word_level_char_bbox
+        )
+
+        if self.vis_opt:
+            saveInput(
+                self.img_names[index],
+                self.vis_test_dir,
+                image,
+                region_score,
+                affinity_score,
+                confidence_mask,
+            )
+
+        region_score = self.resize_to_half(region_score, interpolation=cv2.INTER_CUBIC)
+        affinity_score = self.resize_to_half(
+            affinity_score, interpolation=cv2.INTER_CUBIC
+        )
+        confidence_mask = self.resize_to_half(
+            confidence_mask, interpolation=cv2.INTER_NEAREST
+        )
+
+        image = imgproc.normalizeMeanVariance(
+            np.array(image), mean=self.mean, variance=self.variance
+        )
+        image = image.transpose(2, 0, 1)
+
+        return image, region_score, affinity_score, confidence_mask
+
+
+class SynthTextDataSet(CraftBaseDataset):
+    def __init__(
+        self,
+        output_size,
+        data_dir,
+        saved_gt_dir,
+        mean,
+        variance,
+        gauss_init_size,
+        gauss_sigma,
+        enlarge_region,
+        enlarge_affinity,
+        aug,
+        vis_test_dir,
+        vis_opt,
+        sample,
+    ):
+        super().__init__(
+            output_size,
+            data_dir,
+            saved_gt_dir,
+            mean,
+            variance,
+            gauss_init_size,
+            gauss_sigma,
+            enlarge_region,
+            enlarge_affinity,
+            aug,
+            vis_test_dir,
+            vis_opt,
+            sample,
+        )
+        self.img_names, self.char_bbox, self.img_words = self.load_data()
+        self.vis_index = list(range(1000))
+
+    def load_data(self, bbox="char"):
+
+        gt = scio.loadmat(os.path.join(self.data_dir, "gt.mat"))
+        img_names = gt["imnames"][0]
+        img_words = gt["txt"][0]
+
+        if bbox == "char":
+            img_bbox = gt["charBB"][0]
+        else:
+            img_bbox = gt["wordBB"][0]  # word bbox needed for test
+
+        return img_names, img_bbox, img_words
+
+    def dilate_img_to_output_size(self, image, char_bbox):
+        h, w, _ = image.shape
+        if min(h, w) <= self.output_size:
+            scale = float(self.output_size) / min(h, w)
+        else:
+            scale = 1.0
+        image = cv2.resize(
+            image, dsize=None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC
+        )
+        char_bbox *= scale
+        return image, char_bbox
+
+    def make_gt_score(self, index):
+        img_path = os.path.join(self.data_dir, self.img_names[index][0])
+        image = cv2.imread(img_path, cv2.IMREAD_COLOR)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        all_char_bbox = self.char_bbox[index].transpose(
+            (2, 1, 0)
+        )  # shape : (Number of characters in image, 4, 2)
+
+        img_h, img_w, _ = image.shape
+
+        confidence_mask = np.ones((img_h, img_w), dtype=np.float32)
+
+        words = [
+            re.split(" \n|\n |\n| ", word.strip()) for word in self.img_words[index]
+        ]
+        words = list(itertools.chain(*words))
+        words = [word for word in words if len(word) > 0]
+
+        word_level_char_bbox = []
+        char_idx = 0
+
+        for i in range(len(words)):
+            length_of_word = len(words[i])
+            word_bbox = all_char_bbox[char_idx : char_idx + length_of_word]
+            assert len(word_bbox) == length_of_word
+            char_idx += length_of_word
+            word_bbox = np.array(word_bbox)
+            word_level_char_bbox.append(word_bbox)
+
+        region_score = self.gaussian_builder.generate_region(
+            img_h,
+            img_w,
+            word_level_char_bbox,
+            horizontal_text_bools=[True for _ in range(len(words))],
+        )
+        affinity_score, all_affinity_bbox = self.gaussian_builder.generate_affinity(
+            img_h,
+            img_w,
+            word_level_char_bbox,
+            horizontal_text_bools=[True for _ in range(len(words))],
+        )
+
+        return (
+            image,
+            region_score,
+            affinity_score,
+            confidence_mask,
+            word_level_char_bbox,
+            all_affinity_bbox,
+            words,
+        )
+
+
+class CustomDataset(CraftBaseDataset):
+    def __init__(
+        self,
+        output_size,
+        data_dir,
+        saved_gt_dir,
+        mean,
+        variance,
+        gauss_init_size,
+        gauss_sigma,
+        enlarge_region,
+        enlarge_affinity,
+        aug,
+        vis_test_dir,
+        vis_opt,
+        sample,
+        watershed_param,
+        pseudo_vis_opt,
+        do_not_care_label,
+    ):
+        super().__init__(
+            output_size,
+            data_dir,
+            saved_gt_dir,
+            mean,
+            variance,
+            gauss_init_size,
+            gauss_sigma,
+            enlarge_region,
+            enlarge_affinity,
+            aug,
+            vis_test_dir,
+            vis_opt,
+            sample,
+        )
+        self.pseudo_vis_opt = pseudo_vis_opt
+        self.do_not_care_label = do_not_care_label
+        self.pseudo_charbox_builder = PseudoCharBoxBuilder(
+            watershed_param, vis_test_dir, pseudo_vis_opt, self.gaussian_builder
+        )
+        self.vis_index = list(range(1000))
+        self.img_dir = os.path.join(data_dir, "ch4_training_images")
+        self.img_gt_box_dir = os.path.join(
+            data_dir, "ch4_training_localization_transcription_gt"
+        )
+        self.img_names = os.listdir(self.img_dir)
+
+    def update_model(self, net):
+        self.net = net
+
+    def update_device(self, gpu):
+        self.gpu = gpu
+
+    def load_img_gt_box(self, img_gt_box_path):
+        lines = open(img_gt_box_path, encoding="utf-8").readlines()
+        word_bboxes = []
+        words = []
+        for line in lines:
+            box_info = line.strip().encode("utf-8").decode("utf-8-sig").split(",")
+            box_points = [int(box_info[i]) for i in range(8)]
+            box_points = np.array(box_points, np.float32).reshape(4, 2)
+            word = box_info[8:]
+            word = ",".join(word)
+            if word in self.do_not_care_label:
+                words.append(self.do_not_care_label[0])
+                word_bboxes.append(box_points)
+                continue
+            word_bboxes.append(box_points)
+            words.append(word)
+        return np.array(word_bboxes), words
+
+    def load_data(self, index):
+        img_name = self.img_names[index]
+        img_path = os.path.join(self.img_dir, img_name)
+        image = cv2.imread(img_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        img_gt_box_path = os.path.join(
+            self.img_gt_box_dir, "gt_%s.txt" % os.path.splitext(img_name)[0]
+        )
+        word_bboxes, words = self.load_img_gt_box(
+            img_gt_box_path
+        )  # shape : (Number of word bbox, 4, 2)
+        confidence_mask = np.ones((image.shape[0], image.shape[1]), np.float32)
+
+        word_level_char_bbox = []
+        do_care_words = []
+        horizontal_text_bools = []
+
+        if len(word_bboxes) == 0:
+            return (
+                image,
+                word_level_char_bbox,
+                do_care_words,
+                confidence_mask,
+                horizontal_text_bools,
+            )
+        _word_bboxes = word_bboxes.copy()
+        for i in range(len(word_bboxes)):
+            if words[i] in self.do_not_care_label:
+                cv2.fillPoly(confidence_mask, [np.int32(_word_bboxes[i])], 0)
+                continue
+
+            (
+                pseudo_char_bbox,
+                confidence,
+                horizontal_text_bool,
+            ) = self.pseudo_charbox_builder.build_char_box(
+                self.net, self.gpu, image, word_bboxes[i], words[i], img_name=img_name
+            )
+
+            cv2.fillPoly(confidence_mask, [np.int32(_word_bboxes[i])], confidence)
+            do_care_words.append(words[i])
+            word_level_char_bbox.append(pseudo_char_bbox)
+            horizontal_text_bools.append(horizontal_text_bool)
+
+        return (
+            image,
+            word_level_char_bbox,
+            do_care_words,
+            confidence_mask,
+            horizontal_text_bools,
+        )
+
+    def make_gt_score(self, index):
+        """
+        Make region, affinity scores using pseudo character-level GT bounding box
+        word_level_char_bbox's shape : [word_num, [char_num_in_one_word, 4, 2]]
+        :rtype region_score: np.float32
+        :rtype affinity_score: np.float32
+        :rtype confidence_mask: np.float32
+        :rtype word_level_char_bbox: np.float32
+        :rtype words: list
+        """
+        (
+            image,
+            word_level_char_bbox,
+            words,
+            confidence_mask,
+            horizontal_text_bools,
+        ) = self.load_data(index)
+        img_h, img_w, _ = image.shape
+
+        if len(word_level_char_bbox) == 0:
+            region_score = np.zeros((img_h, img_w), dtype=np.float32)
+            affinity_score = np.zeros((img_h, img_w), dtype=np.float32)
+            all_affinity_bbox = []
+        else:
+            region_score = self.gaussian_builder.generate_region(
+                img_h, img_w, word_level_char_bbox, horizontal_text_bools
+            )
+            affinity_score, all_affinity_bbox = self.gaussian_builder.generate_affinity(
+                img_h, img_w, word_level_char_bbox, horizontal_text_bools
+            )
+
+        return (
+            image,
+            region_score,
+            affinity_score,
+            confidence_mask,
+            word_level_char_bbox,
+            all_affinity_bbox,
+            words,
+        )
+
+    def load_saved_gt_score(self, index):
+        """
+        Load pre-saved official CRAFT model's region, affinity scores to train
+        word_level_char_bbox's shape : [word_num, [char_num_in_one_word, 4, 2]]
+        :rtype region_score: np.float32
+        :rtype affinity_score: np.float32
+        :rtype confidence_mask: np.float32
+        :rtype word_level_char_bbox: np.float32
+        :rtype words: list
+        """
+        img_name = self.img_names[index]
+        img_path = os.path.join(self.img_dir, img_name)
+        image = cv2.imread(img_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        img_gt_box_path = os.path.join(
+            self.img_gt_box_dir, "gt_%s.txt" % os.path.splitext(img_name)[0]
+        )
+        word_bboxes, words = self.load_img_gt_box(img_gt_box_path)
+        image, word_bboxes = rescale(image, word_bboxes)
+        img_h, img_w, _ = image.shape
+
+        query_idx = int(self.img_names[index].split(".")[0].split("_")[1])
+
+        saved_region_scores_path = os.path.join(
+            self.saved_gt_dir, f"res_img_{query_idx}_region.jpg"
+        )
+        saved_affi_scores_path = os.path.join(
+            self.saved_gt_dir, f"res_img_{query_idx}_affi.jpg"
+        )
+        saved_cf_mask_path = os.path.join(
+            self.saved_gt_dir, f"res_img_{query_idx}_cf_mask_thresh_0.6.jpg"
+        )
+        region_score = cv2.imread(saved_region_scores_path, cv2.IMREAD_GRAYSCALE)
+        affinity_score = cv2.imread(saved_affi_scores_path, cv2.IMREAD_GRAYSCALE)
+        confidence_mask = cv2.imread(saved_cf_mask_path, cv2.IMREAD_GRAYSCALE)
+
+        region_score = cv2.resize(region_score, (img_w, img_h))
+        affinity_score = cv2.resize(affinity_score, (img_w, img_h))
+        confidence_mask = cv2.resize(
+            confidence_mask, (img_w, img_h), interpolation=cv2.INTER_NEAREST
+        )
+
+        region_score = region_score.astype(np.float32) / 255
+        affinity_score = affinity_score.astype(np.float32) / 255
+        confidence_mask = confidence_mask.astype(np.float32) / 255
+
+        # NOTE : Even though word_level_char_bbox is not necessary, align bbox format with make_gt_score()
+        word_level_char_bbox = []
+
+        for i in range(len(word_bboxes)):
+            word_level_char_bbox.append(np.expand_dims(word_bboxes[i], 0))
+
+        return (
+            image,
+            region_score,
+            affinity_score,
+            confidence_mask,
+            word_level_char_bbox,
+            words,
+        )

trainer/craft/data/gaussian.py

@@ -0,0 +1,192 @@
+import numpy as np
+import cv2
+
+from data.boxEnlarge import enlargebox
+
+
+class GaussianBuilder(object):
+    def __init__(self, init_size, sigma, enlarge_region, enlarge_affinity):
+        self.init_size = init_size
+        self.sigma = sigma
+        self.enlarge_region = enlarge_region
+        self.enlarge_affinity = enlarge_affinity
+        self.gaussian_map, self.gaussian_map_color = self.generate_gaussian_map()
+
+    def generate_gaussian_map(self):
+        circle_mask = self.generate_circle_mask()
+
+        gaussian_map = np.zeros((self.init_size, self.init_size), np.float32)
+
+        for i in range(self.init_size):
+            for j in range(self.init_size):
+                gaussian_map[i, j] = (
+                    1
+                    / 2
+                    / np.pi
+                    / (self.sigma ** 2)
+                    * np.exp(
+                        -1
+                        / 2
+                        * (
+                            (i - self.init_size / 2) ** 2 / (self.sigma ** 2)
+                            + (j - self.init_size / 2) ** 2 / (self.sigma ** 2)
+                        )
+                    )
+                )
+
+        gaussian_map = gaussian_map * circle_mask
+        gaussian_map = (gaussian_map / np.max(gaussian_map)).astype(np.float32)
+
+        gaussian_map_color = (gaussian_map * 255).astype(np.uint8)
+        gaussian_map_color = cv2.applyColorMap(gaussian_map_color, cv2.COLORMAP_JET)
+        return gaussian_map, gaussian_map_color
+
+    def generate_circle_mask(self):
+
+        zero_arr = np.zeros((self.init_size, self.init_size), np.float32)
+        circle_mask = cv2.circle(
+            img=zero_arr,
+            center=(self.init_size // 2, self.init_size // 2),
+            radius=self.init_size // 2,
+            color=1,
+            thickness=-1,
+        )
+
+        return circle_mask
+
+    def four_point_transform(self, bbox):
+        """
+        Using the bbox, standard 2D gaussian map, returns Transformed 2d Gaussian map
+        """
+        width, height = (
+            np.max(bbox[:, 0]).astype(np.int32),
+            np.max(bbox[:, 1]).astype(np.int32),
+        )
+        init_points = np.array(
+            [
+                [0, 0],
+                [self.init_size, 0],
+                [self.init_size, self.init_size],
+                [0, self.init_size],
+            ],
+            dtype="float32",
+        )
+
+        M = cv2.getPerspectiveTransform(init_points, bbox)
+        warped_gaussian_map = cv2.warpPerspective(self.gaussian_map, M, (width, height))
+        return warped_gaussian_map, width, height
+
+    def add_gaussian_map_to_score_map(
+        self, score_map, bbox, enlarge_size, horizontal_text_bool, map_type=None
+    ):
+        """
+        Mapping 2D Gaussian to the character box coordinates of the score_map.
+
+        :param score_map: Target map to put 2D gaussian on character box
+        :type score_map: np.float32
+        :param bbox: character boxes
+        :type bbox: np.float32
+        :param enlarge_size: Enlarge size of gaussian map to fit character shape
+        :type enlarge_size: list of enlarge size [x dim, y dim]
+        :param horizontal_text_bool: Flag that bbox is horizontal text or not
+        :type horizontal_text_bool: bool
+        :param map_type: Whether map's type is "region" | "affinity"
+        :type map_type: str
+        :return score_map: score map that all 2D gaussian put on character box
+        :rtype: np.float32
+        """
+
+        map_h, map_w = score_map.shape
+        bbox = enlargebox(bbox, map_h, map_w, enlarge_size, horizontal_text_bool)
+
+        # If any one point of character bbox is out of range, don't put in on map
+        if np.any(bbox < 0) or np.any(bbox[:, 0] > map_w) or np.any(bbox[:, 1] > map_h):
+            return score_map
+
+        bbox_left, bbox_top = np.array([np.min(bbox[:, 0]), np.min(bbox[:, 1])]).astype(
+            np.int32
+        )
+        bbox -= (bbox_left, bbox_top)
+        warped_gaussian_map, width, height = self.four_point_transform(
+            bbox.astype(np.float32)
+        )
+
+        try:
+            bbox_area_of_image = score_map[
+                bbox_top : bbox_top + height, bbox_left : bbox_left + width,
+            ]
+            high_value_score = np.where(
+                warped_gaussian_map > bbox_area_of_image,
+                warped_gaussian_map,
+                bbox_area_of_image,
+            )
+            score_map[
+                bbox_top : bbox_top + height, bbox_left : bbox_left + width,
+            ] = high_value_score
+
+        except Exception as e:
+            print("Error : {}".format(e))
+            print(
+                "On generating {} map, strange box came out. (width: {}, height: {})".format(
+                    map_type, width, height
+                )
+            )
+
+        return score_map
+
+    def calculate_affinity_box_points(self, bbox_1, bbox_2, vertical=False):
+        center_1, center_2 = np.mean(bbox_1, axis=0), np.mean(bbox_2, axis=0)
+        if vertical:
+            tl = (bbox_1[0] + bbox_1[-1] + center_1) / 3
+            tr = (bbox_1[1:3].sum(0) + center_1) / 3
+            br = (bbox_2[1:3].sum(0) + center_2) / 3
+            bl = (bbox_2[0] + bbox_2[-1] + center_2) / 3
+        else:
+            tl = (bbox_1[0:2].sum(0) + center_1) / 3
+            tr = (bbox_2[0:2].sum(0) + center_2) / 3
+            br = (bbox_2[2:4].sum(0) + center_2) / 3
+            bl = (bbox_1[2:4].sum(0) + center_1) / 3
+        affinity_box = np.array([tl, tr, br, bl]).astype(np.float32)
+        return affinity_box
+
+    def generate_region(
+        self, img_h, img_w, word_level_char_bbox, horizontal_text_bools
+    ):
+        region_map = np.zeros([img_h, img_w], dtype=np.float32)
+        for i in range(
+            len(word_level_char_bbox)
+        ):  # shape : [word_num, [char_num_in_one_word, 4, 2]]
+            for j in range(len(word_level_char_bbox[i])):
+                region_map = self.add_gaussian_map_to_score_map(
+                    region_map,
+                    word_level_char_bbox[i][j].copy(),
+                    self.enlarge_region,
+                    horizontal_text_bools[i],
+                    map_type="region",
+                )
+        return region_map
+
+    def generate_affinity(
+        self, img_h, img_w, word_level_char_bbox, horizontal_text_bools
+    ):
+
+        affinity_map = np.zeros([img_h, img_w], dtype=np.float32)
+        all_affinity_bbox = []
+        for i in range(len(word_level_char_bbox)):
+            for j in range(len(word_level_char_bbox[i]) - 1):
+                affinity_bbox = self.calculate_affinity_box_points(
+                    word_level_char_bbox[i][j], word_level_char_bbox[i][j + 1]
+                )
+
+                affinity_map = self.add_gaussian_map_to_score_map(
+                    affinity_map,
+                    affinity_bbox.copy(),
+                    self.enlarge_affinity,
+                    horizontal_text_bools[i],
+                    map_type="affinity",
+                )
+                all_affinity_bbox.append(np.expand_dims(affinity_bbox, axis=0))
+
+        if len(all_affinity_bbox) > 0:
+            all_affinity_bbox = np.concatenate(all_affinity_bbox, axis=0)
+        return affinity_map, all_affinity_bbox

trainer/craft/data/imgaug.py

@@ -0,0 +1,175 @@
+import random
+
+import cv2
+import numpy as np
+from PIL import Image
+from torchvision.transforms.functional import resized_crop, crop
+from torchvision.transforms import RandomResizedCrop, RandomCrop
+from torchvision.transforms import InterpolationMode
+
+
+def rescale(img, bboxes, target_size=2240):
+    h, w = img.shape[0:2]
+    scale = target_size / max(h, w)
+    img = cv2.resize(img, dsize=None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
+    bboxes = bboxes * scale
+    return img, bboxes
+
+
+def random_resize_crop_synth(augment_targets, size):
+    image, region_score, affinity_score, confidence_mask = augment_targets
+
+    image = Image.fromarray(image)
+    region_score = Image.fromarray(region_score)
+    affinity_score = Image.fromarray(affinity_score)
+    confidence_mask = Image.fromarray(confidence_mask)
+
+    short_side = min(image.size)
+    i, j, h, w = RandomCrop.get_params(image, output_size=(short_side, short_side))
+
+    image = resized_crop(
+        image, i, j, h, w, size=(size, size), interpolation=InterpolationMode.BICUBIC
+    )
+    region_score = resized_crop(
+        region_score, i, j, h, w, (size, size), interpolation=InterpolationMode.BICUBIC
+    )
+    affinity_score = resized_crop(
+        affinity_score,
+        i,
+        j,
+        h,
+        w,
+        (size, size),
+        interpolation=InterpolationMode.BICUBIC,
+    )
+    confidence_mask = resized_crop(
+        confidence_mask,
+        i,
+        j,
+        h,
+        w,
+        (size, size),
+        interpolation=InterpolationMode.NEAREST,
+    )
+
+    image = np.array(image)
+    region_score = np.array(region_score)
+    affinity_score = np.array(affinity_score)
+    confidence_mask = np.array(confidence_mask)
+    augment_targets = [image, region_score, affinity_score, confidence_mask]
+
+    return augment_targets
+
+
+def random_resize_crop(
+    augment_targets, scale, ratio, size, threshold, pre_crop_area=None
+):
+    image, region_score, affinity_score, confidence_mask = augment_targets
+
+    image = Image.fromarray(image)
+    region_score = Image.fromarray(region_score)
+    affinity_score = Image.fromarray(affinity_score)
+    confidence_mask = Image.fromarray(confidence_mask)
+
+    if pre_crop_area != None:
+        i, j, h, w = pre_crop_area
+
+    else:
+        if random.random() < threshold:
+            i, j, h, w = RandomResizedCrop.get_params(image, scale=scale, ratio=ratio)
+        else:
+            i, j, h, w = RandomResizedCrop.get_params(
+                image, scale=(1.0, 1.0), ratio=(1.0, 1.0)
+            )
+
+    image = resized_crop(
+        image, i, j, h, w, size=(size, size), interpolation=InterpolationMode.BICUBIC
+    )
+    region_score = resized_crop(
+        region_score, i, j, h, w, (size, size), interpolation=InterpolationMode.BICUBIC
+    )
+    affinity_score = resized_crop(
+        affinity_score,
+        i,
+        j,
+        h,
+        w,
+        (size, size),
+        interpolation=InterpolationMode.BICUBIC,
+    )
+    confidence_mask = resized_crop(
+        confidence_mask,
+        i,
+        j,
+        h,
+        w,
+        (size, size),
+        interpolation=InterpolationMode.NEAREST,
+    )
+
+    image = np.array(image)
+    region_score = np.array(region_score)
+    affinity_score = np.array(affinity_score)
+    confidence_mask = np.array(confidence_mask)
+    augment_targets = [image, region_score, affinity_score, confidence_mask]
+
+    return augment_targets
+
+
+def random_crop(augment_targets, size):
+    image, region_score, affinity_score, confidence_mask = augment_targets
+
+    image = Image.fromarray(image)
+    region_score = Image.fromarray(region_score)
+    affinity_score = Image.fromarray(affinity_score)
+    confidence_mask = Image.fromarray(confidence_mask)
+
+    i, j, h, w = RandomCrop.get_params(image, output_size=(size, size))
+
+    image = crop(image, i, j, h, w)
+    region_score = crop(region_score, i, j, h, w)
+    affinity_score = crop(affinity_score, i, j, h, w)
+    confidence_mask = crop(confidence_mask, i, j, h, w)
+
+    image = np.array(image)
+    region_score = np.array(region_score)
+    affinity_score = np.array(affinity_score)
+    confidence_mask = np.array(confidence_mask)
+    augment_targets = [image, region_score, affinity_score, confidence_mask]
+
+    return augment_targets
+
+
+def random_horizontal_flip(imgs):
+    if random.random() < 0.5:
+        for i in range(len(imgs)):
+            imgs[i] = np.flip(imgs[i], axis=1).copy()
+    return imgs
+
+
+def random_scale(images, word_level_char_bbox, scale_range):
+    scale = random.sample(scale_range, 1)[0]
+
+    for i in range(len(images)):
+        images[i] = cv2.resize(images[i], dsize=None, fx=scale, fy=scale)
+
+    for i in range(len(word_level_char_bbox)):
+        word_level_char_bbox[i] *= scale
+
+    return images
+
+
+def random_rotate(images, max_angle):
+    angle = random.random() * 2 * max_angle - max_angle
+    for i in range(len(images)):
+        img = images[i]
+        w, h = img.shape[:2]
+        rotation_matrix = cv2.getRotationMatrix2D((h / 2, w / 2), angle, 1)
+        if i == len(images) - 1:
+            img_rotation = cv2.warpAffine(
+                img, M=rotation_matrix, dsize=(h, w), flags=cv2.INTER_NEAREST
+            )
+        else:
+            img_rotation = cv2.warpAffine(img, rotation_matrix, (h, w))
+        images[i] = img_rotation
+    return images

trainer/craft/data/imgproc.py

@@ -0,0 +1,91 @@
+"""  
+Copyright (c) 2019-present NAVER Corp.
+MIT License
+"""
+
+# -*- coding: utf-8 -*-
+import numpy as np
+
+import cv2
+from skimage import io
+
+
+def loadImage(img_file):
+    img = io.imread(img_file)  # RGB order
+    if img.shape[0] == 2:
+        img = img[0]
+    if len(img.shape) == 2:
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+    if img.shape[2] == 4:
+        img = img[:, :, :3]
+    img = np.array(img)
+
+    return img
+
+
+def normalizeMeanVariance(
+    in_img, mean=(0.485, 0.456, 0.406), variance=(0.229, 0.224, 0.225)
+):
+    # should be RGB order
+    img = in_img.copy().astype(np.float32)
+
+    img -= np.array(
+        [mean[0] * 255.0, mean[1] * 255.0, mean[2] * 255.0], dtype=np.float32
+    )
+    img /= np.array(
+        [variance[0] * 255.0, variance[1] * 255.0, variance[2] * 255.0],
+        dtype=np.float32,
+    )
+    return img
+
+
+def denormalizeMeanVariance(
+    in_img, mean=(0.485, 0.456, 0.406), variance=(0.229, 0.224, 0.225)
+):
+    # should be RGB order
+    img = in_img.copy()
+    img *= variance
+    img += mean
+    img *= 255.0
+    img = np.clip(img, 0, 255).astype(np.uint8)
+    return img
+
+
+def resize_aspect_ratio(img, square_size, interpolation, mag_ratio=1):
+    height, width, channel = img.shape
+
+    # magnify image size
+    target_size = mag_ratio * max(height, width)
+
+    # set original image size
+    if target_size > square_size:
+        target_size = square_size
+
+    ratio = target_size / max(height, width)
+
+    target_h, target_w = int(height * ratio), int(width * ratio)
+
+    # NOTE
+    valid_size_heatmap = (int(target_h / 2), int(target_w / 2))
+
+    proc = cv2.resize(img, (target_w, target_h), interpolation=interpolation)
+
+    # make canvas and paste image
+    target_h32, target_w32 = target_h, target_w
+    if target_h % 32 != 0:
+        target_h32 = target_h + (32 - target_h % 32)
+    if target_w % 32 != 0:
+        target_w32 = target_w + (32 - target_w % 32)
+    resized = np.zeros((target_h32, target_w32, channel), dtype=np.float32)
+    resized[0:target_h, 0:target_w, :] = proc
+
+    # target_h, target_w = target_h32, target_w32
+    # size_heatmap = (int(target_w/2), int(target_h/2))
+
+    return resized, ratio, valid_size_heatmap
+
+
+def cvt2HeatmapImg(img):
+    img = (np.clip(img, 0, 1) * 255).astype(np.uint8)
+    img = cv2.applyColorMap(img, cv2.COLORMAP_JET)
+    return img

trainer/craft/data/pseudo_label/make_charbox.py

@@ -0,0 +1,263 @@
+import os
+import random
+import math
+
+import numpy as np
+import cv2
+import torch
+
+from data import imgproc
+from data.pseudo_label.watershed import exec_watershed_by_version
+
+
+class PseudoCharBoxBuilder:
+    def __init__(self, watershed_param, vis_test_dir, pseudo_vis_opt, gaussian_builder):
+        self.watershed_param = watershed_param
+        self.vis_test_dir = vis_test_dir
+        self.pseudo_vis_opt = pseudo_vis_opt
+        self.gaussian_builder = gaussian_builder
+        self.cnt = 0
+        self.flag = False
+
+    def crop_image_by_bbox(self, image, box, word):
+        w = max(
+            int(np.linalg.norm(box[0] - box[1])), int(np.linalg.norm(box[2] - box[3]))
+        )
+        h = max(
+            int(np.linalg.norm(box[0] - box[3])), int(np.linalg.norm(box[1] - box[2]))
+        )
+        try:
+            word_ratio = h / w
+        except:
+            import ipdb
+
+            ipdb.set_trace()
+
+        one_char_ratio = min(h, w) / (max(h, w) / len(word))
+
+        # NOTE: criterion to split vertical word in here is set to work properly on IC15 dataset
+        if word_ratio > 2 or (word_ratio > 1.6 and one_char_ratio > 2.4):
+            # warping method of vertical word (classified by upper condition)
+            horizontal_text_bool = False
+            long_side = h
+            short_side = w
+            M = cv2.getPerspectiveTransform(
+                np.float32(box),
+                np.float32(
+                    np.array(
+                        [
+                            [long_side, 0],
+                            [long_side, short_side],
+                            [0, short_side],
+                            [0, 0],
+                        ]
+                    )
+                ),
+            )
+            self.flag = True
+        else:
+            # warping method of horizontal word
+            horizontal_text_bool = True
+            long_side = w
+            short_side = h
+            M = cv2.getPerspectiveTransform(
+                np.float32(box),
+                np.float32(
+                    np.array(
+                        [
+                            [0, 0],
+                            [long_side, 0],
+                            [long_side, short_side],
+                            [0, short_side],
+                        ]
+                    )
+                ),
+            )
+            self.flag = False
+
+        warped = cv2.warpPerspective(image, M, (long_side, short_side))
+        return warped, M, horizontal_text_bool
+
+    def inference_word_box(self, net, gpu, word_image):
+        if net.training:
+            net.eval()
+
+        with torch.no_grad():
+            word_img_torch = torch.from_numpy(
+                imgproc.normalizeMeanVariance(
+                    word_image,
+                    mean=(0.485, 0.456, 0.406),
+                    variance=(0.229, 0.224, 0.225),
+                )
+            )
+            word_img_torch = word_img_torch.permute(2, 0, 1).unsqueeze(0)
+            word_img_torch = word_img_torch.type(torch.FloatTensor).cuda(gpu)
+            with torch.cuda.amp.autocast():
+                word_img_scores, _ = net(word_img_torch)
+        return word_img_scores
+
+    def visualize_pseudo_label(
+        self, word_image, region_score, watershed_box, pseudo_char_bbox, img_name,
+    ):
+        word_img_h, word_img_w, _ = word_image.shape
+        word_img_cp1 = word_image.copy()
+        word_img_cp2 = word_image.copy()
+        _watershed_box = np.int32(watershed_box)
+        _pseudo_char_bbox = np.int32(pseudo_char_bbox)
+
+        region_score_color = cv2.applyColorMap(np.uint8(region_score), cv2.COLORMAP_JET)
+        region_score_color = cv2.resize(region_score_color, (word_img_w, word_img_h))
+
+        for box in _watershed_box:
+            cv2.polylines(
+                np.uint8(word_img_cp1),
+                [np.reshape(box, (-1, 1, 2))],
+                True,
+                (255, 0, 0),
+            )
+
+        for box in _pseudo_char_bbox:
+            cv2.polylines(
+                np.uint8(word_img_cp2), [np.reshape(box, (-1, 1, 2))], True, (255, 0, 0)
+            )
+
+        # NOTE: Just for visualize, put gaussian map on char box
+        pseudo_gt_region_score = self.gaussian_builder.generate_region(
+            word_img_h, word_img_w, [_pseudo_char_bbox], [True]
+        )
+
+        pseudo_gt_region_score = cv2.applyColorMap(
+            (pseudo_gt_region_score * 255).astype("uint8"), cv2.COLORMAP_JET
+        )
+
+        overlay_img = cv2.addWeighted(
+            word_image[:, :, ::-1], 0.7, pseudo_gt_region_score, 0.3, 5
+        )
+        vis_result = np.hstack(
+            [
+                word_image[:, :, ::-1],
+                region_score_color,
+                word_img_cp1[:, :, ::-1],
+                word_img_cp2[:, :, ::-1],
+                pseudo_gt_region_score,
+                overlay_img,
+            ]
+        )
+
+        if not os.path.exists(os.path.dirname(self.vis_test_dir)):
+            os.makedirs(os.path.dirname(self.vis_test_dir))
+        cv2.imwrite(
+            os.path.join(
+                self.vis_test_dir,
+                "{}_{}".format(
+                    img_name, f"pseudo_char_bbox_{random.randint(0,100)}.jpg"
+                ),
+            ),
+            vis_result,
+        )
+
+    def clip_into_boundary(self, box, bound):
+        if len(box) == 0:
+            return box
+        else:
+            box[:, :, 0] = np.clip(box[:, :, 0], 0, bound[1])
+            box[:, :, 1] = np.clip(box[:, :, 1], 0, bound[0])
+            return box
+
+    def get_confidence(self, real_len, pseudo_len):
+        if pseudo_len == 0:
+            return 0.0
+        return (real_len - min(real_len, abs(real_len - pseudo_len))) / real_len
+
+    def split_word_equal_gap(self, word_img_w, word_img_h, word):
+        width = word_img_w
+        height = word_img_h
+
+        width_per_char = width / len(word)
+        bboxes = []
+        for j, char in enumerate(word):
+            if char == " ":
+                continue
+            left = j * width_per_char
+            right = (j + 1) * width_per_char
+            bbox = np.array([[left, 0], [right, 0], [right, height], [left, height]])
+            bboxes.append(bbox)
+
+        bboxes = np.array(bboxes, np.float32)
+        return bboxes
+
+    def cal_angle(self, v1):
+        theta = np.arccos(min(1, v1[0] / (np.linalg.norm(v1) + 10e-8)))
+        return 2 * math.pi - theta if v1[1] < 0 else theta
+
+    def clockwise_sort(self, points):
+        # returns 4x2 [[x1,y1],[x2,y2],[x3,y3],[x4,y4]] ndarray
+        v1, v2, v3, v4 = points
+        center = (v1 + v2 + v3 + v4) / 4
+        theta = np.array(
+            [
+                self.cal_angle(v1 - center),
+                self.cal_angle(v2 - center),
+                self.cal_angle(v3 - center),
+                self.cal_angle(v4 - center),
+            ]
+        )
+        index = np.argsort(theta)
+        return np.array([v1, v2, v3, v4])[index, :]
+
+    def build_char_box(self, net, gpu, image, word_bbox, word, img_name=""):
+        word_image, M, horizontal_text_bool = self.crop_image_by_bbox(
+            image, word_bbox, word
+        )
+        real_word_without_space = word.replace("\s", "")
+        real_char_len = len(real_word_without_space)
+
+        scale = 128.0 / word_image.shape[0]
+
+        word_image = cv2.resize(word_image, None, fx=scale, fy=scale)
+        word_img_h, word_img_w, _ = word_image.shape
+
+        scores = self.inference_word_box(net, gpu, word_image)
+        region_score = scores[0, :, :, 0].cpu().data.numpy()
+        region_score = np.uint8(np.clip(region_score, 0, 1) * 255)
+
+        region_score_rgb = cv2.resize(region_score, (word_img_w, word_img_h))
+        region_score_rgb = cv2.cvtColor(region_score_rgb, cv2.COLOR_GRAY2RGB)
+
+        pseudo_char_bbox = exec_watershed_by_version(
+            self.watershed_param, region_score, word_image, self.pseudo_vis_opt
+        )
+
+        # Used for visualize only
+        watershed_box = pseudo_char_bbox.copy()
+
+        pseudo_char_bbox = self.clip_into_boundary(
+            pseudo_char_bbox, region_score_rgb.shape
+        )
+
+        confidence = self.get_confidence(real_char_len, len(pseudo_char_bbox))
+
+        if confidence <= 0.5:
+            pseudo_char_bbox = self.split_word_equal_gap(word_img_w, word_img_h, word)
+            confidence = 0.5
+
+        if self.pseudo_vis_opt and self.flag:
+            self.visualize_pseudo_label(
+                word_image, region_score, watershed_box, pseudo_char_bbox, img_name,
+            )
+
+        if len(pseudo_char_bbox) != 0:
+            index = np.argsort(pseudo_char_bbox[:, 0, 0])
+            pseudo_char_bbox = pseudo_char_bbox[index]
+
+        pseudo_char_bbox /= scale
+
+        M_inv = np.linalg.pinv(M)
+        for i in range(len(pseudo_char_bbox)):
+            pseudo_char_bbox[i] = cv2.perspectiveTransform(
+                pseudo_char_bbox[i][None, :, :], M_inv
+            )
+
+        pseudo_char_bbox = self.clip_into_boundary(pseudo_char_bbox, image.shape)
+
+        return pseudo_char_bbox, confidence, horizontal_text_bool

trainer/craft/data/pseudo_label/watershed.py

@@ -0,0 +1,45 @@
+import cv2
+import numpy as np
+from skimage.segmentation import watershed
+
+
+def segment_region_score(watershed_param, region_score, word_image, pseudo_vis_opt):
+    region_score = np.float32(region_score) / 255
+    fore = np.uint8(region_score > 0.75)
+    back = np.uint8(region_score < 0.05)
+    unknown = 1 - (fore + back)
+    ret, markers = cv2.connectedComponents(fore)
+    markers += 1
+    markers[unknown == 1] = 0
+
+    labels = watershed(-region_score, markers)
+    boxes = []
+    for label in range(2, ret + 1):
+        y, x = np.where(labels == label)
+        x_max = x.max()
+        y_max = y.max()
+        x_min = x.min()
+        y_min = y.min()
+        box = [[x_min, y_min], [x_max, y_min], [x_max, y_max], [x_min, y_max]]
+        box = np.array(box)
+        box *= 2
+        boxes.append(box)
+    return np.array(boxes, dtype=np.float32)
+
+
+def exec_watershed_by_version(
+    watershed_param, region_score, word_image, pseudo_vis_opt
+):
+
+    func_name_map_dict = {
+        "skimage": segment_region_score,
+    }
+
+    try:
+        return func_name_map_dict[watershed_param.version](
+            watershed_param, region_score, word_image, pseudo_vis_opt
+        )
+    except:
+        print(
+            f"Watershed version {watershed_param.version} does not exist in func_name_map_dict."
+        )

trainer/craft/data_root_dir/folder.txt

@@ -0,0 +1 @@
+place dataset folder here

trainer/craft/eval.py

@@ -0,0 +1,381 @@
+# -*- coding: utf-8 -*-
+
+import argparse
+import os
+
+import cv2
+import numpy as np
+import torch
+import torch.backends.cudnn as cudnn
+from tqdm import tqdm
+import wandb
+
+from config.load_config import load_yaml, DotDict
+from model.craft import CRAFT
+from metrics.eval_det_iou import DetectionIoUEvaluator
+from utils.inference_boxes import (
+    test_net,
+    load_icdar2015_gt,
+    load_icdar2013_gt,
+    load_synthtext_gt,
+)
+from utils.util import copyStateDict
+
+
+
+def save_result_synth(img_file, img, pre_output, pre_box, gt_box=None, result_dir=""):
+
+    img = np.array(img)
+    img_copy = img.copy()
+    region = pre_output[0]
+    affinity = pre_output[1]
+
+    # make result file list
+    filename, file_ext = os.path.splitext(os.path.basename(img_file))
+
+    # draw bounding boxes for prediction, color green
+    for i, box in enumerate(pre_box):
+        poly = np.array(box).astype(np.int32).reshape((-1))
+        poly = poly.reshape(-1, 2)
+        try:
+            cv2.polylines(
+                img, [poly.reshape((-1, 1, 2))], True, color=(0, 255, 0), thickness=2
+            )
+        except:
+            pass
+
+    # draw bounding boxes for gt, color red
+    if gt_box is not None:
+        for j in range(len(gt_box)):
+            cv2.polylines(
+                img,
+                [np.array(gt_box[j]["points"]).astype(np.int32).reshape((-1, 1, 2))],
+                True,
+                color=(0, 0, 255),
+                thickness=2,
+            )
+
+    # draw overlay image
+    overlay_img = overlay(img_copy, region, affinity, pre_box)
+
+    # Save result image
+    res_img_path = result_dir + "/res_" + filename + ".jpg"
+    cv2.imwrite(res_img_path, img)
+
+    overlay_image_path = result_dir + "/res_" + filename + "_box.jpg"
+    cv2.imwrite(overlay_image_path, overlay_img)
+
+
+def save_result_2015(img_file, img, pre_output, pre_box, gt_box, result_dir):
+
+    img = np.array(img)
+    img_copy = img.copy()
+    region = pre_output[0]
+    affinity = pre_output[1]
+
+    # make result file list
+    filename, file_ext = os.path.splitext(os.path.basename(img_file))
+
+    for i, box in enumerate(pre_box):
+        poly = np.array(box).astype(np.int32).reshape((-1))
+        poly = poly.reshape(-1, 2)
+        try:
+            cv2.polylines(
+                img, [poly.reshape((-1, 1, 2))], True, color=(0, 255, 0), thickness=2
+            )
+        except:
+            pass
+
+    if gt_box is not None:
+        for j in range(len(gt_box)):
+            _gt_box = np.array(gt_box[j]["points"]).reshape(-1, 2).astype(np.int32)
+            if gt_box[j]["text"] == "###":
+                cv2.polylines(img, [_gt_box], True, color=(128, 128, 128), thickness=2)
+            else:
+                cv2.polylines(img, [_gt_box], True, color=(0, 0, 255), thickness=2)
+
+    # draw overlay image
+    overlay_img = overlay(img_copy, region, affinity, pre_box)
+
+    # Save result image
+    res_img_path = result_dir + "/res_" + filename + ".jpg"
+    cv2.imwrite(res_img_path, img)
+
+    overlay_image_path = result_dir + "/res_" + filename + "_box.jpg"
+    cv2.imwrite(overlay_image_path, overlay_img)
+
+
+def save_result_2013(img_file, img, pre_output, pre_box, gt_box=None, result_dir=""):
+
+    img = np.array(img)
+    img_copy = img.copy()
+    region = pre_output[0]
+    affinity = pre_output[1]
+
+    # make result file list
+    filename, file_ext = os.path.splitext(os.path.basename(img_file))
+
+    # draw bounding boxes for prediction, color green
+    for i, box in enumerate(pre_box):
+        poly = np.array(box).astype(np.int32).reshape((-1))
+        poly = poly.reshape(-1, 2)
+        try:
+            cv2.polylines(
+                img, [poly.reshape((-1, 1, 2))], True, color=(0, 255, 0), thickness=2
+            )
+        except:
+            pass
+
+    # draw bounding boxes for gt, color red
+    if gt_box is not None:
+        for j in range(len(gt_box)):
+            cv2.polylines(
+                img,
+                [np.array(gt_box[j]["points"]).reshape((-1, 1, 2))],
+                True,
+                color=(0, 0, 255),
+                thickness=2,
+            )
+
+    # draw overlay image
+    overlay_img = overlay(img_copy, region, affinity, pre_box)
+
+    # Save result image
+    res_img_path = result_dir + "/res_" + filename + ".jpg"
+    cv2.imwrite(res_img_path, img)
+
+    overlay_image_path = result_dir + "/res_" + filename + "_box.jpg"
+    cv2.imwrite(overlay_image_path, overlay_img)
+
+
+def overlay(image, region, affinity, single_img_bbox):
+
+    height, width, channel = image.shape
+
+    region_score = cv2.resize(region, (width, height))
+    affinity_score = cv2.resize(affinity, (width, height))
+
+    overlay_region = cv2.addWeighted(image.copy(), 0.4, region_score, 0.6, 5)
+    overlay_aff = cv2.addWeighted(image.copy(), 0.4, affinity_score, 0.6, 5)
+
+    boxed_img = image.copy()
+    for word_box in single_img_bbox:
+        cv2.polylines(
+            boxed_img,
+            [word_box.astype(np.int32).reshape((-1, 1, 2))],
+            True,
+            color=(0, 255, 0),
+            thickness=3,
+        )
+
+    temp1 = np.hstack([image, boxed_img])
+    temp2 = np.hstack([overlay_region, overlay_aff])
+    temp3 = np.vstack([temp1, temp2])
+
+    return temp3
+
+
+def load_test_dataset_iou(test_folder_name, config):
+
+    if test_folder_name == "synthtext":
+        total_bboxes_gt, total_img_path = load_synthtext_gt(config.test_data_dir)
+
+    elif test_folder_name == "icdar2013":
+        total_bboxes_gt, total_img_path = load_icdar2013_gt(
+            dataFolder=config.test_data_dir
+        )
+
+    elif test_folder_name == "icdar2015":
+        total_bboxes_gt, total_img_path = load_icdar2015_gt(
+            dataFolder=config.test_data_dir
+        )
+
+    elif test_folder_name == "custom_data":
+        total_bboxes_gt, total_img_path = load_icdar2015_gt(
+            dataFolder=config.test_data_dir
+        )
+
+    else:
+        print("not found test dataset")
+        return None, None
+
+    return total_bboxes_gt, total_img_path
+
+
+def viz_test(img, pre_output, pre_box, gt_box, img_name, result_dir, test_folder_name):
+
+    if test_folder_name == "synthtext":
+        save_result_synth(
+            img_name, img[:, :, ::-1].copy(), pre_output, pre_box, gt_box, result_dir
+        )
+    elif test_folder_name == "icdar2013":
+        save_result_2013(
+            img_name, img[:, :, ::-1].copy(), pre_output, pre_box, gt_box, result_dir
+        )
+    elif test_folder_name == "icdar2015":
+        save_result_2015(
+            img_name, img[:, :, ::-1].copy(), pre_output, pre_box, gt_box, result_dir
+        )
+    elif test_folder_name == "custom_data":
+        save_result_2015(
+            img_name, img[:, :, ::-1].copy(), pre_output, pre_box, gt_box, result_dir
+        )
+    else:
+        print("not found test dataset")
+
+
+def main_eval(model_path, backbone, config, evaluator, result_dir, buffer, model, mode):
+
+    if not os.path.exists(result_dir):
+        os.makedirs(result_dir, exist_ok=True)
+
+    total_imgs_bboxes_gt, total_imgs_path = load_test_dataset_iou("custom_data", config)
+
+    if mode == "weak_supervision" and torch.cuda.device_count() != 1:
+        gpu_count = torch.cuda.device_count() // 2
+    else:
+        gpu_count = torch.cuda.device_count()
+    gpu_idx = torch.cuda.current_device()
+    torch.cuda.set_device(gpu_idx)
+
+    # Only evaluation time
+    if model is None:
+        piece_imgs_path = total_imgs_path
+
+        if backbone == "vgg":
+            model = CRAFT()
+        else:
+            raise Exception("Undefined architecture")
+
+        print("Loading weights from checkpoint (" + model_path + ")")
+        net_param = torch.load(model_path, map_location=f"cuda:{gpu_idx}")
+        model.load_state_dict(copyStateDict(net_param["craft"]))
+
+        if config.cuda:
+            model = model.cuda()
+            cudnn.benchmark = False
+
+    # Distributed evaluation in the middle of training time
+    else:
+        if buffer is not None:
+            # check all buffer value is None for distributed evaluation
+            assert all(
+                v is None for v in buffer
+            ), "Buffer already filled with another value."
+        slice_idx = len(total_imgs_bboxes_gt) // gpu_count
+
+        # last gpu
+        if gpu_idx == gpu_count - 1:
+            piece_imgs_path = total_imgs_path[gpu_idx * slice_idx :]
+            # piece_imgs_bboxes_gt = total_imgs_bboxes_gt[gpu_idx * slice_idx:]
+        else:
+            piece_imgs_path = total_imgs_path[
+                gpu_idx * slice_idx : (gpu_idx + 1) * slice_idx
+            ]
+            # piece_imgs_bboxes_gt = total_imgs_bboxes_gt[gpu_idx * slice_idx: (gpu_idx + 1) * slice_idx]
+
+    model.eval()
+
+    # -----------------------------------------------------------------------------------------------------------------#
+    total_imgs_bboxes_pre = []
+    for k, img_path in enumerate(tqdm(piece_imgs_path)):
+        image = cv2.imread(img_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        single_img_bbox = []
+        bboxes, polys, score_text = test_net(
+            model,
+            image,
+            config.text_threshold,
+            config.link_threshold,
+            config.low_text,
+            config.cuda,
+            config.poly,
+            config.canvas_size,
+            config.mag_ratio,
+        )
+
+        for box in bboxes:
+            box_info = {"points": box, "text": "###", "ignore": False}
+            single_img_bbox.append(box_info)
+        total_imgs_bboxes_pre.append(single_img_bbox)
+        # Distributed evaluation -------------------------------------------------------------------------------------#
+        if buffer is not None:
+            buffer[gpu_idx * slice_idx + k] = single_img_bbox
+        # print(sum([element is not None for element in buffer]))
+        # -------------------------------------------------------------------------------------------------------------#
+
+        if config.vis_opt:
+            viz_test(
+                image,
+                score_text,
+                pre_box=polys,
+                gt_box=total_imgs_bboxes_gt[k],
+                img_name=img_path,
+                result_dir=result_dir,
+                test_folder_name="custom_data",
+            )
+
+    # When distributed evaluation mode, wait until buffer is full filled
+    if buffer is not None:
+        while None in buffer:
+            continue
+        assert all(v is not None for v in buffer), "Buffer not filled"
+        total_imgs_bboxes_pre = buffer
+
+    results = []
+    for i, (gt, pred) in enumerate(zip(total_imgs_bboxes_gt, total_imgs_bboxes_pre)):
+        perSampleMetrics_dict = evaluator.evaluate_image(gt, pred)
+        results.append(perSampleMetrics_dict)
+
+    metrics = evaluator.combine_results(results)
+    print(metrics)
+    return metrics
+
+def cal_eval(config, data, res_dir_name, opt, mode):
+    evaluator = DetectionIoUEvaluator()
+    test_config = DotDict(config.test[data])
+    res_dir = os.path.join(os.path.join("exp", args.yaml), "{}".format(res_dir_name))
+
+    if opt == "iou_eval":
+        main_eval(
+            config.test.trained_model,
+            config.train.backbone,
+            test_config,
+            evaluator,
+            res_dir,
+            buffer=None,
+            model=None,
+            mode=mode,
+        )
+    else:
+        print("Undefined evaluation")
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="CRAFT Text Detection Eval")
+    parser.add_argument(
+        "--yaml",
+        "--yaml_file_name",
+        default="custom_data_train",
+        type=str,
+        help="Load configuration",
+    )
+    args = parser.parse_args()
+
+    # load configure
+    config = load_yaml(args.yaml)
+    config = DotDict(config)
+
+    if config["wandb_opt"]:
+        wandb.init(project="evaluation", entity="gmuffiness", name=args.yaml)
+        wandb.config.update(config)
+
+    val_result_dir_name = args.yaml
+    cal_eval(
+        config,
+        "custom_data",
+        val_result_dir_name + "-ic15-iou",
+        opt="iou_eval",
+        mode=None,
+    )

trainer/craft/exp/folder.txt

@@ -0,0 +1 @@
+trained model will be saved here

trainer/craft/loss/mseloss.py

@@ -0,0 +1,172 @@
+import torch
+import torch.nn as nn
+
+
+class Loss(nn.Module):
+    def __init__(self):
+        super(Loss, self).__init__()
+
+    def forward(self, gt_region, gt_affinity, pred_region, pred_affinity, conf_map):
+        loss = torch.mean(
+            ((gt_region - pred_region).pow(2) + (gt_affinity - pred_affinity).pow(2))
+            * conf_map
+        )
+        return loss
+
+
+class Maploss_v2(nn.Module):
+    def __init__(self):
+
+        super(Maploss_v2, self).__init__()
+
+    def batch_image_loss(self, pred_score, label_score, neg_rto, n_min_neg):
+
+        # positive_loss
+        positive_pixel = (label_score > 0.1).float()
+        positive_pixel_number = torch.sum(positive_pixel)
+
+        positive_loss_region = pred_score * positive_pixel
+
+        # negative_loss
+        negative_pixel = (label_score <= 0.1).float()
+        negative_pixel_number = torch.sum(negative_pixel)
+        negative_loss_region = pred_score * negative_pixel
+
+        if positive_pixel_number != 0:
+            if negative_pixel_number < neg_rto * positive_pixel_number:
+                negative_loss = (
+                    torch.sum(
+                        torch.topk(
+                            negative_loss_region.view(-1), n_min_neg, sorted=False
+                        )[0]
+                    )
+                    / n_min_neg
+                )
+            else:
+                negative_loss = torch.sum(
+                    torch.topk(
+                        negative_loss_region.view(-1),
+                        int(neg_rto * positive_pixel_number),
+                        sorted=False,
+                    )[0]
+                ) / (positive_pixel_number * neg_rto)
+            positive_loss = torch.sum(positive_loss_region) / positive_pixel_number
+        else:
+            # only negative pixel
+            negative_loss = (
+                torch.sum(
+                    torch.topk(negative_loss_region.view(-1), n_min_neg, sorted=False)[
+                        0
+                    ]
+                )
+                / n_min_neg
+            )
+            positive_loss = 0.0
+        total_loss = positive_loss + negative_loss
+        return total_loss
+
+    def forward(
+        self,
+        region_scores_label,
+        affinity_socres_label,
+        region_scores_pre,
+        affinity_scores_pre,
+        mask,
+        neg_rto,
+        n_min_neg,
+    ):
+        loss_fn = torch.nn.MSELoss(reduce=False, size_average=False)
+        assert (
+            region_scores_label.size() == region_scores_pre.size()
+            and affinity_socres_label.size() == affinity_scores_pre.size()
+        )
+        loss1 = loss_fn(region_scores_pre, region_scores_label)
+        loss2 = loss_fn(affinity_scores_pre, affinity_socres_label)
+
+        loss_region = torch.mul(loss1, mask)
+        loss_affinity = torch.mul(loss2, mask)
+
+        char_loss = self.batch_image_loss(
+            loss_region, region_scores_label, neg_rto, n_min_neg
+        )
+        affi_loss = self.batch_image_loss(
+            loss_affinity, affinity_socres_label, neg_rto, n_min_neg
+        )
+        return char_loss + affi_loss
+
+
+class Maploss_v3(nn.Module):
+    def __init__(self):
+
+        super(Maploss_v3, self).__init__()
+
+    def single_image_loss(self, pre_loss, loss_label, neg_rto, n_min_neg):
+
+        batch_size = pre_loss.shape[0]
+
+        positive_loss, negative_loss = 0, 0
+        for single_loss, single_label in zip(pre_loss, loss_label):
+
+            # positive_loss
+            pos_pixel = (single_label >= 0.1).float()
+            n_pos_pixel = torch.sum(pos_pixel)
+            pos_loss_region = single_loss * pos_pixel
+            positive_loss += torch.sum(pos_loss_region) / max(n_pos_pixel, 1e-12)
+
+            # negative_loss
+            neg_pixel = (single_label < 0.1).float()
+            n_neg_pixel = torch.sum(neg_pixel)
+            neg_loss_region = single_loss * neg_pixel
+
+            if n_pos_pixel != 0:
+                if n_neg_pixel < neg_rto * n_pos_pixel:
+                    negative_loss += torch.sum(neg_loss_region) / n_neg_pixel
+                else:
+                    n_hard_neg = max(n_min_neg, neg_rto * n_pos_pixel)
+                    # n_hard_neg = neg_rto*n_pos_pixel
+                    negative_loss += (
+                        torch.sum(
+                            torch.topk(neg_loss_region.view(-1), int(n_hard_neg))[0]
+                        )
+                        / n_hard_neg
+                    )
+            else:
+                # only negative pixel
+                negative_loss += (
+                    torch.sum(torch.topk(neg_loss_region.view(-1), n_min_neg)[0])
+                    / n_min_neg
+                )
+
+        total_loss = (positive_loss + negative_loss) / batch_size
+
+        return total_loss
+
+    def forward(
+        self,
+        region_scores_label,
+        affinity_scores_label,
+        region_scores_pre,
+        affinity_scores_pre,
+        mask,
+        neg_rto,
+        n_min_neg,
+    ):
+        loss_fn = torch.nn.MSELoss(reduce=False, size_average=False)
+
+        assert (
+            region_scores_label.size() == region_scores_pre.size()
+            and affinity_scores_label.size() == affinity_scores_pre.size()
+        )
+        loss1 = loss_fn(region_scores_pre, region_scores_label)
+        loss2 = loss_fn(affinity_scores_pre, affinity_scores_label)
+
+        loss_region = torch.mul(loss1, mask)
+        loss_affinity = torch.mul(loss2, mask)
+        char_loss = self.single_image_loss(
+            loss_region, region_scores_label, neg_rto, n_min_neg
+        )
+        affi_loss = self.single_image_loss(
+            loss_affinity, affinity_scores_label, neg_rto, n_min_neg
+        )
+
+        return char_loss + affi_loss

trainer/craft/metrics/eval_det_iou.py

@@ -0,0 +1,244 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+from collections import namedtuple
+import numpy as np
+from shapely.geometry import Polygon
+"""
+cite from:
+PaddleOCR, github: https://github.com/PaddlePaddle/PaddleOCR
+PaddleOCR reference from :
+https://github.com/MhLiao/DB/blob/3c32b808d4412680310d3d28eeb6a2d5bf1566c5/concern/icdar2015_eval/detection/iou.py#L8
+"""
+
+
+class DetectionIoUEvaluator(object):
+    def __init__(self, iou_constraint=0.5, area_precision_constraint=0.5):
+        self.iou_constraint = iou_constraint
+        self.area_precision_constraint = area_precision_constraint
+
+    def evaluate_image(self, gt, pred):
+        def get_union(pD, pG):
+            return Polygon(pD).union(Polygon(pG)).area
+
+        def get_intersection_over_union(pD, pG):
+            return get_intersection(pD, pG) / get_union(pD, pG)
+
+        def get_intersection(pD, pG):
+            return Polygon(pD).intersection(Polygon(pG)).area
+
+        def compute_ap(confList, matchList, numGtCare):
+            correct = 0
+            AP = 0
+            if len(confList) > 0:
+                confList = np.array(confList)
+                matchList = np.array(matchList)
+                sorted_ind = np.argsort(-confList)
+                confList = confList[sorted_ind]
+                matchList = matchList[sorted_ind]
+                for n in range(len(confList)):
+                    match = matchList[n]
+                    if match:
+                        correct += 1
+                        AP += float(correct) / (n + 1)
+
+                if numGtCare > 0:
+                    AP /= numGtCare
+
+            return AP
+
+        perSampleMetrics = {}
+
+        matchedSum = 0
+
+        Rectangle = namedtuple('Rectangle', 'xmin ymin xmax ymax')
+
+        numGlobalCareGt = 0
+        numGlobalCareDet = 0
+
+        arrGlobalConfidences = []
+        arrGlobalMatches = []
+
+        recall = 0
+        precision = 0
+        hmean = 0
+
+        detMatched = 0
+
+        iouMat = np.empty([1, 1])
+
+        gtPols = []
+        detPols = []
+
+        gtPolPoints = []
+        detPolPoints = []
+
+        # Array of Ground Truth Polygons' keys marked as don't Care
+        gtDontCarePolsNum = []
+        # Array of Detected Polygons' matched with a don't Care GT
+        detDontCarePolsNum = []
+
+        pairs = []
+        detMatchedNums = []
+
+        arrSampleConfidences = []
+        arrSampleMatch = []
+
+        evaluationLog = ""
+
+        # print(len(gt))
+
+        for n in range(len(gt)):
+            points = gt[n]['points']
+            # transcription = gt[n]['text']
+            dontCare = gt[n]['ignore']
+            #             points = Polygon(points)
+            #             points = points.buffer(0)
+            try:
+                if not Polygon(points).is_valid or not Polygon(points).is_simple:
+                    continue
+            except:
+                import ipdb;
+                ipdb.set_trace()
+
+            #import ipdb;ipdb.set_trace()
+            gtPol = points
+            gtPols.append(gtPol)
+            gtPolPoints.append(points)
+            if dontCare:
+                gtDontCarePolsNum.append(len(gtPols) - 1)
+
+        evaluationLog += "GT polygons: " + str(len(gtPols)) + (
+            " (" + str(len(gtDontCarePolsNum)) + " don't care)\n"
+            if len(gtDontCarePolsNum) > 0 else "\n")
+
+        for n in range(len(pred)):
+            points = pred[n]['points']
+            #             points = Polygon(points)
+            #             points = points.buffer(0)
+            if not Polygon(points).is_valid or not Polygon(points).is_simple:
+                continue
+
+            detPol = points
+            detPols.append(detPol)
+            detPolPoints.append(points)
+            if len(gtDontCarePolsNum) > 0:
+                for dontCarePol in gtDontCarePolsNum:
+                    dontCarePol = gtPols[dontCarePol]
+                    intersected_area = get_intersection(dontCarePol, detPol)
+                    pdDimensions = Polygon(detPol).area
+                    precision = 0 if pdDimensions == 0 else intersected_area / pdDimensions
+                    if (precision > self.area_precision_constraint):
+                        detDontCarePolsNum.append(len(detPols) - 1)
+                        break
+
+        evaluationLog += "DET polygons: " + str(len(detPols)) + (
+            " (" + str(len(detDontCarePolsNum)) + " don't care)\n"
+            if len(detDontCarePolsNum) > 0 else "\n")
+
+        if len(gtPols) > 0 and len(detPols) > 0:
+            # Calculate IoU and precision matrices
+            outputShape = [len(gtPols), len(detPols)]
+            iouMat = np.empty(outputShape)
+            gtRectMat = np.zeros(len(gtPols), np.int8)
+            detRectMat = np.zeros(len(detPols), np.int8)
+            for gtNum in range(len(gtPols)):
+                for detNum in range(len(detPols)):
+                    pG = gtPols[gtNum]
+                    pD = detPols[detNum]
+                    iouMat[gtNum, detNum] = get_intersection_over_union(pD, pG)
+
+            for gtNum in range(len(gtPols)):
+                for detNum in range(len(detPols)):
+                    if gtRectMat[gtNum] == 0 and detRectMat[
+                            detNum] == 0 and gtNum not in gtDontCarePolsNum and detNum not in detDontCarePolsNum:
+                        if iouMat[gtNum, detNum] > self.iou_constraint:
+                            gtRectMat[gtNum] = 1
+                            detRectMat[detNum] = 1
+                            detMatched += 1
+                            pairs.append({'gt': gtNum, 'det': detNum})
+                            detMatchedNums.append(detNum)
+                            evaluationLog += "Match GT #" + \
+                                             str(gtNum) + " with Det #" + str(detNum) + "\n"
+
+        numGtCare = (len(gtPols) - len(gtDontCarePolsNum))
+        numDetCare = (len(detPols) - len(detDontCarePolsNum))
+        if numGtCare == 0:
+            recall = float(1)
+            precision = float(0) if numDetCare > 0 else float(1)
+        else:
+            recall = float(detMatched) / numGtCare
+            precision = 0 if numDetCare == 0 else float(detMatched) / numDetCare
+
+        hmean = 0 if (precision + recall) == 0 else 2.0 * \
+                                                    precision * recall / (precision + recall)
+
+        matchedSum += detMatched
+        numGlobalCareGt += numGtCare
+        numGlobalCareDet += numDetCare
+
+        perSampleMetrics = {
+            'precision': precision,
+            'recall': recall,
+            'hmean': hmean,
+            'pairs': pairs,
+            'iouMat': [] if len(detPols) > 100 else iouMat.tolist(),
+            'gtPolPoints': gtPolPoints,
+            'detPolPoints': detPolPoints,
+            'gtCare': numGtCare,
+            'detCare': numDetCare,
+            'gtDontCare': gtDontCarePolsNum,
+            'detDontCare': detDontCarePolsNum,
+            'detMatched': detMatched,
+            'evaluationLog': evaluationLog
+        }
+
+        return perSampleMetrics
+
+    def combine_results(self, results):
+        numGlobalCareGt = 0
+        numGlobalCareDet = 0
+        matchedSum = 0
+        for result in results:
+            numGlobalCareGt += result['gtCare']
+            numGlobalCareDet += result['detCare']
+            matchedSum += result['detMatched']
+
+        methodRecall = 0 if numGlobalCareGt == 0 else float(
+            matchedSum) / numGlobalCareGt
+        methodPrecision = 0 if numGlobalCareDet == 0 else float(
+            matchedSum) / numGlobalCareDet
+        methodHmean = 0 if methodRecall + methodPrecision == 0 else 2 * \
+                                                                    methodRecall * methodPrecision / (
+                                                                            methodRecall + methodPrecision)
+        # print(methodRecall, methodPrecision, methodHmean)
+        # sys.exit(-1)
+        methodMetrics = {
+            'precision': methodPrecision,
+            'recall': methodRecall,
+            'hmean': methodHmean
+        }
+
+        return methodMetrics
+
+
+if __name__ == '__main__':
+    evaluator = DetectionIoUEvaluator()
+    gts = [[{
+        'points': [(0, 0), (1, 0), (1, 1), (0, 1)],
+        'text': 1234,
+        'ignore': False,
+    }, {
+        'points': [(2, 2), (3, 2), (3, 3), (2, 3)],
+        'text': 5678,
+        'ignore': False,
+    }]]
+    preds = [[{
+        'points': [(0.1, 0.1), (1, 0), (1, 1), (0, 1)],
+        'text': 123,
+        'ignore': False,
+    }]]
+    results = []
+    for gt, pred in zip(gts, preds):
+        results.append(evaluator.evaluate_image(gt, pred))
+    metrics = evaluator.combine_results(results)
+    print(metrics)

trainer/craft/model/craft.py

@@ -0,0 +1,112 @@
+"""  
+Copyright (c) 2019-present NAVER Corp.
+MIT License
+"""
+
+# -*- coding: utf-8 -*-
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from model.vgg16_bn import vgg16_bn, init_weights
+
+class double_conv(nn.Module):
+    def __init__(self, in_ch, mid_ch, out_ch):
+        super(double_conv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_ch + mid_ch, mid_ch, kernel_size=1),
+            nn.BatchNorm2d(mid_ch),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_ch, out_ch, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class CRAFT(nn.Module):
+    def __init__(self, pretrained=True, freeze=False, amp=False):
+        super(CRAFT, self).__init__()
+
+        self.amp = amp
+
+        """ Base network """
+        self.basenet = vgg16_bn(pretrained, freeze)
+
+        """ U network """
+        self.upconv1 = double_conv(1024, 512, 256)
+        self.upconv2 = double_conv(512, 256, 128)
+        self.upconv3 = double_conv(256, 128, 64)
+        self.upconv4 = double_conv(128, 64, 32)
+
+        num_class = 2
+        self.conv_cls = nn.Sequential(
+            nn.Conv2d(32, 32, kernel_size=3, padding=1), nn.ReLU(inplace=True),
+            nn.Conv2d(32, 32, kernel_size=3, padding=1), nn.ReLU(inplace=True),
+            nn.Conv2d(32, 16, kernel_size=3, padding=1), nn.ReLU(inplace=True),
+            nn.Conv2d(16, 16, kernel_size=1), nn.ReLU(inplace=True),
+            nn.Conv2d(16, num_class, kernel_size=1),
+        )
+
+        init_weights(self.upconv1.modules())
+        init_weights(self.upconv2.modules())
+        init_weights(self.upconv3.modules())
+        init_weights(self.upconv4.modules())
+        init_weights(self.conv_cls.modules())
+        
+    def forward(self, x):
+        """ Base network """
+        if self.amp:
+            with torch.cuda.amp.autocast():
+                sources = self.basenet(x)
+
+                """ U network """
+                y = torch.cat([sources[0], sources[1]], dim=1)
+                y = self.upconv1(y)
+
+                y = F.interpolate(y, size=sources[2].size()[2:], mode='bilinear', align_corners=False)
+                y = torch.cat([y, sources[2]], dim=1)
+                y = self.upconv2(y)
+
+                y = F.interpolate(y, size=sources[3].size()[2:], mode='bilinear', align_corners=False)
+                y = torch.cat([y, sources[3]], dim=1)
+                y = self.upconv3(y)
+
+                y = F.interpolate(y, size=sources[4].size()[2:], mode='bilinear', align_corners=False)
+                y = torch.cat([y, sources[4]], dim=1)
+                feature = self.upconv4(y)
+
+                y = self.conv_cls(feature)
+
+                return y.permute(0,2,3,1), feature
+        else:
+
+            sources = self.basenet(x)
+
+            """ U network """
+            y = torch.cat([sources[0], sources[1]], dim=1)
+            y = self.upconv1(y)
+
+            y = F.interpolate(y, size=sources[2].size()[2:], mode='bilinear', align_corners=False)
+            y = torch.cat([y, sources[2]], dim=1)
+            y = self.upconv2(y)
+
+            y = F.interpolate(y, size=sources[3].size()[2:], mode='bilinear', align_corners=False)
+            y = torch.cat([y, sources[3]], dim=1)
+            y = self.upconv3(y)
+
+            y = F.interpolate(y, size=sources[4].size()[2:], mode='bilinear', align_corners=False)
+            y = torch.cat([y, sources[4]], dim=1)
+            feature = self.upconv4(y)
+
+            y = self.conv_cls(feature)
+
+            return y.permute(0, 2, 3, 1), feature
+
+if __name__ == '__main__':
+    model = CRAFT(pretrained=True).cuda()
+    output, _ = model(torch.randn(1, 3, 768, 768).cuda())
+    print(output.shape)

trainer/craft/model/vgg16_bn.py

@@ -0,0 +1,77 @@
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+import torchvision
+from torchvision import models
+from packaging import version
+
+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=True):
+        super(vgg16_bn, self).__init__()
+        if version.parse(torchvision.__version__) >= version.parse('0.13'):
+            vgg_pretrained_features = models.vgg16_bn(
+                weights=models.VGG16_BN_Weights.DEFAULT if pretrained else None
+            ).features
+        else: # torchvision.__version__ < 0.13
+            models.vgg.model_urls['vgg16_bn'] = models.vgg.model_urls['vgg16_bn'].replace('https://', 'http://')
+            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
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        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
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        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)
+        )
+
+        if not pretrained:
+            init_weights(self.slice1.modules())
+            init_weights(self.slice2.modules())
+            init_weights(self.slice3.modules())
+            init_weights(self.slice4.modules())
+
+        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
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu2_2 = h
+        h = self.slice2(h)
+        h_relu3_2 = h
+        h = self.slice3(h)
+        h_relu4_3 = h
+        h = self.slice4(h)
+        h_relu5_3 = h
+        h = self.slice5(h)
+        h_fc7 = h
+        return h_fc7, h_relu5_3, h_relu4_3, h_relu3_2, h_relu2_2

trainer/craft/requirements.txt

@@ -0,0 +1,10 @@
+conda==4.10.3
+opencv-python==4.5.3.56
+Pillow==9.3.0
+Polygon3==3.0.9.1
+PyYAML==5.4.1
+scikit-image==0.17.2
+Shapely==1.8.0
+torch==1.13.1
+torchvision==0.10.0
+wandb==0.12.9

trainer/craft/scripts/run_cde.sh

@@ -0,0 +1,7 @@
+# sed -i -e 's/\r$//' scripts/run_cde.sh
+EXP_NAME=custom_data_release_test_3
+yaml_path="config/$EXP_NAME.yaml"
+cp config/custom_data_train.yaml $yaml_path
+#CUDA_VISIBLE_DEVICES=0,1 python3 train_distributed.py --yaml=$EXP_NAME --port=2468
+CUDA_VISIBLE_DEVICES=0 python3 train.py --yaml=$EXP_NAME --port=2468
+rm "config/$EXP_NAME.yaml"

trainer/craft/train.py

@@ -0,0 +1,479 @@
+# -*- coding: utf-8 -*-
+import argparse
+import os
+import shutil
+import time
+import multiprocessing as mp
+import yaml
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import wandb
+
+from config.load_config import load_yaml, DotDict
+from data.dataset import SynthTextDataSet, CustomDataset
+from loss.mseloss import Maploss_v2, Maploss_v3
+from model.craft import CRAFT
+from eval import main_eval
+from metrics.eval_det_iou import DetectionIoUEvaluator
+from utils.util import copyStateDict, save_parser
+
+
+class Trainer(object):
+    def __init__(self, config, gpu, mode):
+
+        self.config = config
+        self.gpu = gpu
+        self.mode = mode
+        self.net_param = self.get_load_param(gpu)
+
+    def get_synth_loader(self):
+
+        dataset = SynthTextDataSet(
+            output_size=self.config.train.data.output_size,
+            data_dir=self.config.train.synth_data_dir,
+            saved_gt_dir=None,
+            mean=self.config.train.data.mean,
+            variance=self.config.train.data.variance,
+            gauss_init_size=self.config.train.data.gauss_init_size,
+            gauss_sigma=self.config.train.data.gauss_sigma,
+            enlarge_region=self.config.train.data.enlarge_region,
+            enlarge_affinity=self.config.train.data.enlarge_affinity,
+            aug=self.config.train.data.syn_aug,
+            vis_test_dir=self.config.vis_test_dir,
+            vis_opt=self.config.train.data.vis_opt,
+            sample=self.config.train.data.syn_sample,
+        )
+
+        syn_loader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=self.config.train.batch_size // self.config.train.synth_ratio,
+            shuffle=False,
+            num_workers=self.config.train.num_workers,
+            drop_last=True,
+            pin_memory=True,
+        )
+        return syn_loader
+
+    def get_custom_dataset(self):
+
+        custom_dataset = CustomDataset(
+            output_size=self.config.train.data.output_size,
+            data_dir=self.config.data_root_dir,
+            saved_gt_dir=None,
+            mean=self.config.train.data.mean,
+            variance=self.config.train.data.variance,
+            gauss_init_size=self.config.train.data.gauss_init_size,
+            gauss_sigma=self.config.train.data.gauss_sigma,
+            enlarge_region=self.config.train.data.enlarge_region,
+            enlarge_affinity=self.config.train.data.enlarge_affinity,
+            watershed_param=self.config.train.data.watershed,
+            aug=self.config.train.data.custom_aug,
+            vis_test_dir=self.config.vis_test_dir,
+            sample=self.config.train.data.custom_sample,
+            vis_opt=self.config.train.data.vis_opt,
+            pseudo_vis_opt=self.config.train.data.pseudo_vis_opt,
+            do_not_care_label=self.config.train.data.do_not_care_label,
+        )
+
+        return custom_dataset
+
+    def get_load_param(self, gpu):
+
+        if self.config.train.ckpt_path is not None:
+            map_location = "cuda:%d" % gpu
+            param = torch.load(self.config.train.ckpt_path, map_location=map_location)
+        else:
+            param = None
+
+        return param
+
+    def adjust_learning_rate(self, optimizer, gamma, step, lr):
+        lr = lr * (gamma ** step)
+        for param_group in optimizer.param_groups:
+            param_group["lr"] = lr
+        return param_group["lr"]
+
+    def get_loss(self):
+        if self.config.train.loss == 2:
+            criterion = Maploss_v2()
+        elif self.config.train.loss == 3:
+            criterion = Maploss_v3()
+        else:
+            raise Exception("Undefined loss")
+        return criterion
+
+    def iou_eval(self, dataset, train_step, buffer, model):
+        test_config = DotDict(self.config.test[dataset])
+
+        val_result_dir = os.path.join(
+            self.config.results_dir, "{}/{}".format(dataset + "_iou", str(train_step))
+        )
+
+        evaluator = DetectionIoUEvaluator()
+
+        metrics = main_eval(
+            None,
+            self.config.train.backbone,
+            test_config,
+            evaluator,
+            val_result_dir,
+            buffer,
+            model,
+            self.mode,
+        )
+        if self.gpu == 0 and self.config.wandb_opt:
+            wandb.log(
+                {
+                    "{} iou Recall".format(dataset): np.round(metrics["recall"], 3),
+                    "{} iou Precision".format(dataset): np.round(
+                        metrics["precision"], 3
+                    ),
+                    "{} iou F1-score".format(dataset): np.round(metrics["hmean"], 3),
+                }
+            )
+
+    def train(self, buffer_dict):
+
+        torch.cuda.set_device(self.gpu)
+
+        # MODEL -------------------------------------------------------------------------------------------------------#
+        # SUPERVISION model
+        if self.config.mode == "weak_supervision":
+            if self.config.train.backbone == "vgg":
+                supervision_model = CRAFT(pretrained=False, amp=self.config.train.amp)
+            else:
+                raise Exception("Undefined architecture")
+
+            supervision_device = self.gpu
+            if self.config.train.ckpt_path is not None:
+                supervision_param = self.get_load_param(supervision_device)
+                supervision_model.load_state_dict(
+                    copyStateDict(supervision_param["craft"])
+                )
+                supervision_model = supervision_model.to(f"cuda:{supervision_device}")
+            print(f"Supervision model loading on : gpu {supervision_device}")
+        else:
+            supervision_model, supervision_device = None, None
+
+        # TRAIN model
+        if self.config.train.backbone == "vgg":
+            craft = CRAFT(pretrained=False, amp=self.config.train.amp)
+        else:
+            raise Exception("Undefined architecture")
+
+        if self.config.train.ckpt_path is not None:
+            craft.load_state_dict(copyStateDict(self.net_param["craft"]))
+
+        craft = craft.cuda()
+        craft = torch.nn.DataParallel(craft)
+
+        torch.backends.cudnn.benchmark = True
+
+        # DATASET -----------------------------------------------------------------------------------------------------#
+
+        if self.config.train.use_synthtext:
+            trn_syn_loader = self.get_synth_loader()
+            batch_syn = iter(trn_syn_loader)
+
+        if self.config.train.real_dataset == "custom":
+            trn_real_dataset = self.get_custom_dataset()
+        else:
+            raise Exception("Undefined dataset")
+
+        if self.config.mode == "weak_supervision":
+            trn_real_dataset.update_model(supervision_model)
+            trn_real_dataset.update_device(supervision_device)
+
+        trn_real_loader = torch.utils.data.DataLoader(
+            trn_real_dataset,
+            batch_size=self.config.train.batch_size,
+            shuffle=False,
+            num_workers=self.config.train.num_workers,
+            drop_last=False,
+            pin_memory=True,
+        )
+
+        # OPTIMIZER ---------------------------------------------------------------------------------------------------#
+        optimizer = optim.Adam(
+            craft.parameters(),
+            lr=self.config.train.lr,
+            weight_decay=self.config.train.weight_decay,
+        )
+
+        if self.config.train.ckpt_path is not None and self.config.train.st_iter != 0:
+            optimizer.load_state_dict(copyStateDict(self.net_param["optimizer"]))
+            self.config.train.st_iter = self.net_param["optimizer"]["state"][0]["step"]
+            self.config.train.lr = self.net_param["optimizer"]["param_groups"][0]["lr"]
+
+        # LOSS --------------------------------------------------------------------------------------------------------#
+        # mixed precision
+        if self.config.train.amp:
+            scaler = torch.cuda.amp.GradScaler()
+
+            if (
+                    self.config.train.ckpt_path is not None
+                    and self.config.train.st_iter != 0
+            ):
+                scaler.load_state_dict(copyStateDict(self.net_param["scaler"]))
+        else:
+            scaler = None
+
+        criterion = self.get_loss()
+
+        # TRAIN -------------------------------------------------------------------------------------------------------#
+        train_step = self.config.train.st_iter
+        whole_training_step = self.config.train.end_iter
+        update_lr_rate_step = 0
+        training_lr = self.config.train.lr
+        loss_value = 0
+        batch_time = 0
+        start_time = time.time()
+
+        print(
+            "================================ Train start ================================"
+        )
+        while train_step < whole_training_step:
+            for (
+                    index,
+                    (
+                            images,
+                            region_scores,
+                            affinity_scores,
+                            confidence_masks,
+                    ),
+            ) in enumerate(trn_real_loader):
+                craft.train()
+                if train_step > 0 and train_step % self.config.train.lr_decay == 0:
+                    update_lr_rate_step += 1
+                    training_lr = self.adjust_learning_rate(
+                        optimizer,
+                        self.config.train.gamma,
+                        update_lr_rate_step,
+                        self.config.train.lr,
+                    )
+
+                images = images.cuda(non_blocking=True)
+                region_scores = region_scores.cuda(non_blocking=True)
+                affinity_scores = affinity_scores.cuda(non_blocking=True)
+                confidence_masks = confidence_masks.cuda(non_blocking=True)
+
+                if self.config.train.use_synthtext:
+                    # Synth image load
+                    syn_image, syn_region_label, syn_affi_label, syn_confidence_mask = next(
+                        batch_syn
+                    )
+                    syn_image = syn_image.cuda(non_blocking=True)
+                    syn_region_label = syn_region_label.cuda(non_blocking=True)
+                    syn_affi_label = syn_affi_label.cuda(non_blocking=True)
+                    syn_confidence_mask = syn_confidence_mask.cuda(non_blocking=True)
+
+                    # concat syn & custom image
+                    images = torch.cat((syn_image, images), 0)
+                    region_image_label = torch.cat(
+                        (syn_region_label, region_scores), 0
+                    )
+                    affinity_image_label = torch.cat((syn_affi_label, affinity_scores), 0)
+                    confidence_mask_label = torch.cat(
+                        (syn_confidence_mask, confidence_masks), 0
+                    )
+                else:
+                    region_image_label = region_scores
+                    affinity_image_label = affinity_scores
+                    confidence_mask_label = confidence_masks
+
+                if self.config.train.amp:
+                    with torch.cuda.amp.autocast():
+
+                        output, _ = craft(images)
+                        out1 = output[:, :, :, 0]
+                        out2 = output[:, :, :, 1]
+
+                        loss = criterion(
+                            region_image_label,
+                            affinity_image_label,
+                            out1,
+                            out2,
+                            confidence_mask_label,
+                            self.config.train.neg_rto,
+                            self.config.train.n_min_neg,
+                        )
+
+                    optimizer.zero_grad()
+                    scaler.scale(loss).backward()
+                    scaler.step(optimizer)
+                    scaler.update()
+
+                else:
+                    output, _ = craft(images)
+                    out1 = output[:, :, :, 0]
+                    out2 = output[:, :, :, 1]
+                    loss = criterion(
+                        region_image_label,
+                        affinity_image_label,
+                        out1,
+                        out2,
+                        confidence_mask_label,
+                        self.config.train.neg_rto,
+                    )
+
+                    optimizer.zero_grad()
+                    loss.backward()
+                    optimizer.step()
+
+                end_time = time.time()
+                loss_value += loss.item()
+                batch_time += end_time - start_time
+
+                if train_step > 0 and train_step % 5 == 0:
+                    mean_loss = loss_value / 5
+                    loss_value = 0
+                    avg_batch_time = batch_time / 5
+                    batch_time = 0
+
+                    print(
+                        "{}, training_step: {}|{}, learning rate: {:.8f}, "
+                        "training_loss: {:.5f}, avg_batch_time: {:.5f}".format(
+                            time.strftime(
+                                "%Y-%m-%d:%H:%M:%S", time.localtime(time.time())
+                            ),
+                            train_step,
+                            whole_training_step,
+                            training_lr,
+                            mean_loss,
+                            avg_batch_time,
+                        )
+                    )
+
+                    if self.config.wandb_opt:
+                        wandb.log({"train_step": train_step, "mean_loss": mean_loss})
+
+                if (
+                        train_step % self.config.train.eval_interval == 0
+                        and train_step != 0
+                ):
+
+                    craft.eval()
+
+                    print("Saving state, index:", train_step)
+                    save_param_dic = {
+                        "iter": train_step,
+                        "craft": craft.state_dict(),
+                        "optimizer": optimizer.state_dict(),
+                    }
+                    save_param_path = (
+                            self.config.results_dir
+                            + "/CRAFT_clr_"
+                            + repr(train_step)
+                            + ".pth"
+                    )
+
+                    if self.config.train.amp:
+                        save_param_dic["scaler"] = scaler.state_dict()
+                        save_param_path = (
+                                self.config.results_dir
+                                + "/CRAFT_clr_amp_"
+                                + repr(train_step)
+                                + ".pth"
+                        )
+
+                    torch.save(save_param_dic, save_param_path)
+
+                    # validation
+                    self.iou_eval(
+                        "custom_data",
+                        train_step,
+                        buffer_dict["custom_data"],
+                        craft,
+                    )
+
+                train_step += 1
+                if train_step >= whole_training_step:
+                    break
+
+            if self.config.mode == "weak_supervision":
+                state_dict = craft.module.state_dict()
+                supervision_model.load_state_dict(state_dict)
+                trn_real_dataset.update_model(supervision_model)
+
+        # save last model
+        save_param_dic = {
+            "iter": train_step,
+            "craft": craft.state_dict(),
+            "optimizer": optimizer.state_dict(),
+        }
+        save_param_path = (
+                self.config.results_dir + "/CRAFT_clr_" + repr(train_step) + ".pth"
+        )
+
+        if self.config.train.amp:
+            save_param_dic["scaler"] = scaler.state_dict()
+            save_param_path = (
+                    self.config.results_dir
+                    + "/CRAFT_clr_amp_"
+                    + repr(train_step)
+                    + ".pth"
+            )
+        torch.save(save_param_dic, save_param_path)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="CRAFT custom data train")
+    parser.add_argument(
+        "--yaml",
+        "--yaml_file_name",
+        default="custom_data_train",
+        type=str,
+        help="Load configuration",
+    )
+    parser.add_argument(
+        "--port", "--use ddp port", default="2346", type=str, help="Port number"
+    )
+
+    args = parser.parse_args()
+
+    # load configure
+    exp_name = args.yaml
+    config = load_yaml(args.yaml)
+
+    print("-" * 20 + " Options " + "-" * 20)
+    print(yaml.dump(config))
+    print("-" * 40)
+
+    # Make result_dir
+    res_dir = os.path.join(config["results_dir"], args.yaml)
+    config["results_dir"] = res_dir
+    if not os.path.exists(res_dir):
+        os.makedirs(res_dir)
+
+    # Duplicate yaml file to result_dir
+    shutil.copy(
+        "config/" + args.yaml + ".yaml", os.path.join(res_dir, args.yaml) + ".yaml"
+    )
+
+    if config["mode"] == "weak_supervision":
+        mode = "weak_supervision"
+    else:
+        mode = None
+
+
+    # Apply config to wandb
+    if config["wandb_opt"]:
+        wandb.init(project="craft-stage2", entity="user_name", name=exp_name)
+        wandb.config.update(config)
+
+    config = DotDict(config)
+
+    # Start train
+    buffer_dict = {"custom_data":None}
+    trainer = Trainer(config, 0, mode)
+    trainer.train(buffer_dict)
+
+    if config["wandb_opt"]:
+        wandb.finish()
+
+
+if __name__ == "__main__":
+    main()

trainer/craft/trainSynth.py

@@ -0,0 +1,409 @@
+# -*- coding: utf-8 -*-
+import argparse
+import os
+import shutil
+import time
+import yaml
+import multiprocessing as mp
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import wandb
+
+from config.load_config import load_yaml, DotDict
+from data.dataset import SynthTextDataSet
+from loss.mseloss import Maploss_v2, Maploss_v3
+from model.craft import CRAFT
+from metrics.eval_det_iou import DetectionIoUEvaluator
+from eval import main_eval
+from utils.util import copyStateDict, save_parser
+
+
+class Trainer(object):
+    def __init__(self, config, gpu):
+
+        self.config = config
+        self.gpu = gpu
+        self.mode = None
+        self.trn_loader, self.trn_sampler = self.get_trn_loader()
+        self.net_param = self.get_load_param(gpu)
+
+    def get_trn_loader(self):
+
+        dataset = SynthTextDataSet(
+            output_size=self.config.train.data.output_size,
+            data_dir=self.config.data_dir.synthtext,
+            saved_gt_dir=None,
+            mean=self.config.train.data.mean,
+            variance=self.config.train.data.variance,
+            gauss_init_size=self.config.train.data.gauss_init_size,
+            gauss_sigma=self.config.train.data.gauss_sigma,
+            enlarge_region=self.config.train.data.enlarge_region,
+            enlarge_affinity=self.config.train.data.enlarge_affinity,
+            aug=self.config.train.data.syn_aug,
+            vis_test_dir=self.config.vis_test_dir,
+            vis_opt=self.config.train.data.vis_opt,
+            sample=self.config.train.data.syn_sample,
+        )
+
+        trn_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+
+        trn_loader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=self.config.train.batch_size,
+            shuffle=False,
+            num_workers=self.config.train.num_workers,
+            sampler=trn_sampler,
+            drop_last=True,
+            pin_memory=True,
+        )
+        return trn_loader, trn_sampler
+
+    def get_load_param(self, gpu):
+        if self.config.train.ckpt_path is not None:
+            map_location = {"cuda:%d" % 0: "cuda:%d" % gpu}
+            param = torch.load(self.config.train.ckpt_path, map_location=map_location)
+        else:
+            param = None
+        return param
+
+    def adjust_learning_rate(self, optimizer, gamma, step, lr):
+        lr = lr * (gamma ** step)
+        for param_group in optimizer.param_groups:
+            param_group["lr"] = lr
+        return param_group["lr"]
+
+    def get_loss(self):
+        if self.config.train.loss == 2:
+            criterion = Maploss_v2()
+        elif self.config.train.loss == 3:
+            criterion = Maploss_v3()
+        else:
+            raise Exception("Undefined loss")
+        return criterion
+
+    def iou_eval(self, dataset, train_step, save_param_path, buffer, model):
+        test_config = DotDict(self.config.test[dataset])
+
+        val_result_dir = os.path.join(
+            self.config.results_dir, "{}/{}".format(dataset + "_iou", str(train_step))
+        )
+
+        evaluator = DetectionIoUEvaluator()
+
+        metrics = main_eval(
+            save_param_path,
+            self.config.train.backbone,
+            test_config,
+            evaluator,
+            val_result_dir,
+            buffer,
+            model,
+            self.mode,
+        )
+        if self.gpu == 0 and self.config.wandb_opt:
+            wandb.log(
+                {
+                    "{} IoU Recall".format(dataset): np.round(metrics["recall"], 3),
+                    "{} IoU Precision".format(dataset): np.round(
+                        metrics["precision"], 3
+                    ),
+                    "{} IoU F1-score".format(dataset): np.round(metrics["hmean"], 3),
+                }
+            )
+
+    def train(self, buffer_dict):
+        torch.cuda.set_device(self.gpu)
+
+        # DATASET -----------------------------------------------------------------------------------------------------#
+        trn_loader = self.trn_loader
+
+        # MODEL -------------------------------------------------------------------------------------------------------#
+        if self.config.train.backbone == "vgg":
+            craft = CRAFT(pretrained=True, amp=self.config.train.amp)
+        else:
+            raise Exception("Undefined architecture")
+
+        if self.config.train.ckpt_path is not None:
+            craft.load_state_dict(copyStateDict(self.net_param["craft"]))
+        craft = nn.SyncBatchNorm.convert_sync_batchnorm(craft)
+        craft = craft.cuda()
+        craft = torch.nn.parallel.DistributedDataParallel(craft, device_ids=[self.gpu])
+
+        torch.backends.cudnn.benchmark = True
+
+        # OPTIMIZER----------------------------------------------------------------------------------------------------#
+
+        optimizer = optim.Adam(
+            craft.parameters(),
+            lr=self.config.train.lr,
+            weight_decay=self.config.train.weight_decay,
+        )
+
+        if self.config.train.ckpt_path is not None and self.config.train.st_iter != 0:
+            optimizer.load_state_dict(copyStateDict(self.net_param["optimizer"]))
+            self.config.train.st_iter = self.net_param["optimizer"]["state"][0]["step"]
+            self.config.train.lr = self.net_param["optimizer"]["param_groups"][0]["lr"]
+
+        # LOSS --------------------------------------------------------------------------------------------------------#
+        # mixed precision
+        if self.config.train.amp:
+            scaler = torch.cuda.amp.GradScaler()
+
+            # load model
+            if (
+                self.config.train.ckpt_path is not None
+                and self.config.train.st_iter != 0
+            ):
+                scaler.load_state_dict(copyStateDict(self.net_param["scaler"]))
+        else:
+            scaler = None
+
+        criterion = self.get_loss()
+
+        # TRAIN -------------------------------------------------------------------------------------------------------#
+        train_step = self.config.train.st_iter
+        whole_training_step = self.config.train.end_iter
+        update_lr_rate_step = 0
+        training_lr = self.config.train.lr
+        loss_value = 0
+        batch_time = 0
+        epoch = 0
+        start_time = time.time()
+
+        while train_step < whole_training_step:
+            self.trn_sampler.set_epoch(train_step)
+            for (
+                index,
+                (image, region_image, affinity_image, confidence_mask,),
+            ) in enumerate(trn_loader):
+                craft.train()
+                if train_step > 0 and train_step % self.config.train.lr_decay == 0:
+                    update_lr_rate_step += 1
+                    training_lr = self.adjust_learning_rate(
+                        optimizer,
+                        self.config.train.gamma,
+                        update_lr_rate_step,
+                        self.config.train.lr,
+                    )
+
+                images = image.cuda(non_blocking=True)
+                region_image_label = region_image.cuda(non_blocking=True)
+                affinity_image_label = affinity_image.cuda(non_blocking=True)
+                confidence_mask_label = confidence_mask.cuda(non_blocking=True)
+
+                if self.config.train.amp:
+                    with torch.cuda.amp.autocast():
+
+                        output, _ = craft(images)
+                        out1 = output[:, :, :, 0]
+                        out2 = output[:, :, :, 1]
+
+                        loss = criterion(
+                            region_image_label,
+                            affinity_image_label,
+                            out1,
+                            out2,
+                            confidence_mask_label,
+                            self.config.train.neg_rto,
+                            self.config.train.n_min_neg,
+                        )
+
+                    optimizer.zero_grad()
+                    scaler.scale(loss).backward()
+                    scaler.step(optimizer)
+                    scaler.update()
+
+                else:
+                    output, _ = craft(images)
+                    out1 = output[:, :, :, 0]
+                    out2 = output[:, :, :, 1]
+                    loss = criterion(
+                        region_image_label,
+                        affinity_image_label,
+                        out1,
+                        out2,
+                        confidence_mask_label,
+                        self.config.train.neg_rto,
+                    )
+
+                    optimizer.zero_grad()
+                    loss.backward()
+                    optimizer.step()
+
+                end_time = time.time()
+                loss_value += loss.item()
+                batch_time += end_time - start_time
+
+                if train_step > 0 and train_step % 5 == 0 and self.gpu == 0:
+                    mean_loss = loss_value / 5
+                    loss_value = 0
+                    avg_batch_time = batch_time / 5
+                    batch_time = 0
+
+                    print(
+                        "{}, training_step: {}|{}, learning rate: {:.8f}, "
+                        "training_loss: {:.5f}, avg_batch_time: {:.5f}".format(
+                            time.strftime(
+                                "%Y-%m-%d:%H:%M:%S", time.localtime(time.time())
+                            ),
+                            train_step,
+                            whole_training_step,
+                            training_lr,
+                            mean_loss,
+                            avg_batch_time,
+                        )
+                    )
+                    if self.gpu == 0 and self.config.wandb_opt:
+                        wandb.log({"train_step": train_step, "mean_loss": mean_loss})
+
+                if (
+                    train_step % self.config.train.eval_interval == 0
+                    and train_step != 0
+                ):
+
+                    # initialize all buffer value with zero
+                    if self.gpu == 0:
+                        for buffer in buffer_dict.values():
+                            for i in range(len(buffer)):
+                                buffer[i] = None
+
+                    print("Saving state, index:", train_step)
+                    save_param_dic = {
+                        "iter": train_step,
+                        "craft": craft.state_dict(),
+                        "optimizer": optimizer.state_dict(),
+                    }
+                    save_param_path = (
+                        self.config.results_dir
+                        + "/CRAFT_clr_"
+                        + repr(train_step)
+                        + ".pth"
+                    )
+
+                    if self.config.train.amp:
+                        save_param_dic["scaler"] = scaler.state_dict()
+                        save_param_path = (
+                            self.config.results_dir
+                            + "/CRAFT_clr_amp_"
+                            + repr(train_step)
+                            + ".pth"
+                        )
+                        
+                    if self.gpu == 0:
+                        torch.save(save_param_dic, save_param_path)
+
+                    # validation
+                    self.iou_eval(
+                        "icdar2013",
+                        train_step,
+                        save_param_path,
+                        buffer_dict["icdar2013"],
+                        craft,
+                    )
+
+                train_step += 1
+                if train_step >= whole_training_step:
+                    break
+            epoch += 1
+
+        # save last model
+        if self.gpu == 0:
+            save_param_dic = {
+                "iter": train_step,
+                "craft": craft.state_dict(),
+                "optimizer": optimizer.state_dict(),
+            }
+            save_param_path = (
+                self.config.results_dir + "/CRAFT_clr_" + repr(train_step) + ".pth"
+            )
+
+            if self.config.train.amp:
+                save_param_dic["scaler"] = scaler.state_dict()
+                save_param_path = (
+                    self.config.results_dir
+                    + "/CRAFT_clr_amp_"
+                    + repr(train_step)
+                    + ".pth"
+                )
+            torch.save(save_param_dic, save_param_path)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="CRAFT SynthText Train")
+    parser.add_argument(
+        "--yaml",
+        "--yaml_file_name",
+        default="syn_train",
+        type=str,
+        help="Load configuration",
+    )
+    parser.add_argument(
+        "--port", "--use ddp port", default="2646", type=str, help="Load configuration"
+    )
+
+    args = parser.parse_args()
+
+    # load configure
+    exp_name = args.yaml
+    config = load_yaml(args.yaml)
+
+    print("-" * 20 + " Options " + "-" * 20)
+    print(yaml.dump(config))
+    print("-" * 40)
+
+    # Make result_dir
+    res_dir = os.path.join(config["results_dir"], args.yaml)
+    config["results_dir"] = res_dir
+    if not os.path.exists(res_dir):
+        os.makedirs(res_dir)
+
+    # Duplicate yaml file to result_dir
+    shutil.copy(
+        "config/" + args.yaml + ".yaml", os.path.join(res_dir, args.yaml) + ".yaml"
+    )
+
+    ngpus_per_node = torch.cuda.device_count()
+    print(f"Total device num : {ngpus_per_node}")
+
+    manager = mp.Manager()
+    buffer1 = manager.list([None] * config["test"]["icdar2013"]["test_set_size"])
+    buffer_dict = {"icdar2013": buffer1}
+    torch.multiprocessing.spawn(
+        main_worker,
+        nprocs=ngpus_per_node,
+        args=(args.port, ngpus_per_node, config, buffer_dict, exp_name,),
+    )
+    print('flag5')
+
+
+def main_worker(gpu, port, ngpus_per_node, config, buffer_dict, exp_name):
+
+    torch.distributed.init_process_group(
+        backend="nccl",
+        init_method="tcp://127.0.0.1:" + port,
+        world_size=ngpus_per_node,
+        rank=gpu,
+    )
+
+    # Apply config to wandb
+    if gpu == 0 and config["wandb_opt"]:
+        wandb.init(project="craft-stage1", entity="gmuffiness", name=exp_name)
+        wandb.config.update(config)
+
+    batch_size = int(config["train"]["batch_size"] / ngpus_per_node)
+    config["train"]["batch_size"] = batch_size
+    config = DotDict(config)
+
+    # Start train
+    trainer = Trainer(config, gpu)
+    trainer.train(buffer_dict)
+
+    if gpu == 0 and config["wandb_opt"]:
+        wandb.finish()
+    torch.distributed.destroy_process_group()
+
+if __name__ == "__main__":
+    main()

trainer/craft/train_distributed.py

@@ -0,0 +1,523 @@
+# -*- coding: utf-8 -*-
+import argparse
+import os
+import shutil
+import time
+import multiprocessing as mp
+import yaml
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import wandb
+
+from config.load_config import load_yaml, DotDict
+from data.dataset import SynthTextDataSet, CustomDataset
+from loss.mseloss import Maploss_v2, Maploss_v3
+from model.craft import CRAFT
+from eval import main_eval
+from metrics.eval_det_iou import DetectionIoUEvaluator
+from utils.util import copyStateDict, save_parser
+
+
+class Trainer(object):
+    def __init__(self, config, gpu, mode):
+
+        self.config = config
+        self.gpu = gpu
+        self.mode = mode
+        self.net_param = self.get_load_param(gpu)
+
+    def get_synth_loader(self):
+
+        dataset = SynthTextDataSet(
+            output_size=self.config.train.data.output_size,
+            data_dir=self.config.train.synth_data_dir,
+            saved_gt_dir=None,
+            mean=self.config.train.data.mean,
+            variance=self.config.train.data.variance,
+            gauss_init_size=self.config.train.data.gauss_init_size,
+            gauss_sigma=self.config.train.data.gauss_sigma,
+            enlarge_region=self.config.train.data.enlarge_region,
+            enlarge_affinity=self.config.train.data.enlarge_affinity,
+            aug=self.config.train.data.syn_aug,
+            vis_test_dir=self.config.vis_test_dir,
+            vis_opt=self.config.train.data.vis_opt,
+            sample=self.config.train.data.syn_sample,
+        )
+
+        syn_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+
+        syn_loader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=self.config.train.batch_size // self.config.train.synth_ratio,
+            shuffle=False,
+            num_workers=self.config.train.num_workers,
+            sampler=syn_sampler,
+            drop_last=True,
+            pin_memory=True,
+        )
+        return syn_loader
+
+    def get_custom_dataset(self):
+
+        custom_dataset = CustomDataset(
+            output_size=self.config.train.data.output_size,
+            data_dir=self.config.data_root_dir,
+            saved_gt_dir=None,
+            mean=self.config.train.data.mean,
+            variance=self.config.train.data.variance,
+            gauss_init_size=self.config.train.data.gauss_init_size,
+            gauss_sigma=self.config.train.data.gauss_sigma,
+            enlarge_region=self.config.train.data.enlarge_region,
+            enlarge_affinity=self.config.train.data.enlarge_affinity,
+            watershed_param=self.config.train.data.watershed,
+            aug=self.config.train.data.custom_aug,
+            vis_test_dir=self.config.vis_test_dir,
+            sample=self.config.train.data.custom_sample,
+            vis_opt=self.config.train.data.vis_opt,
+            pseudo_vis_opt=self.config.train.data.pseudo_vis_opt,
+            do_not_care_label=self.config.train.data.do_not_care_label,
+        )
+
+        return custom_dataset
+
+    def get_load_param(self, gpu):
+
+        if self.config.train.ckpt_path is not None:
+            map_location = "cuda:%d" % gpu
+            param = torch.load(self.config.train.ckpt_path, map_location=map_location)
+        else:
+            param = None
+
+        return param
+
+    def adjust_learning_rate(self, optimizer, gamma, step, lr):
+        lr = lr * (gamma ** step)
+        for param_group in optimizer.param_groups:
+            param_group["lr"] = lr
+        return param_group["lr"]
+
+    def get_loss(self):
+        if self.config.train.loss == 2:
+            criterion = Maploss_v2()
+        elif self.config.train.loss == 3:
+            criterion = Maploss_v3()
+        else:
+            raise Exception("Undefined loss")
+        return criterion
+
+    def iou_eval(self, dataset, train_step, buffer, model):
+        test_config = DotDict(self.config.test[dataset])
+
+        val_result_dir = os.path.join(
+            self.config.results_dir, "{}/{}".format(dataset + "_iou", str(train_step))
+        )
+
+        evaluator = DetectionIoUEvaluator()
+
+        metrics = main_eval(
+            None,
+            self.config.train.backbone,
+            test_config,
+            evaluator,
+            val_result_dir,
+            buffer,
+            model,
+            self.mode,
+        )
+        if self.gpu == 0 and self.config.wandb_opt:
+            wandb.log(
+                {
+                    "{} iou Recall".format(dataset): np.round(metrics["recall"], 3),
+                    "{} iou Precision".format(dataset): np.round(
+                        metrics["precision"], 3
+                    ),
+                    "{} iou F1-score".format(dataset): np.round(metrics["hmean"], 3),
+                }
+            )
+
+    def train(self, buffer_dict):
+
+        torch.cuda.set_device(self.gpu)
+        total_gpu_num = torch.cuda.device_count()
+
+        # MODEL -------------------------------------------------------------------------------------------------------#
+        # SUPERVISION model
+        if self.config.mode == "weak_supervision":
+            if self.config.train.backbone == "vgg":
+                supervision_model = CRAFT(pretrained=False, amp=self.config.train.amp)
+            else:
+                raise Exception("Undefined architecture")
+
+            # NOTE: only work on half GPU assign train / half GPU assign supervision setting
+            supervision_device = total_gpu_num // 2 + self.gpu
+            if self.config.train.ckpt_path is not None:
+                supervision_param = self.get_load_param(supervision_device)
+                supervision_model.load_state_dict(
+                    copyStateDict(supervision_param["craft"])
+                )
+                supervision_model = supervision_model.to(f"cuda:{supervision_device}")
+            print(f"Supervision model loading on : gpu {supervision_device}")
+        else:
+            supervision_model, supervision_device = None, None
+
+        # TRAIN model
+        if self.config.train.backbone == "vgg":
+            craft = CRAFT(pretrained=False, amp=self.config.train.amp)
+        else:
+            raise Exception("Undefined architecture")
+
+        if self.config.train.ckpt_path is not None:
+            craft.load_state_dict(copyStateDict(self.net_param["craft"]))
+
+        craft = nn.SyncBatchNorm.convert_sync_batchnorm(craft)
+        craft = craft.cuda()
+        craft = torch.nn.parallel.DistributedDataParallel(craft, device_ids=[self.gpu])
+
+        torch.backends.cudnn.benchmark = True
+
+        # DATASET -----------------------------------------------------------------------------------------------------#
+
+        if self.config.train.use_synthtext:
+            trn_syn_loader = self.get_synth_loader()
+            batch_syn = iter(trn_syn_loader)
+
+        if self.config.train.real_dataset == "custom":
+            trn_real_dataset = self.get_custom_dataset()
+        else:
+            raise Exception("Undefined dataset")
+
+        if self.config.mode == "weak_supervision":
+            trn_real_dataset.update_model(supervision_model)
+            trn_real_dataset.update_device(supervision_device)
+
+        trn_real_sampler = torch.utils.data.distributed.DistributedSampler(
+            trn_real_dataset
+        )
+        trn_real_loader = torch.utils.data.DataLoader(
+            trn_real_dataset,
+            batch_size=self.config.train.batch_size,
+            shuffle=False,
+            num_workers=self.config.train.num_workers,
+            sampler=trn_real_sampler,
+            drop_last=False,
+            pin_memory=True,
+        )
+
+        # OPTIMIZER ---------------------------------------------------------------------------------------------------#
+        optimizer = optim.Adam(
+            craft.parameters(),
+            lr=self.config.train.lr,
+            weight_decay=self.config.train.weight_decay,
+        )
+
+        if self.config.train.ckpt_path is not None and self.config.train.st_iter != 0:
+            optimizer.load_state_dict(copyStateDict(self.net_param["optimizer"]))
+            self.config.train.st_iter = self.net_param["optimizer"]["state"][0]["step"]
+            self.config.train.lr = self.net_param["optimizer"]["param_groups"][0]["lr"]
+
+        # LOSS --------------------------------------------------------------------------------------------------------#
+        # mixed precision
+        if self.config.train.amp:
+            scaler = torch.cuda.amp.GradScaler()
+
+            if (
+                self.config.train.ckpt_path is not None
+                and self.config.train.st_iter != 0
+            ):
+                scaler.load_state_dict(copyStateDict(self.net_param["scaler"]))
+        else:
+            scaler = None
+
+        criterion = self.get_loss()
+
+        # TRAIN -------------------------------------------------------------------------------------------------------#
+        train_step = self.config.train.st_iter
+        whole_training_step = self.config.train.end_iter
+        update_lr_rate_step = 0
+        training_lr = self.config.train.lr
+        loss_value = 0
+        batch_time = 0
+        start_time = time.time()
+
+        print(
+            "================================ Train start ================================"
+        )
+        while train_step < whole_training_step:
+            trn_real_sampler.set_epoch(train_step)
+            for (
+                index,
+                (
+                    images,
+                    region_scores,
+                    affinity_scores,
+                    confidence_masks,
+                ),
+            ) in enumerate(trn_real_loader):
+                craft.train()
+                if train_step > 0 and train_step % self.config.train.lr_decay == 0:
+                    update_lr_rate_step += 1
+                    training_lr = self.adjust_learning_rate(
+                        optimizer,
+                        self.config.train.gamma,
+                        update_lr_rate_step,
+                        self.config.train.lr,
+                    )
+
+                images = images.cuda(non_blocking=True)
+                region_scores = region_scores.cuda(non_blocking=True)
+                affinity_scores = affinity_scores.cuda(non_blocking=True)
+                confidence_masks = confidence_masks.cuda(non_blocking=True)
+
+                if self.config.train.use_synthtext:
+                    # Synth image load
+                    syn_image, syn_region_label, syn_affi_label, syn_confidence_mask = next(
+                        batch_syn
+                    )
+                    syn_image = syn_image.cuda(non_blocking=True)
+                    syn_region_label = syn_region_label.cuda(non_blocking=True)
+                    syn_affi_label = syn_affi_label.cuda(non_blocking=True)
+                    syn_confidence_mask = syn_confidence_mask.cuda(non_blocking=True)
+
+                    # concat syn & custom image
+                    images = torch.cat((syn_image, images), 0)
+                    region_image_label = torch.cat(
+                        (syn_region_label, region_scores), 0
+                    )
+                    affinity_image_label = torch.cat((syn_affi_label, affinity_scores), 0)
+                    confidence_mask_label = torch.cat(
+                        (syn_confidence_mask, confidence_masks), 0
+                    )
+                else:
+                    region_image_label = region_scores
+                    affinity_image_label = affinity_scores
+                    confidence_mask_label = confidence_masks
+
+                if self.config.train.amp:
+                    with torch.cuda.amp.autocast():
+
+                        output, _ = craft(images)
+                        out1 = output[:, :, :, 0]
+                        out2 = output[:, :, :, 1]
+
+                        loss = criterion(
+                            region_image_label,
+                            affinity_image_label,
+                            out1,
+                            out2,
+                            confidence_mask_label,
+                            self.config.train.neg_rto,
+                            self.config.train.n_min_neg,
+                        )
+
+                    optimizer.zero_grad()
+                    scaler.scale(loss).backward()
+                    scaler.step(optimizer)
+                    scaler.update()
+
+                else:
+                    output, _ = craft(images)
+                    out1 = output[:, :, :, 0]
+                    out2 = output[:, :, :, 1]
+                    loss = criterion(
+                        region_image_label,
+                        affinity_image_label,
+                        out1,
+                        out2,
+                        confidence_mask_label,
+                        self.config.train.neg_rto,
+                    )
+
+                    optimizer.zero_grad()
+                    loss.backward()
+                    optimizer.step()
+
+                end_time = time.time()
+                loss_value += loss.item()
+                batch_time += end_time - start_time
+
+                if train_step > 0 and train_step % 5 == 0 and self.gpu == 0:
+                    mean_loss = loss_value / 5
+                    loss_value = 0
+                    avg_batch_time = batch_time / 5
+                    batch_time = 0
+
+                    print(
+                        "{}, training_step: {}|{}, learning rate: {:.8f}, "
+                        "training_loss: {:.5f}, avg_batch_time: {:.5f}".format(
+                            time.strftime(
+                                "%Y-%m-%d:%H:%M:%S", time.localtime(time.time())
+                            ),
+                            train_step,
+                            whole_training_step,
+                            training_lr,
+                            mean_loss,
+                            avg_batch_time,
+                        )
+                    )
+
+                    if self.gpu == 0 and self.config.wandb_opt:
+                        wandb.log({"train_step": train_step, "mean_loss": mean_loss})
+
+                if (
+                    train_step % self.config.train.eval_interval == 0
+                    and train_step != 0
+                ):
+
+                    craft.eval()
+                    # initialize all buffer value with zero
+                    if self.gpu == 0:
+                        for buffer in buffer_dict.values():
+                            for i in range(len(buffer)):
+                                buffer[i] = None
+
+                        print("Saving state, index:", train_step)
+                        save_param_dic = {
+                            "iter": train_step,
+                            "craft": craft.state_dict(),
+                            "optimizer": optimizer.state_dict(),
+                        }
+                        save_param_path = (
+                            self.config.results_dir
+                            + "/CRAFT_clr_"
+                            + repr(train_step)
+                            + ".pth"
+                        )
+
+                        if self.config.train.amp:
+                            save_param_dic["scaler"] = scaler.state_dict()
+                            save_param_path = (
+                                self.config.results_dir
+                                + "/CRAFT_clr_amp_"
+                                + repr(train_step)
+                                + ".pth"
+                            )
+
+                        torch.save(save_param_dic, save_param_path)
+
+                    # validation
+                    self.iou_eval(
+                        "custom_data",
+                        train_step,
+                        buffer_dict["custom_data"],
+                        craft,
+                    )
+
+                train_step += 1
+                if train_step >= whole_training_step:
+                    break
+
+            if self.config.mode == "weak_supervision":
+                state_dict = craft.module.state_dict()
+                supervision_model.load_state_dict(state_dict)
+                trn_real_dataset.update_model(supervision_model)
+
+        # save last model
+        if self.gpu == 0:
+            save_param_dic = {
+                "iter": train_step,
+                "craft": craft.state_dict(),
+                "optimizer": optimizer.state_dict(),
+            }
+            save_param_path = (
+                self.config.results_dir + "/CRAFT_clr_" + repr(train_step) + ".pth"
+            )
+
+            if self.config.train.amp:
+                save_param_dic["scaler"] = scaler.state_dict()
+                save_param_path = (
+                    self.config.results_dir
+                    + "/CRAFT_clr_amp_"
+                    + repr(train_step)
+                    + ".pth"
+                )
+            torch.save(save_param_dic, save_param_path)
+
+def main():
+    parser = argparse.ArgumentParser(description="CRAFT custom data train")
+    parser.add_argument(
+        "--yaml",
+        "--yaml_file_name",
+        default="custom_data_train",
+        type=str,
+        help="Load configuration",
+    )
+    parser.add_argument(
+        "--port", "--use ddp port", default="2346", type=str, help="Port number"
+    )
+
+    args = parser.parse_args()
+
+    # load configure
+    exp_name = args.yaml
+    config = load_yaml(args.yaml)
+
+    print("-" * 20 + " Options " + "-" * 20)
+    print(yaml.dump(config))
+    print("-" * 40)
+
+    # Make result_dir
+    res_dir = os.path.join(config["results_dir"], args.yaml)
+    config["results_dir"] = res_dir
+    if not os.path.exists(res_dir):
+        os.makedirs(res_dir)
+
+    # Duplicate yaml file to result_dir
+    shutil.copy(
+        "config/" + args.yaml + ".yaml", os.path.join(res_dir, args.yaml) + ".yaml"
+    )
+
+    if config["mode"] == "weak_supervision":
+        # NOTE: half GPU assign train / half GPU assign supervision setting
+        ngpus_per_node = torch.cuda.device_count() // 2
+        mode = "weak_supervision"
+    else:
+        ngpus_per_node = torch.cuda.device_count()
+        mode = None
+
+    print(f"Total process num : {ngpus_per_node}")
+
+    manager = mp.Manager()
+    buffer1 = manager.list([None] * config["test"]["custom_data"]["test_set_size"])
+
+    buffer_dict = {"custom_data": buffer1}
+    torch.multiprocessing.spawn(
+        main_worker,
+        nprocs=ngpus_per_node,
+        args=(args.port, ngpus_per_node, config, buffer_dict, exp_name, mode,),
+    )
+
+
+def main_worker(gpu, port, ngpus_per_node, config, buffer_dict, exp_name, mode):
+
+    torch.distributed.init_process_group(
+        backend="nccl",
+        init_method="tcp://127.0.0.1:" + port,
+        world_size=ngpus_per_node,
+        rank=gpu,
+    )
+
+    # Apply config to wandb
+    if gpu == 0 and config["wandb_opt"]:
+        wandb.init(project="craft-stage2", entity="user_name", name=exp_name)
+        wandb.config.update(config)
+
+    batch_size = int(config["train"]["batch_size"] / ngpus_per_node)
+    config["train"]["batch_size"] = batch_size
+    config = DotDict(config)
+
+    # Start train
+    trainer = Trainer(config, gpu, mode)
+    trainer.train(buffer_dict)
+
+    if gpu == 0:
+        if config["wandb_opt"]:
+            wandb.finish()
+
+    torch.distributed.barrier()
+    torch.distributed.destroy_process_group()
+
+if __name__ == "__main__":
+    main()

trainer/craft/utils/craft_utils.py

@@ -0,0 +1,345 @@
+
+# -*- coding: utf-8 -*-
+import os
+import torch
+import cv2
+import math
+import numpy as np
+from data import imgproc
+
+""" auxilary functions """
+# unwarp corodinates
+
+
+
+
+def warpCoord(Minv, pt):
+    out = np.matmul(Minv, (pt[0], pt[1], 1))
+    return np.array([out[0]/out[2], out[1]/out[2]])
+""" end of auxilary functions """
+
+def test():
+    print('pass')
+    
+
+def getDetBoxes_core(textmap, linkmap, text_threshold, link_threshold, low_text):
+    # prepare data
+    linkmap = linkmap.copy()
+    textmap = textmap.copy()
+    img_h, img_w = textmap.shape
+
+    """ labeling method """
+    ret, text_score = cv2.threshold(textmap, low_text, 1, 0)
+    ret, link_score = cv2.threshold(linkmap, link_threshold, 1, 0)
+
+    text_score_comb = np.clip(text_score + link_score, 0, 1)
+    nLabels, labels, stats, centroids = \
+        cv2.connectedComponentsWithStats(text_score_comb.astype(np.uint8), connectivity=4)
+
+    det = []
+    mapper = []
+    for k in range(1,nLabels):
+        # size filtering
+        size = stats[k, cv2.CC_STAT_AREA]
+        if size < 10: continue
+
+        # thresholding
+        if np.max(textmap[labels==k]) < text_threshold: continue
+
+        # make segmentation map
+        segmap = np.zeros(textmap.shape, dtype=np.uint8)
+        segmap[labels==k] = 255
+        segmap[np.logical_and(link_score==1, text_score==0)] = 0   # remove link area
+        x, y = stats[k, cv2.CC_STAT_LEFT], stats[k, cv2.CC_STAT_TOP]
+        w, h = stats[k, cv2.CC_STAT_WIDTH], stats[k, cv2.CC_STAT_HEIGHT]
+        niter = int(math.sqrt(size * min(w, h) / (w * h)) * 2)
+        sx, ex, sy, ey = x - niter, x + w + niter + 1, y - niter, y + h + niter + 1
+        # boundary check
+        if sx < 0 : sx = 0
+        if sy < 0 : sy = 0
+        if ex >= img_w: ex = img_w
+        if ey >= img_h: ey = img_h
+        kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(1 + niter, 1 + niter))
+        segmap[sy:ey, sx:ex] = cv2.dilate(segmap[sy:ey, sx:ex], kernel, iterations=1)
+        #kernel1 = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 5))
+        #segmap[sy:ey, sx:ex] = cv2.dilate(segmap[sy:ey, sx:ex], kernel1, iterations=1)
+
+
+        # make box
+        np_contours = np.roll(np.array(np.where(segmap!=0)),1,axis=0).transpose().reshape(-1,2)
+        rectangle = cv2.minAreaRect(np_contours)
+        box = cv2.boxPoints(rectangle)
+
+        # align diamond-shape
+        w, h = np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[1] - box[2])
+        box_ratio = max(w, h) / (min(w, h) + 1e-5)
+        if abs(1 - box_ratio) <= 0.1:
+            l, r = min(np_contours[:,0]), max(np_contours[:,0])
+            t, b = min(np_contours[:,1]), max(np_contours[:,1])
+            box = np.array([[l, t], [r, t], [r, b], [l, b]], dtype=np.float32)
+
+        # make clock-wise order
+        startidx = box.sum(axis=1).argmin()
+        box = np.roll(box, 4-startidx, 0)
+        box = np.array(box)
+
+        det.append(box)
+        mapper.append(k)
+
+    return det, labels, mapper
+
+def getPoly_core(boxes, labels, mapper, linkmap):
+    # configs
+    num_cp = 5
+    max_len_ratio = 0.7
+    expand_ratio = 1.45
+    max_r = 2.0
+    step_r = 0.2
+
+    polys = []  
+    for k, box in enumerate(boxes):
+        # size filter for small instance
+        w, h = int(np.linalg.norm(box[0] - box[1]) + 1), int(np.linalg.norm(box[1] - box[2]) + 1)
+        if w < 30 or h < 30:
+            polys.append(None); continue
+
+        # warp image
+        tar = np.float32([[0,0],[w,0],[w,h],[0,h]])
+        M = cv2.getPerspectiveTransform(box, tar)
+        word_label = cv2.warpPerspective(labels, M, (w, h), flags=cv2.INTER_NEAREST)
+        try:
+            Minv = np.linalg.inv(M)
+        except:
+            polys.append(None); continue
+
+        # binarization for selected label
+        cur_label = mapper[k]
+        word_label[word_label != cur_label] = 0
+        word_label[word_label > 0] = 1
+
+        """ Polygon generation """
+        # find top/bottom contours
+        cp = []
+        max_len = -1
+        for i in range(w):
+            region = np.where(word_label[:,i] != 0)[0]
+            if len(region) < 2 : continue
+            cp.append((i, region[0], region[-1]))
+            length = region[-1] - region[0] + 1
+            if length > max_len: max_len = length
+
+        # pass if max_len is similar to h
+        if h * max_len_ratio < max_len:
+            polys.append(None); continue
+
+        # get pivot points with fixed length
+        tot_seg = num_cp * 2 + 1
+        seg_w = w / tot_seg     # segment width
+        pp = [None] * num_cp    # init pivot points
+        cp_section = [[0, 0]] * tot_seg
+        seg_height = [0] * num_cp
+        seg_num = 0
+        num_sec = 0
+        prev_h = -1
+        for i in range(0,len(cp)):
+            (x, sy, ey) = cp[i]
+            if (seg_num + 1) * seg_w <= x and seg_num <= tot_seg:
+                # average previous segment
+                if num_sec == 0: break
+                cp_section[seg_num] = [cp_section[seg_num][0] / num_sec, cp_section[seg_num][1] / num_sec]
+                num_sec = 0
+
+                # reset variables
+                seg_num += 1
+                prev_h = -1
+
+            # accumulate center points
+            cy = (sy + ey) * 0.5
+            cur_h = ey - sy + 1
+            cp_section[seg_num] = [cp_section[seg_num][0] + x, cp_section[seg_num][1] + cy]
+            num_sec += 1
+
+            if seg_num % 2 == 0: continue # No polygon area
+
+            if prev_h < cur_h:
+                pp[int((seg_num - 1)/2)] = (x, cy)
+                seg_height[int((seg_num - 1)/2)] = cur_h
+                prev_h = cur_h
+
+        # processing last segment
+        if num_sec != 0:
+            cp_section[-1] = [cp_section[-1][0] / num_sec, cp_section[-1][1] / num_sec]
+
+        # pass if num of pivots is not sufficient or segment widh is smaller than character height 
+        if None in pp or seg_w < np.max(seg_height) * 0.25:
+            polys.append(None); continue
+
+        # calc median maximum of pivot points
+        half_char_h = np.median(seg_height) * expand_ratio / 2
+
+        # calc gradiant and apply to make horizontal pivots
+        new_pp = []
+        for i, (x, cy) in enumerate(pp):
+            dx = cp_section[i * 2 + 2][0] - cp_section[i * 2][0]
+            dy = cp_section[i * 2 + 2][1] - cp_section[i * 2][1]
+            if dx == 0:     # gradient if zero
+                new_pp.append([x, cy - half_char_h, x, cy + half_char_h])
+                continue
+            rad = - math.atan2(dy, dx)
+            c, s = half_char_h * math.cos(rad), half_char_h * math.sin(rad)
+            new_pp.append([x - s, cy - c, x + s, cy + c])
+
+        # get edge points to cover character heatmaps
+        isSppFound, isEppFound = False, False
+        grad_s = (pp[1][1] - pp[0][1]) / (pp[1][0] - pp[0][0]) + (pp[2][1] - pp[1][1]) / (pp[2][0] - pp[1][0])
+        grad_e = (pp[-2][1] - pp[-1][1]) / (pp[-2][0] - pp[-1][0]) + (pp[-3][1] - pp[-2][1]) / (pp[-3][0] - pp[-2][0])
+        for r in np.arange(0.5, max_r, step_r):
+            dx = 2 * half_char_h * r
+            if not isSppFound:
+                line_img = np.zeros(word_label.shape, dtype=np.uint8)
+                dy = grad_s * dx
+                p = np.array(new_pp[0]) - np.array([dx, dy, dx, dy])
+                cv2.line(line_img, (int(p[0]), int(p[1])), (int(p[2]), int(p[3])), 1, thickness=1)
+                if np.sum(np.logical_and(word_label, line_img)) == 0 or r + 2 * step_r >= max_r:
+                    spp = p
+                    isSppFound = True
+            if not isEppFound:
+                line_img = np.zeros(word_label.shape, dtype=np.uint8)
+                dy = grad_e * dx
+                p = np.array(new_pp[-1]) + np.array([dx, dy, dx, dy])
+                cv2.line(line_img, (int(p[0]), int(p[1])), (int(p[2]), int(p[3])), 1, thickness=1)
+                if np.sum(np.logical_and(word_label, line_img)) == 0 or r + 2 * step_r >= max_r:
+                    epp = p
+                    isEppFound = True
+            if isSppFound and isEppFound:
+                break
+
+        # pass if boundary of polygon is not found
+        if not (isSppFound and isEppFound):
+            polys.append(None); continue
+
+        # make final polygon
+        poly = []
+        poly.append(warpCoord(Minv, (spp[0], spp[1])))
+        for p in new_pp:
+            poly.append(warpCoord(Minv, (p[0], p[1])))
+        poly.append(warpCoord(Minv, (epp[0], epp[1])))
+        poly.append(warpCoord(Minv, (epp[2], epp[3])))
+        for p in reversed(new_pp):
+            poly.append(warpCoord(Minv, (p[2], p[3])))
+        poly.append(warpCoord(Minv, (spp[2], spp[3])))
+
+        # add to final result
+        polys.append(np.array(poly))
+
+    return polys
+
+def getDetBoxes(textmap, linkmap, text_threshold, link_threshold, low_text, poly=False):
+    boxes, labels, mapper = getDetBoxes_core(textmap, linkmap, text_threshold, link_threshold, low_text)
+
+    if poly:
+        polys = getPoly_core(boxes, labels, mapper, linkmap)
+    else:
+        polys = [None] * len(boxes)
+
+    return boxes, polys
+
+def adjustResultCoordinates(polys, ratio_w, ratio_h, ratio_net = 2):
+    if len(polys) > 0:
+        polys = np.array(polys)
+        for k in range(len(polys)):
+            if polys[k] is not None:
+                polys[k] *= (ratio_w * ratio_net, ratio_h * ratio_net)
+    return polys
+
+def save_outputs(image, region_scores, affinity_scores, text_threshold, link_threshold,
+                                           low_text, outoput_path, confidence_mask = None):
+    """save image, region_scores, and affinity_scores in a single image. region_scores and affinity_scores must be
+    cpu numpy arrays. You can convert GPU Tensors to CPU numpy arrays like this:
+    >>> array = tensor.cpu().data.numpy()
+    When saving outputs of the network during training, make sure you convert ALL tensors (image, region_score,
+    affinity_score) to numpy array first.
+    :param image: numpy array
+    :param region_scores: [] 2D numpy array with each element between 0~1.
+    :param affinity_scores: same as region_scores
+    :param text_threshold: 0 ~ 1. Closer to 0, characters with lower confidence will also be considered a word and be boxed
+    :param link_threshold: 0 ~ 1. Closer to 0, links with lower confidence will also be considered a word and be boxed
+    :param low_text: 0 ~ 1. Closer to 0, boxes will be more loosely drawn.
+    :param outoput_path:
+    :param confidence_mask:
+    :return:
+    """
+
+    assert region_scores.shape == affinity_scores.shape
+    assert len(image.shape) - 1 == len(region_scores.shape)
+
+    boxes, polys = getDetBoxes(region_scores, affinity_scores, text_threshold, link_threshold,
+                                           low_text, False)
+    boxes = np.array(boxes, np.int32) * 2
+    if len(boxes) > 0:
+        np.clip(boxes[:, :, 0], 0, image.shape[1])
+        np.clip(boxes[:, :, 1], 0, image.shape[0])
+        for box in boxes:
+            cv2.polylines(image, [np.reshape(box, (-1, 1, 2))], True, (0, 0, 255))
+
+    target_gaussian_heatmap_color = imgproc.cvt2HeatmapImg(region_scores)
+    target_gaussian_affinity_heatmap_color = imgproc.cvt2HeatmapImg(affinity_scores)
+
+    if confidence_mask is not None:
+        confidence_mask_gray = imgproc.cvt2HeatmapImg(confidence_mask)
+        gt_scores = np.hstack([target_gaussian_heatmap_color, target_gaussian_affinity_heatmap_color])
+        confidence_mask_gray = np.hstack([np.zeros_like(confidence_mask_gray), confidence_mask_gray])
+        output = np.concatenate([gt_scores, confidence_mask_gray], axis=0)
+        output = np.hstack([image, output])
+
+    else:
+        gt_scores = np.concatenate([target_gaussian_heatmap_color, target_gaussian_affinity_heatmap_color], axis=0)
+        output = np.hstack([image, gt_scores])
+
+    cv2.imwrite(outoput_path, output)
+    return output
+
+
+def save_outputs_from_tensors(images, region_scores, affinity_scores, text_threshold, link_threshold,
+                                           low_text, output_dir, image_names, confidence_mask = None):
+
+    """takes images, region_scores, and affinity_scores as tensors (cab be GPU).
+    :param images: 4D tensor
+    :param region_scores: 3D tensor with values between 0 ~ 1
+    :param affinity_scores: 3D tensor with values between 0 ~ 1
+    :param text_threshold:
+    :param link_threshold:
+    :param low_text:
+    :param output_dir: direcotry to save the output images. Will be joined with base names of image_names
+    :param image_names: names of each image. Doesn't have to be the base name (image file names)
+    :param confidence_mask:
+    :return:
+    """
+    #import ipdb;ipdb.set_trace()
+    #images = images.cpu().permute(0, 2, 3, 1).contiguous().data.numpy()
+    if type(images) == torch.Tensor:
+        images = np.array(images)
+
+    region_scores = region_scores.cpu().data.numpy()
+    affinity_scores = affinity_scores.cpu().data.numpy()
+
+    batch_size = images.shape[0]
+    assert batch_size == region_scores.shape[0] and batch_size == affinity_scores.shape[0] and batch_size == len(image_names), \
+        "The first dimension (i.e. batch size) of images, region scores, and affinity scores must be equal"
+
+    output_images = []
+
+    for i in range(batch_size):
+        image = images[i]
+        region_score = region_scores[i]
+        affinity_score = affinity_scores[i]
+
+        image_name = os.path.basename(image_names[i])
+        outoput_path = os.path.join(output_dir,image_name)
+
+        output_image = save_outputs(image, region_score, affinity_score, text_threshold, link_threshold,
+                                           low_text, outoput_path, confidence_mask=confidence_mask)
+
+        output_images.append(output_image)
+
+    return output_images

trainer/craft/utils/inference_boxes.py

@@ -0,0 +1,361 @@
+import os
+import re
+import itertools
+
+import cv2
+import time
+import numpy as np
+import torch
+from torch.autograd import Variable
+
+from utils.craft_utils import getDetBoxes, adjustResultCoordinates
+from data import imgproc
+from data.dataset import SynthTextDataSet
+import math
+import xml.etree.ElementTree as elemTree
+
+
+#-------------------------------------------------------------------------------------------------------------------#
+def rotatePoint(xc, yc, xp, yp, theta):
+    xoff = xp - xc
+    yoff = yp - yc
+
+    cosTheta = math.cos(theta)
+    sinTheta = math.sin(theta)
+    pResx = cosTheta * xoff + sinTheta * yoff
+    pResy = - sinTheta * xoff + cosTheta * yoff
+    # pRes = (xc + pResx, yc + pResy)
+    return int(xc + pResx), int(yc + pResy)
+
+def addRotatedShape(cx, cy, w, h, angle):
+    p0x, p0y = rotatePoint(cx, cy, cx - w / 2, cy - h / 2, -angle)
+    p1x, p1y = rotatePoint(cx, cy, cx + w / 2, cy - h / 2, -angle)
+    p2x, p2y = rotatePoint(cx, cy, cx + w / 2, cy + h / 2, -angle)
+    p3x, p3y = rotatePoint(cx, cy, cx - w / 2, cy + h / 2, -angle)
+
+    points = [[p0x, p0y], [p1x, p1y], [p2x, p2y], [p3x, p3y]]
+
+    return points
+
+def xml_parsing(xml):
+    tree = elemTree.parse(xml)
+
+    annotations = []  # Initialize the list to store labels
+    iter_element = tree.iter(tag="object")
+
+    for element in iter_element:
+        annotation = {}  # Initialize the dict to store labels
+
+        annotation['name'] = element.find("name").text  # Save the name tag value
+
+        box_coords = element.iter(tag="robndbox")
+
+        for box_coord in box_coords:
+            cx = float(box_coord.find("cx").text)
+            cy = float(box_coord.find("cy").text)
+            w = float(box_coord.find("w").text)
+            h = float(box_coord.find("h").text)
+            angle = float(box_coord.find("angle").text)
+
+            convertcoodi = addRotatedShape(cx, cy, w, h, angle)
+
+            annotation['box_coodi'] = convertcoodi
+            annotations.append(annotation)
+
+        box_coords = element.iter(tag="bndbox")
+
+        for box_coord in box_coords:
+            xmin = int(box_coord.find("xmin").text)
+            ymin = int(box_coord.find("ymin").text)
+            xmax = int(box_coord.find("xmax").text)
+            ymax = int(box_coord.find("ymax").text)
+            # annotation['bndbox'] = [xmin,ymin,xmax,ymax]
+
+            annotation['box_coodi'] = [[xmin, ymin], [xmax, ymin], [xmax, ymax],
+                                       [xmin, ymax]]
+            annotations.append(annotation)
+
+
+
+
+    bounds = []
+    for i in range(len(annotations)):
+        box_info_dict = {"points": None, "text": None, "ignore": None}
+
+        box_info_dict["points"] = np.array(annotations[i]['box_coodi'])
+        if annotations[i]['name'] == "dnc":
+            box_info_dict["text"] = "###"
+            box_info_dict["ignore"] = True
+        else:
+            box_info_dict["text"] = annotations[i]['name']
+            box_info_dict["ignore"] = False
+
+        bounds.append(box_info_dict)
+
+
+
+    return bounds
+
+#-------------------------------------------------------------------------------------------------------------------#
+
+def load_prescription_gt(dataFolder):
+
+
+    total_img_path = []
+    total_imgs_bboxes = []
+    for (root, directories, files) in os.walk(dataFolder):
+        for file in files:
+            if '.jpg' in file:
+                img_path = os.path.join(root, file)
+                total_img_path.append(img_path)
+            if '.xml' in file:
+                gt_path = os.path.join(root, file)
+                total_imgs_bboxes.append(gt_path)
+
+
+    total_imgs_parsing_bboxes = []
+    for img_path, bbox in zip(sorted(total_img_path), sorted(total_imgs_bboxes)):
+        # check file
+
+        assert img_path.split(".jpg")[0] == bbox.split(".xml")[0]
+
+        result_label = xml_parsing(bbox)
+        total_imgs_parsing_bboxes.append(result_label)
+
+
+    return total_imgs_parsing_bboxes, sorted(total_img_path)
+
+
+# NOTE
+def load_prescription_cleval_gt(dataFolder):
+
+
+    total_img_path = []
+    total_gt_path = []
+    for (root, directories, files) in os.walk(dataFolder):
+        for file in files:
+            if '.jpg' in file:
+                img_path = os.path.join(root, file)
+                total_img_path.append(img_path)
+            if '_cl.txt' in file:
+                gt_path = os.path.join(root, file)
+                total_gt_path.append(gt_path)
+
+
+    total_imgs_parsing_bboxes = []
+    for img_path, gt_path in zip(sorted(total_img_path), sorted(total_gt_path)):
+        # check file
+
+        assert img_path.split(".jpg")[0] == gt_path.split('_label_cl.txt')[0]
+
+        lines = open(gt_path, encoding="utf-8").readlines()
+        word_bboxes = []
+
+        for line in lines:
+            box_info_dict = {"points": None, "text": None, "ignore": None}
+            box_info = line.strip().encode("utf-8").decode("utf-8-sig").split(",")
+
+            box_points = [int(box_info[i]) for i in range(8)]
+            box_info_dict["points"] = np.array(box_points)
+
+            word_bboxes.append(box_info_dict)
+        total_imgs_parsing_bboxes.append(word_bboxes)
+
+    return total_imgs_parsing_bboxes, sorted(total_img_path)
+
+
+def load_synthtext_gt(data_folder):
+
+    synth_dataset = SynthTextDataSet(
+        output_size=768, data_dir=data_folder, saved_gt_dir=data_folder, logging=False
+    )
+    img_names, img_bbox, img_words = synth_dataset.load_data(bbox="word")
+
+    total_img_path = []
+    total_imgs_bboxes = []
+    for index in range(len(img_bbox[:100])):
+        img_path = os.path.join(data_folder, img_names[index][0])
+        total_img_path.append(img_path)
+        try:
+            wordbox = img_bbox[index].transpose((2, 1, 0))
+        except:
+            wordbox = np.expand_dims(img_bbox[index], axis=0)
+            wordbox = wordbox.transpose((0, 2, 1))
+
+        words = [re.split(" \n|\n |\n| ", t.strip()) for t in img_words[index]]
+        words = list(itertools.chain(*words))
+        words = [t for t in words if len(t) > 0]
+
+        if len(words) != len(wordbox):
+            import ipdb
+
+            ipdb.set_trace()
+
+        single_img_bboxes = []
+        for j in range(len(words)):
+            box_info_dict = {"points": None, "text": None, "ignore": None}
+            box_info_dict["points"] = wordbox[j]
+            box_info_dict["text"] = words[j]
+            box_info_dict["ignore"] = False
+            single_img_bboxes.append(box_info_dict)
+
+        total_imgs_bboxes.append(single_img_bboxes)
+
+    return total_imgs_bboxes, total_img_path
+
+
+def load_icdar2015_gt(dataFolder, isTraing=False):
+    if isTraing:
+        img_folderName = "ch4_training_images"
+        gt_folderName = "ch4_training_localization_transcription_gt"
+    else:
+        img_folderName = "ch4_test_images"
+        gt_folderName = "ch4_test_localization_transcription_gt"
+
+    gt_folder_path = os.listdir(os.path.join(dataFolder, gt_folderName))
+    total_imgs_bboxes = []
+    total_img_path = []
+    for gt_path in gt_folder_path:
+        gt_path = os.path.join(os.path.join(dataFolder, gt_folderName), gt_path)
+        img_path = (
+            gt_path.replace(gt_folderName, img_folderName)
+            .replace(".txt", ".jpg")
+            .replace("gt_", "")
+        )
+        image = cv2.imread(img_path)
+        lines = open(gt_path, encoding="utf-8").readlines()
+        single_img_bboxes = []
+        for line in lines:
+            box_info_dict = {"points": None, "text": None, "ignore": None}
+
+            box_info = line.strip().encode("utf-8").decode("utf-8-sig").split(",")
+            box_points = [int(box_info[j]) for j in range(8)]
+            word = box_info[8:]
+            word = ",".join(word)
+            box_points = np.array(box_points, np.int32).reshape(4, 2)
+            cv2.polylines(
+                image, [np.array(box_points).astype(np.int)], True, (0, 0, 255), 1
+            )
+            box_info_dict["points"] = box_points
+            box_info_dict["text"] = word
+            if word == "###":
+                box_info_dict["ignore"] = True
+            else:
+                box_info_dict["ignore"] = False
+
+            single_img_bboxes.append(box_info_dict)
+        total_imgs_bboxes.append(single_img_bboxes)
+        total_img_path.append(img_path)
+    return total_imgs_bboxes, total_img_path
+
+
+def load_icdar2013_gt(dataFolder, isTraing=False):
+
+    # choose test dataset
+    if isTraing:
+        img_folderName = "Challenge2_Test_Task12_Images"
+        gt_folderName = "Challenge2_Test_Task1_GT"
+    else:
+        img_folderName = "Challenge2_Test_Task12_Images"
+        gt_folderName = "Challenge2_Test_Task1_GT"
+
+    gt_folder_path = os.listdir(os.path.join(dataFolder, gt_folderName))
+
+    total_imgs_bboxes = []
+    total_img_path = []
+    for gt_path in gt_folder_path:
+        gt_path = os.path.join(os.path.join(dataFolder, gt_folderName), gt_path)
+        img_path = (
+            gt_path.replace(gt_folderName, img_folderName)
+            .replace(".txt", ".jpg")
+            .replace("gt_", "")
+        )
+        image = cv2.imread(img_path)
+        lines = open(gt_path, encoding="utf-8").readlines()
+        single_img_bboxes = []
+        for line in lines:
+            box_info_dict = {"points": None, "text": None, "ignore": None}
+
+            box_info = line.strip().encode("utf-8").decode("utf-8-sig").split(",")
+            box = [int(box_info[j]) for j in range(4)]
+            word = box_info[4:]
+            word = ",".join(word)
+            box = [
+                [box[0], box[1]],
+                [box[2], box[1]],
+                [box[2], box[3]],
+                [box[0], box[3]],
+            ]
+
+            box_info_dict["points"] = box
+            box_info_dict["text"] = word
+            if word == "###":
+                box_info_dict["ignore"] = True
+            else:
+                box_info_dict["ignore"] = False
+
+            single_img_bboxes.append(box_info_dict)
+
+        total_imgs_bboxes.append(single_img_bboxes)
+        total_img_path.append(img_path)
+
+    return total_imgs_bboxes, total_img_path
+
+
+def test_net(
+    net,
+    image,
+    text_threshold,
+    link_threshold,
+    low_text,
+    cuda,
+    poly,
+    canvas_size=1280,
+    mag_ratio=1.5,
+):
+    # resize
+
+    img_resized, target_ratio, size_heatmap = imgproc.resize_aspect_ratio(
+        image, canvas_size, interpolation=cv2.INTER_LINEAR, mag_ratio=mag_ratio
+    )
+    ratio_h = ratio_w = 1 / target_ratio
+
+    # preprocessing
+    x = imgproc.normalizeMeanVariance(img_resized)
+    x = torch.from_numpy(x).permute(2, 0, 1)  # [h, w, c] to [c, h, w]
+    x = Variable(x.unsqueeze(0))  # [c, h, w] to [b, c, h, w]
+    if cuda:
+        x = x.cuda()
+
+    # forward pass
+    with torch.no_grad():
+        y, feature = net(x)
+
+    # make score and link map
+    score_text = y[0, :, :, 0].cpu().data.numpy().astype(np.float32)
+    score_link = y[0, :, :, 1].cpu().data.numpy().astype(np.float32)
+
+    # NOTE
+    score_text = score_text[: size_heatmap[0], : size_heatmap[1]]
+    score_link = score_link[: size_heatmap[0], : size_heatmap[1]]
+
+    # Post-processing
+    boxes, polys = getDetBoxes(
+        score_text, score_link, text_threshold, link_threshold, low_text, poly
+    )
+
+    # coordinate adjustment
+    boxes = adjustResultCoordinates(boxes, ratio_w, ratio_h)
+    polys = adjustResultCoordinates(polys, ratio_w, ratio_h)
+    for k in range(len(polys)):
+        if polys[k] is None:
+            polys[k] = boxes[k]
+
+    # render results (optional)
+    score_text = score_text.copy()
+    render_score_text = imgproc.cvt2HeatmapImg(score_text)
+    render_score_link = imgproc.cvt2HeatmapImg(score_link)
+    render_img = [render_score_text, render_score_link]
+    # ret_score_text = imgproc.cvt2HeatmapImg(render_img)
+
+    return boxes, polys, render_img

trainer/craft/utils/util.py

@@ -0,0 +1,142 @@
+from collections import OrderedDict
+import os
+
+import cv2
+import numpy as np
+
+from data import imgproc
+from utils import craft_utils
+
+
+def copyStateDict(state_dict):
+    if list(state_dict.keys())[0].startswith("module"):
+        start_idx = 1
+    else:
+        start_idx = 0
+    new_state_dict = OrderedDict()
+    for k, v in state_dict.items():
+        name = ".".join(k.split(".")[start_idx:])
+        new_state_dict[name] = v
+    return new_state_dict
+
+
+def saveInput(
+    imagename, vis_dir, image, region_scores, affinity_scores, confidence_mask
+):
+    image = np.uint8(image.copy())
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+
+    boxes, polys = craft_utils.getDetBoxes(
+        region_scores, affinity_scores, 0.85, 0.2, 0.5, False
+    )
+
+    if image.shape[0] / region_scores.shape[0] >= 2:
+        boxes = np.array(boxes, np.int32) * 2
+    else:
+        boxes = np.array(boxes, np.int32)
+
+    if len(boxes) > 0:
+        np.clip(boxes[:, :, 0], 0, image.shape[1])
+        np.clip(boxes[:, :, 1], 0, image.shape[0])
+        for box in boxes:
+            cv2.polylines(image, [np.reshape(box, (-1, 1, 2))], True, (0, 0, 255))
+    target_gaussian_heatmap_color = imgproc.cvt2HeatmapImg(region_scores)
+    target_gaussian_affinity_heatmap_color = imgproc.cvt2HeatmapImg(affinity_scores)
+    confidence_mask_gray = imgproc.cvt2HeatmapImg(confidence_mask)
+
+    # overlay
+    height, width, channel = image.shape
+    overlay_region = cv2.resize(target_gaussian_heatmap_color, (width, height))
+    overlay_aff = cv2.resize(target_gaussian_affinity_heatmap_color, (width, height))
+    confidence_mask_gray = cv2.resize(
+        confidence_mask_gray, (width, height), interpolation=cv2.INTER_NEAREST
+    )
+    overlay_region = cv2.addWeighted(image, 0.4, overlay_region, 0.6, 5)
+    overlay_aff = cv2.addWeighted(image, 0.4, overlay_aff, 0.7, 6)
+
+    gt_scores = np.concatenate([overlay_region, overlay_aff], axis=1)
+
+    output = np.concatenate([gt_scores, confidence_mask_gray], axis=1)
+
+    output = np.hstack([image, output])
+
+    # synthtext
+    if type(imagename) is not str:
+        imagename = imagename[0].split("/")[-1][:-4]
+
+    outpath = vis_dir + f"/{imagename}_input.jpg"
+    if not os.path.exists(os.path.dirname(outpath)):
+        os.makedirs(os.path.dirname(outpath), exist_ok=True)
+    cv2.imwrite(outpath, output)
+    # print(f'Logging train input into {outpath}')
+
+
+def saveImage(
+    imagename,
+    vis_dir,
+    image,
+    bboxes,
+    affi_bboxes,
+    region_scores,
+    affinity_scores,
+    confidence_mask,
+):
+    output_image = np.uint8(image.copy())
+    output_image = cv2.cvtColor(output_image, cv2.COLOR_RGB2BGR)
+    if len(bboxes) > 0:
+        for i in range(len(bboxes)):
+            _bboxes = np.int32(bboxes[i])
+            for j in range(_bboxes.shape[0]):
+                cv2.polylines(
+                    output_image,
+                    [np.reshape(_bboxes[j], (-1, 1, 2))],
+                    True,
+                    (0, 0, 255),
+                )
+
+        for i in range(len(affi_bboxes)):
+            cv2.polylines(
+                output_image,
+                [np.reshape(affi_bboxes[i].astype(np.int32), (-1, 1, 2))],
+                True,
+                (255, 0, 0),
+            )
+
+    target_gaussian_heatmap_color = imgproc.cvt2HeatmapImg(region_scores)
+    target_gaussian_affinity_heatmap_color = imgproc.cvt2HeatmapImg(affinity_scores)
+    confidence_mask_gray = imgproc.cvt2HeatmapImg(confidence_mask)
+
+    # overlay
+    height, width, channel = image.shape
+    overlay_region = cv2.resize(target_gaussian_heatmap_color, (width, height))
+    overlay_aff = cv2.resize(target_gaussian_affinity_heatmap_color, (width, height))
+
+    overlay_region = cv2.addWeighted(image.copy(), 0.4, overlay_region, 0.6, 5)
+    overlay_aff = cv2.addWeighted(image.copy(), 0.4, overlay_aff, 0.6, 5)
+
+    heat_map = np.concatenate([overlay_region, overlay_aff], axis=1)
+
+    # synthtext
+    if type(imagename) is not str:
+        imagename = imagename[0].split("/")[-1][:-4]
+
+    output = np.concatenate([output_image, heat_map, confidence_mask_gray], axis=1)
+    outpath = vis_dir + f"/{imagename}.jpg"
+    if not os.path.exists(os.path.dirname(outpath)):
+        os.makedirs(os.path.dirname(outpath), exist_ok=True)
+
+    cv2.imwrite(outpath, output)
+    # print(f'Logging original image into {outpath}')
+
+
+def save_parser(args):
+
+    """ final options """
+    with open(f"{args.results_dir}/opt.txt", "a", encoding="utf-8") as opt_file:
+        opt_log = "------------ Options -------------\n"
+        arg = vars(args)
+        for k, v in arg.items():
+            opt_log += f"{str(k)}: {str(v)}\n"
+        opt_log += "---------------------------------------\n"
+        print(opt_log)
+        opt_file.write(opt_log)

trainer/dataset.py

@@ -0,0 +1,283 @@
+import os
+import sys
+import re
+import six
+import math
+import torch
+import pandas  as pd
+
+from natsort import natsorted
+from PIL import Image
+import numpy as np
+from torch.utils.data import Dataset, ConcatDataset, Subset
+from torch._utils import _accumulate
+import torchvision.transforms as transforms
+
+def contrast_grey(img):
+    high = np.percentile(img, 90)
+    low  = np.percentile(img, 10)
+    return (high-low)/(high+low), high, low
+
+def adjust_contrast_grey(img, target = 0.4):
+    contrast, high, low = contrast_grey(img)
+    if contrast < target:
+        img = img.astype(int)
+        ratio = 200./(high-low)
+        img = (img - low + 25)*ratio
+        img = np.maximum(np.full(img.shape, 0) ,np.minimum(np.full(img.shape, 255), img)).astype(np.uint8)
+    return img
+
+
+class Batch_Balanced_Dataset(object):
+
+    def __init__(self, opt):
+        """
+        Modulate the data ratio in the batch.
+        For example, when select_data is "MJ-ST" and batch_ratio is "0.5-0.5",
+        the 50% of the batch is filled with MJ and the other 50% of the batch is filled with ST.
+        """
+        log = open(f'./saved_models/{opt.experiment_name}/log_dataset.txt', 'a')
+        dashed_line = '-' * 80
+        print(dashed_line)
+        log.write(dashed_line + '\n')
+        print(f'dataset_root: {opt.train_data}\nopt.select_data: {opt.select_data}\nopt.batch_ratio: {opt.batch_ratio}')
+        log.write(f'dataset_root: {opt.train_data}\nopt.select_data: {opt.select_data}\nopt.batch_ratio: {opt.batch_ratio}\n')
+        assert len(opt.select_data) == len(opt.batch_ratio)
+
+        _AlignCollate = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD, contrast_adjust = opt.contrast_adjust)
+        self.data_loader_list = []
+        self.dataloader_iter_list = []
+        batch_size_list = []
+        Total_batch_size = 0
+        for selected_d, batch_ratio_d in zip(opt.select_data, opt.batch_ratio):
+            _batch_size = max(round(opt.batch_size * float(batch_ratio_d)), 1)
+            print(dashed_line)
+            log.write(dashed_line + '\n')
+            _dataset, _dataset_log = hierarchical_dataset(root=opt.train_data, opt=opt, select_data=[selected_d])
+            total_number_dataset = len(_dataset)
+            log.write(_dataset_log)
+
+            """
+            The total number of data can be modified with opt.total_data_usage_ratio.
+            ex) opt.total_data_usage_ratio = 1 indicates 100% usage, and 0.2 indicates 20% usage.
+            See 4.2 section in our paper.
+            """
+            number_dataset = int(total_number_dataset * float(opt.total_data_usage_ratio))
+            dataset_split = [number_dataset, total_number_dataset - number_dataset]
+            indices = range(total_number_dataset)
+            _dataset, _ = [Subset(_dataset, indices[offset - length:offset])
+                           for offset, length in zip(_accumulate(dataset_split), dataset_split)]
+            selected_d_log = f'num total samples of {selected_d}: {total_number_dataset} x {opt.total_data_usage_ratio} (total_data_usage_ratio) = {len(_dataset)}\n'
+            selected_d_log += f'num samples of {selected_d} per batch: {opt.batch_size} x {float(batch_ratio_d)} (batch_ratio) = {_batch_size}'
+            print(selected_d_log)
+            log.write(selected_d_log + '\n')
+            batch_size_list.append(str(_batch_size))
+            Total_batch_size += _batch_size
+
+            _data_loader = torch.utils.data.DataLoader(
+                _dataset, batch_size=_batch_size,
+                shuffle=True,
+                num_workers=int(opt.workers), #prefetch_factor=2,persistent_workers=True,
+                collate_fn=_AlignCollate, pin_memory=True)
+            self.data_loader_list.append(_data_loader)
+            self.dataloader_iter_list.append(iter(_data_loader))
+
+        Total_batch_size_log = f'{dashed_line}\n'
+        batch_size_sum = '+'.join(batch_size_list)
+        Total_batch_size_log += f'Total_batch_size: {batch_size_sum} = {Total_batch_size}\n'
+        Total_batch_size_log += f'{dashed_line}'
+        opt.batch_size = Total_batch_size
+
+        print(Total_batch_size_log)
+        log.write(Total_batch_size_log + '\n')
+        log.close()
+
+    def get_batch(self):
+        balanced_batch_images = []
+        balanced_batch_texts = []
+
+        for i, data_loader_iter in enumerate(self.dataloader_iter_list):
+            try:
+                image, text = data_loader_iter.next()
+                balanced_batch_images.append(image)
+                balanced_batch_texts += text
+            except StopIteration:
+                self.dataloader_iter_list[i] = iter(self.data_loader_list[i])
+                image, text = self.dataloader_iter_list[i].next()
+                balanced_batch_images.append(image)
+                balanced_batch_texts += text
+            except ValueError:
+                pass
+
+        balanced_batch_images = torch.cat(balanced_batch_images, 0)
+
+        return balanced_batch_images, balanced_batch_texts
+
+
+def hierarchical_dataset(root, opt, select_data='/'):
+    """ select_data='/' contains all sub-directory of root directory """
+    dataset_list = []
+    dataset_log = f'dataset_root:    {root}\t dataset: {select_data[0]}'
+    print(dataset_log)
+    dataset_log += '\n'
+    for dirpath, dirnames, filenames in os.walk(root+'/'):
+        if not dirnames:
+            select_flag = False
+            for selected_d in select_data:
+                if selected_d in dirpath:
+                    select_flag = True
+                    break
+
+            if select_flag:
+                dataset = OCRDataset(dirpath, opt)
+                sub_dataset_log = f'sub-directory:\t/{os.path.relpath(dirpath, root)}\t num samples: {len(dataset)}'
+                print(sub_dataset_log)
+                dataset_log += f'{sub_dataset_log}\n'
+                dataset_list.append(dataset)
+
+    concatenated_dataset = ConcatDataset(dataset_list)
+
+    return concatenated_dataset, dataset_log
+
+class OCRDataset(Dataset):
+
+    def __init__(self, root, opt):
+
+        self.root = root
+        self.opt = opt
+        print(root)
+        self.df = pd.read_csv(os.path.join(root,'labels.csv'), sep='^([^,]+),', engine='python', usecols=['filename', 'words'], keep_default_na=False)
+        self.nSamples = len(self.df)
+
+        if self.opt.data_filtering_off:
+            self.filtered_index_list = [index + 1 for index in range(self.nSamples)]
+        else:
+            self.filtered_index_list = []
+            for index in range(self.nSamples):
+                label = self.df.at[index,'words']
+                try:
+                    if len(label) > self.opt.batch_max_length:
+                        continue
+                except:
+                    print(label)
+                out_of_char = f'[^{self.opt.character}]'
+                if re.search(out_of_char, label.lower()):
+                    continue
+                self.filtered_index_list.append(index)
+            self.nSamples = len(self.filtered_index_list)
+
+    def __len__(self):
+        return self.nSamples
+
+    def __getitem__(self, index):
+        index = self.filtered_index_list[index]
+        img_fname = self.df.at[index,'filename']
+        img_fpath = os.path.join(self.root, img_fname)
+        label = self.df.at[index,'words']
+
+        if self.opt.rgb:
+            img = Image.open(img_fpath).convert('RGB')  # for color image
+        else:
+            img = Image.open(img_fpath).convert('L')
+
+        if not self.opt.sensitive:
+            label = label.lower()
+
+        # We only train and evaluate on alphanumerics (or pre-defined character set in train.py)
+        out_of_char = f'[^{self.opt.character}]'
+        label = re.sub(out_of_char, '', label)
+
+        return (img, label)
+
+class ResizeNormalize(object):
+
+    def __init__(self, size, interpolation=Image.BICUBIC):
+        self.size = size
+        self.interpolation = interpolation
+        self.toTensor = transforms.ToTensor()
+
+    def __call__(self, img):
+        img = img.resize(self.size, self.interpolation)
+        img = self.toTensor(img)
+        img.sub_(0.5).div_(0.5)
+        return img
+
+
+class NormalizePAD(object):
+
+    def __init__(self, max_size, PAD_type='right'):
+        self.toTensor = transforms.ToTensor()
+        self.max_size = max_size
+        self.max_width_half = math.floor(max_size[2] / 2)
+        self.PAD_type = PAD_type
+
+    def __call__(self, img):
+        img = self.toTensor(img)
+        img.sub_(0.5).div_(0.5)
+        c, h, w = img.size()
+        Pad_img = torch.FloatTensor(*self.max_size).fill_(0)
+        Pad_img[:, :, :w] = img  # right pad
+        if self.max_size[2] != w:  # add border Pad
+            Pad_img[:, :, w:] = img[:, :, w - 1].unsqueeze(2).expand(c, h, self.max_size[2] - w)
+
+        return Pad_img
+
+
+class AlignCollate(object):
+
+    def __init__(self, imgH=32, imgW=100, keep_ratio_with_pad=False, contrast_adjust = 0.):
+        self.imgH = imgH
+        self.imgW = imgW
+        self.keep_ratio_with_pad = keep_ratio_with_pad
+        self.contrast_adjust = contrast_adjust
+
+    def __call__(self, batch):
+        batch = filter(lambda x: x is not None, batch)
+        images, labels = zip(*batch)
+
+        if self.keep_ratio_with_pad:  # same concept with 'Rosetta' paper
+            resized_max_w = self.imgW
+            input_channel = 3 if images[0].mode == 'RGB' else 1
+            transform = NormalizePAD((input_channel, self.imgH, resized_max_w))
+
+            resized_images = []
+            for image in images:
+                w, h = image.size
+
+                #### augmentation here - change contrast
+                if self.contrast_adjust > 0:
+                    image = np.array(image.convert("L"))
+                    image = adjust_contrast_grey(image, target = self.contrast_adjust)
+                    image = Image.fromarray(image, 'L')
+
+                ratio = w / float(h)
+                if math.ceil(self.imgH * ratio) > self.imgW:
+                    resized_w = self.imgW
+                else:
+                    resized_w = math.ceil(self.imgH * ratio)
+
+                resized_image = image.resize((resized_w, self.imgH), Image.BICUBIC)
+                resized_images.append(transform(resized_image))
+                # resized_image.save('./image_test/%d_test.jpg' % w)
+
+            image_tensors = torch.cat([t.unsqueeze(0) for t in resized_images], 0)
+
+        else:
+            transform = ResizeNormalize((self.imgW, self.imgH))
+            image_tensors = [transform(image) for image in images]
+            image_tensors = torch.cat([t.unsqueeze(0) for t in image_tensors], 0)
+
+        return image_tensors, labels
+
+
+def tensor2im(image_tensor, imtype=np.uint8):
+    image_numpy = image_tensor.cpu().float().numpy()
+    if image_numpy.shape[0] == 1:
+        image_numpy = np.tile(image_numpy, (3, 1, 1))
+    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
+    return image_numpy.astype(imtype)
+
+
+def save_image(image_numpy, image_path):
+    image_pil = Image.fromarray(image_numpy)
+    image_pil.save(image_path)

trainer/model.py

@@ -0,0 +1,74 @@
+import torch.nn as nn
+from modules.transformation import TPS_SpatialTransformerNetwork
+from modules.feature_extraction import VGG_FeatureExtractor, RCNN_FeatureExtractor, ResNet_FeatureExtractor
+from modules.sequence_modeling import BidirectionalLSTM
+from modules.prediction import Attention
+
+class Model(nn.Module):
+
+    def __init__(self, opt):
+        super(Model, self).__init__()
+        self.opt = opt
+        self.stages = {'Trans': opt.Transformation, 'Feat': opt.FeatureExtraction,
+                       'Seq': opt.SequenceModeling, 'Pred': opt.Prediction}
+
+        """ Transformation """
+        if opt.Transformation == 'TPS':
+            self.Transformation = TPS_SpatialTransformerNetwork(
+                F=opt.num_fiducial, I_size=(opt.imgH, opt.imgW), I_r_size=(opt.imgH, opt.imgW), I_channel_num=opt.input_channel)
+        else:
+            print('No Transformation module specified')
+
+        """ FeatureExtraction """
+        if opt.FeatureExtraction == 'VGG':
+            self.FeatureExtraction = VGG_FeatureExtractor(opt.input_channel, opt.output_channel)
+        elif opt.FeatureExtraction == 'RCNN':
+            self.FeatureExtraction = RCNN_FeatureExtractor(opt.input_channel, opt.output_channel)
+        elif opt.FeatureExtraction == 'ResNet':
+            self.FeatureExtraction = ResNet_FeatureExtractor(opt.input_channel, opt.output_channel)
+        else:
+            raise Exception('No FeatureExtraction module specified')
+        self.FeatureExtraction_output = opt.output_channel  # int(imgH/16-1) * 512
+        self.AdaptiveAvgPool = nn.AdaptiveAvgPool2d((None, 1))  # Transform final (imgH/16-1) -> 1
+
+        """ Sequence modeling"""
+        if opt.SequenceModeling == 'BiLSTM':
+            self.SequenceModeling = nn.Sequential(
+                BidirectionalLSTM(self.FeatureExtraction_output, opt.hidden_size, opt.hidden_size),
+                BidirectionalLSTM(opt.hidden_size, opt.hidden_size, opt.hidden_size))
+            self.SequenceModeling_output = opt.hidden_size
+        else:
+            print('No SequenceModeling module specified')
+            self.SequenceModeling_output = self.FeatureExtraction_output
+
+        """ Prediction """
+        if opt.Prediction == 'CTC':
+            self.Prediction = nn.Linear(self.SequenceModeling_output, opt.num_class)
+        elif opt.Prediction == 'Attn':
+            self.Prediction = Attention(self.SequenceModeling_output, opt.hidden_size, opt.num_class)
+        else:
+            raise Exception('Prediction is neither CTC or Attn')
+
+    def forward(self, input, text, is_train=True):
+        """ Transformation stage """
+        if not self.stages['Trans'] == "None":
+            input = self.Transformation(input)
+
+        """ Feature extraction stage """
+        visual_feature = self.FeatureExtraction(input)
+        visual_feature = self.AdaptiveAvgPool(visual_feature.permute(0, 3, 1, 2))  # [b, c, h, w] -> [b, w, c, h]
+        visual_feature = visual_feature.squeeze(3)
+
+        """ Sequence modeling stage """
+        if self.stages['Seq'] == 'BiLSTM':
+            contextual_feature = self.SequenceModeling(visual_feature)
+        else:
+            contextual_feature = visual_feature  # for convenience. this is NOT contextually modeled by BiLSTM
+
+        """ Prediction stage """
+        if self.stages['Pred'] == 'CTC':
+            prediction = self.Prediction(contextual_feature.contiguous())
+        else:
+            prediction = self.Prediction(contextual_feature.contiguous(), text, is_train, batch_max_length=self.opt.batch_max_length)
+
+        return prediction

trainer/modules/feature_extraction.py

@@ -0,0 +1,246 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class VGG_FeatureExtractor(nn.Module):
+    """ FeatureExtractor of CRNN (https://arxiv.org/pdf/1507.05717.pdf) """
+
+    def __init__(self, input_channel, output_channel=512):
+        super(VGG_FeatureExtractor, self).__init__()
+        self.output_channel = [int(output_channel / 8), int(output_channel / 4),
+                               int(output_channel / 2), output_channel]  # [64, 128, 256, 512]
+        self.ConvNet = nn.Sequential(
+            nn.Conv2d(input_channel, self.output_channel[0], 3, 1, 1), nn.ReLU(True),
+            nn.MaxPool2d(2, 2),  # 64x16x50
+            nn.Conv2d(self.output_channel[0], self.output_channel[1], 3, 1, 1), nn.ReLU(True),
+            nn.MaxPool2d(2, 2),  # 128x8x25
+            nn.Conv2d(self.output_channel[1], self.output_channel[2], 3, 1, 1), nn.ReLU(True),  # 256x8x25
+            nn.Conv2d(self.output_channel[2], self.output_channel[2], 3, 1, 1), nn.ReLU(True),
+            nn.MaxPool2d((2, 1), (2, 1)),  # 256x4x25
+            nn.Conv2d(self.output_channel[2], self.output_channel[3], 3, 1, 1, bias=False),
+            nn.BatchNorm2d(self.output_channel[3]), nn.ReLU(True),  # 512x4x25
+            nn.Conv2d(self.output_channel[3], self.output_channel[3], 3, 1, 1, bias=False),
+            nn.BatchNorm2d(self.output_channel[3]), nn.ReLU(True),
+            nn.MaxPool2d((2, 1), (2, 1)),  # 512x2x25
+            nn.Conv2d(self.output_channel[3], self.output_channel[3], 2, 1, 0), nn.ReLU(True))  # 512x1x24
+
+    def forward(self, input):
+        return self.ConvNet(input)
+
+
+class RCNN_FeatureExtractor(nn.Module):
+    """ FeatureExtractor of GRCNN (https://papers.nips.cc/paper/6637-gated-recurrent-convolution-neural-network-for-ocr.pdf) """
+
+    def __init__(self, input_channel, output_channel=512):
+        super(RCNN_FeatureExtractor, self).__init__()
+        self.output_channel = [int(output_channel / 8), int(output_channel / 4),
+                               int(output_channel / 2), output_channel]  # [64, 128, 256, 512]
+        self.ConvNet = nn.Sequential(
+            nn.Conv2d(input_channel, self.output_channel[0], 3, 1, 1), nn.ReLU(True),
+            nn.MaxPool2d(2, 2),  # 64 x 16 x 50
+            GRCL(self.output_channel[0], self.output_channel[0], num_iteration=5, kernel_size=3, pad=1),
+            nn.MaxPool2d(2, 2),  # 64 x 8 x 25
+            GRCL(self.output_channel[0], self.output_channel[1], num_iteration=5, kernel_size=3, pad=1),
+            nn.MaxPool2d(2, (2, 1), (0, 1)),  # 128 x 4 x 26
+            GRCL(self.output_channel[1], self.output_channel[2], num_iteration=5, kernel_size=3, pad=1),
+            nn.MaxPool2d(2, (2, 1), (0, 1)),  # 256 x 2 x 27
+            nn.Conv2d(self.output_channel[2], self.output_channel[3], 2, 1, 0, bias=False),
+            nn.BatchNorm2d(self.output_channel[3]), nn.ReLU(True))  # 512 x 1 x 26
+
+    def forward(self, input):
+        return self.ConvNet(input)
+
+
+class ResNet_FeatureExtractor(nn.Module):
+    """ FeatureExtractor of FAN (http://openaccess.thecvf.com/content_ICCV_2017/papers/Cheng_Focusing_Attention_Towards_ICCV_2017_paper.pdf) """
+
+    def __init__(self, input_channel, output_channel=512):
+        super(ResNet_FeatureExtractor, self).__init__()
+        self.ConvNet = ResNet(input_channel, output_channel, BasicBlock, [1, 2, 5, 3])
+
+    def forward(self, input):
+        return self.ConvNet(input)
+
+
+# For Gated RCNN
+class GRCL(nn.Module):
+
+    def __init__(self, input_channel, output_channel, num_iteration, kernel_size, pad):
+        super(GRCL, self).__init__()
+        self.wgf_u = nn.Conv2d(input_channel, output_channel, 1, 1, 0, bias=False)
+        self.wgr_x = nn.Conv2d(output_channel, output_channel, 1, 1, 0, bias=False)
+        self.wf_u = nn.Conv2d(input_channel, output_channel, kernel_size, 1, pad, bias=False)
+        self.wr_x = nn.Conv2d(output_channel, output_channel, kernel_size, 1, pad, bias=False)
+
+        self.BN_x_init = nn.BatchNorm2d(output_channel)
+
+        self.num_iteration = num_iteration
+        self.GRCL = [GRCL_unit(output_channel) for _ in range(num_iteration)]
+        self.GRCL = nn.Sequential(*self.GRCL)
+
+    def forward(self, input):
+        """ The input of GRCL is consistant over time t, which is denoted by u(0)
+        thus wgf_u / wf_u is also consistant over time t.
+        """
+        wgf_u = self.wgf_u(input)
+        wf_u = self.wf_u(input)
+        x = F.relu(self.BN_x_init(wf_u))
+
+        for i in range(self.num_iteration):
+            x = self.GRCL[i](wgf_u, self.wgr_x(x), wf_u, self.wr_x(x))
+
+        return x
+
+
+class GRCL_unit(nn.Module):
+
+    def __init__(self, output_channel):
+        super(GRCL_unit, self).__init__()
+        self.BN_gfu = nn.BatchNorm2d(output_channel)
+        self.BN_grx = nn.BatchNorm2d(output_channel)
+        self.BN_fu = nn.BatchNorm2d(output_channel)
+        self.BN_rx = nn.BatchNorm2d(output_channel)
+        self.BN_Gx = nn.BatchNorm2d(output_channel)
+
+    def forward(self, wgf_u, wgr_x, wf_u, wr_x):
+        G_first_term = self.BN_gfu(wgf_u)
+        G_second_term = self.BN_grx(wgr_x)
+        G = F.sigmoid(G_first_term + G_second_term)
+
+        x_first_term = self.BN_fu(wf_u)
+        x_second_term = self.BN_Gx(self.BN_rx(wr_x) * G)
+        x = F.relu(x_first_term + x_second_term)
+
+        return x
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = self._conv3x3(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = self._conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def _conv3x3(self, in_planes, out_planes, stride=1):
+        "3x3 convolution with padding"
+        return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                         padding=1, bias=False)
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, input_channel, output_channel, block, layers):
+        super(ResNet, self).__init__()
+
+        self.output_channel_block = [int(output_channel / 4), int(output_channel / 2), output_channel, output_channel]
+
+        self.inplanes = int(output_channel / 8)
+        self.conv0_1 = nn.Conv2d(input_channel, int(output_channel / 16),
+                                 kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn0_1 = nn.BatchNorm2d(int(output_channel / 16))
+        self.conv0_2 = nn.Conv2d(int(output_channel / 16), self.inplanes,
+                                 kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn0_2 = nn.BatchNorm2d(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
+        self.layer1 = self._make_layer(block, self.output_channel_block[0], layers[0])
+        self.conv1 = nn.Conv2d(self.output_channel_block[0], self.output_channel_block[
+                               0], kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.output_channel_block[0])
+
+        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
+        self.layer2 = self._make_layer(block, self.output_channel_block[1], layers[1], stride=1)
+        self.conv2 = nn.Conv2d(self.output_channel_block[1], self.output_channel_block[
+                               1], kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(self.output_channel_block[1])
+
+        self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=(2, 1), padding=(0, 1))
+        self.layer3 = self._make_layer(block, self.output_channel_block[2], layers[2], stride=1)
+        self.conv3 = nn.Conv2d(self.output_channel_block[2], self.output_channel_block[
+                               2], kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(self.output_channel_block[2])
+
+        self.layer4 = self._make_layer(block, self.output_channel_block[3], layers[3], stride=1)
+        self.conv4_1 = nn.Conv2d(self.output_channel_block[3], self.output_channel_block[
+                                 3], kernel_size=2, stride=(2, 1), padding=(0, 1), bias=False)
+        self.bn4_1 = nn.BatchNorm2d(self.output_channel_block[3])
+        self.conv4_2 = nn.Conv2d(self.output_channel_block[3], self.output_channel_block[
+                                 3], kernel_size=2, stride=1, padding=0, bias=False)
+        self.bn4_2 = nn.BatchNorm2d(self.output_channel_block[3])
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv0_1(x)
+        x = self.bn0_1(x)
+        x = self.relu(x)
+        x = self.conv0_2(x)
+        x = self.bn0_2(x)
+        x = self.relu(x)
+
+        x = self.maxpool1(x)
+        x = self.layer1(x)
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x = self.maxpool2(x)
+        x = self.layer2(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+
+        x = self.maxpool3(x)
+        x = self.layer3(x)
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = self.relu(x)
+
+        x = self.layer4(x)
+        x = self.conv4_1(x)
+        x = self.bn4_1(x)
+        x = self.relu(x)
+        x = self.conv4_2(x)
+        x = self.bn4_2(x)
+        x = self.relu(x)
+
+        return x

trainer/modules/prediction.py

@@ -0,0 +1,81 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+class Attention(nn.Module):
+
+    def __init__(self, input_size, hidden_size, num_classes):
+        super(Attention, self).__init__()
+        self.attention_cell = AttentionCell(input_size, hidden_size, num_classes)
+        self.hidden_size = hidden_size
+        self.num_classes = num_classes
+        self.generator = nn.Linear(hidden_size, num_classes)
+
+    def _char_to_onehot(self, input_char, onehot_dim=38):
+        input_char = input_char.unsqueeze(1)
+        batch_size = input_char.size(0)
+        one_hot = torch.FloatTensor(batch_size, onehot_dim).zero_().to(device)
+        one_hot = one_hot.scatter_(1, input_char, 1)
+        return one_hot
+
+    def forward(self, batch_H, text, is_train=True, batch_max_length=25):
+        """
+        input:
+            batch_H : contextual_feature H = hidden state of encoder. [batch_size x num_steps x num_classes]
+            text : the text-index of each image. [batch_size x (max_length+1)]. +1 for [GO] token. text[:, 0] = [GO].
+        output: probability distribution at each step [batch_size x num_steps x num_classes]
+        """
+        batch_size = batch_H.size(0)
+        num_steps = batch_max_length + 1  # +1 for [s] at end of sentence.
+
+        output_hiddens = torch.FloatTensor(batch_size, num_steps, self.hidden_size).fill_(0).to(device)
+        hidden = (torch.FloatTensor(batch_size, self.hidden_size).fill_(0).to(device),
+                  torch.FloatTensor(batch_size, self.hidden_size).fill_(0).to(device))
+
+        if is_train:
+            for i in range(num_steps):
+                # one-hot vectors for a i-th char. in a batch
+                char_onehots = self._char_to_onehot(text[:, i], onehot_dim=self.num_classes)
+                # hidden : decoder's hidden s_{t-1}, batch_H : encoder's hidden H, char_onehots : one-hot(y_{t-1})
+                hidden, alpha = self.attention_cell(hidden, batch_H, char_onehots)
+                output_hiddens[:, i, :] = hidden[0]  # LSTM hidden index (0: hidden, 1: Cell)
+            probs = self.generator(output_hiddens)
+
+        else:
+            targets = torch.LongTensor(batch_size).fill_(0).to(device)  # [GO] token
+            probs = torch.FloatTensor(batch_size, num_steps, self.num_classes).fill_(0).to(device)
+
+            for i in range(num_steps):
+                char_onehots = self._char_to_onehot(targets, onehot_dim=self.num_classes)
+                hidden, alpha = self.attention_cell(hidden, batch_H, char_onehots)
+                probs_step = self.generator(hidden[0])
+                probs[:, i, :] = probs_step
+                _, next_input = probs_step.max(1)
+                targets = next_input
+
+        return probs  # batch_size x num_steps x num_classes
+
+
+class AttentionCell(nn.Module):
+
+    def __init__(self, input_size, hidden_size, num_embeddings):
+        super(AttentionCell, self).__init__()
+        self.i2h = nn.Linear(input_size, hidden_size, bias=False)
+        self.h2h = nn.Linear(hidden_size, hidden_size)  # either i2i or h2h should have bias
+        self.score = nn.Linear(hidden_size, 1, bias=False)
+        self.rnn = nn.LSTMCell(input_size + num_embeddings, hidden_size)
+        self.hidden_size = hidden_size
+
+    def forward(self, prev_hidden, batch_H, char_onehots):
+        # [batch_size x num_encoder_step x num_channel] -> [batch_size x num_encoder_step x hidden_size]
+        batch_H_proj = self.i2h(batch_H)
+        prev_hidden_proj = self.h2h(prev_hidden[0]).unsqueeze(1)
+        e = self.score(torch.tanh(batch_H_proj + prev_hidden_proj))  # batch_size x num_encoder_step * 1
+
+        alpha = F.softmax(e, dim=1)
+        context = torch.bmm(alpha.permute(0, 2, 1), batch_H).squeeze(1)  # batch_size x num_channel
+        concat_context = torch.cat([context, char_onehots], 1)  # batch_size x (num_channel + num_embedding)
+        cur_hidden = self.rnn(concat_context, prev_hidden)
+        return cur_hidden, alpha

trainer/modules/sequence_modeling.py

@@ -0,0 +1,22 @@
+import torch.nn as nn
+
+
+class BidirectionalLSTM(nn.Module):
+
+    def __init__(self, input_size, hidden_size, output_size):
+        super(BidirectionalLSTM, self).__init__()
+        self.rnn = nn.LSTM(input_size, hidden_size, bidirectional=True, batch_first=True)
+        self.linear = nn.Linear(hidden_size * 2, output_size)
+
+    def forward(self, input):
+        """
+        input : visual feature [batch_size x T x input_size]
+        output : contextual feature [batch_size x T x output_size]
+        """
+        try:
+            self.rnn.flatten_parameters()
+        except:
+            pass
+        recurrent, _ = self.rnn(input)  # batch_size x T x input_size -> batch_size x T x (2*hidden_size)
+        output = self.linear(recurrent)  # batch_size x T x output_size
+        return output

trainer/modules/transformation.py

@@ -0,0 +1,160 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+class TPS_SpatialTransformerNetwork(nn.Module):
+    """ Rectification Network of RARE, namely TPS based STN """
+
+    def __init__(self, F, I_size, I_r_size, I_channel_num=1):
+        """ Based on RARE TPS
+        input:
+            batch_I: Batch Input Image [batch_size x I_channel_num x I_height x I_width]
+            I_size : (height, width) of the input image I
+            I_r_size : (height, width) of the rectified image I_r
+            I_channel_num : the number of channels of the input image I
+        output:
+            batch_I_r: rectified image [batch_size x I_channel_num x I_r_height x I_r_width]
+        """
+        super(TPS_SpatialTransformerNetwork, self).__init__()
+        self.F = F
+        self.I_size = I_size
+        self.I_r_size = I_r_size  # = (I_r_height, I_r_width)
+        self.I_channel_num = I_channel_num
+        self.LocalizationNetwork = LocalizationNetwork(self.F, self.I_channel_num)
+        self.GridGenerator = GridGenerator(self.F, self.I_r_size)
+
+    def forward(self, batch_I):
+        batch_C_prime = self.LocalizationNetwork(batch_I)  # batch_size x K x 2
+        build_P_prime = self.GridGenerator.build_P_prime(batch_C_prime)  # batch_size x n (= I_r_width x I_r_height) x 2
+        build_P_prime_reshape = build_P_prime.reshape([build_P_prime.size(0), self.I_r_size[0], self.I_r_size[1], 2])
+        batch_I_r = F.grid_sample(batch_I, build_P_prime_reshape, padding_mode='border')
+
+        return batch_I_r
+
+
+class LocalizationNetwork(nn.Module):
+    """ Localization Network of RARE, which predicts C' (K x 2) from I (I_width x I_height) """
+
+    def __init__(self, F, I_channel_num):
+        super(LocalizationNetwork, self).__init__()
+        self.F = F
+        self.I_channel_num = I_channel_num
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels=self.I_channel_num, out_channels=64, kernel_size=3, stride=1, padding=1,
+                      bias=False), nn.BatchNorm2d(64), nn.ReLU(True),
+            nn.MaxPool2d(2, 2),  # batch_size x 64 x I_height/2 x I_width/2
+            nn.Conv2d(64, 128, 3, 1, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(True),
+            nn.MaxPool2d(2, 2),  # batch_size x 128 x I_height/4 x I_width/4
+            nn.Conv2d(128, 256, 3, 1, 1, bias=False), nn.BatchNorm2d(256), nn.ReLU(True),
+            nn.MaxPool2d(2, 2),  # batch_size x 256 x I_height/8 x I_width/8
+            nn.Conv2d(256, 512, 3, 1, 1, bias=False), nn.BatchNorm2d(512), nn.ReLU(True),
+            nn.AdaptiveAvgPool2d(1)  # batch_size x 512
+        )
+
+        self.localization_fc1 = nn.Sequential(nn.Linear(512, 256), nn.ReLU(True))
+        self.localization_fc2 = nn.Linear(256, self.F * 2)
+
+        # Init fc2 in LocalizationNetwork
+        self.localization_fc2.weight.data.fill_(0)
+        """ see RARE paper Fig. 6 (a) """
+        ctrl_pts_x = np.linspace(-1.0, 1.0, int(F / 2))
+        ctrl_pts_y_top = np.linspace(0.0, -1.0, num=int(F / 2))
+        ctrl_pts_y_bottom = np.linspace(1.0, 0.0, num=int(F / 2))
+        ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+        ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+        initial_bias = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0)
+        self.localization_fc2.bias.data = torch.from_numpy(initial_bias).float().view(-1)
+
+    def forward(self, batch_I):
+        """
+        input:     batch_I : Batch Input Image [batch_size x I_channel_num x I_height x I_width]
+        output:    batch_C_prime : Predicted coordinates of fiducial points for input batch [batch_size x F x 2]
+        """
+        batch_size = batch_I.size(0)
+        features = self.conv(batch_I).view(batch_size, -1)
+        batch_C_prime = self.localization_fc2(self.localization_fc1(features)).view(batch_size, self.F, 2)
+        return batch_C_prime
+
+
+class GridGenerator(nn.Module):
+    """ Grid Generator of RARE, which produces P_prime by multiplying T with P """
+
+    def __init__(self, F, I_r_size):
+        """ Generate P_hat and inv_delta_C for later """
+        super(GridGenerator, self).__init__()
+        self.eps = 1e-6
+        self.I_r_height, self.I_r_width = I_r_size
+        self.F = F
+        self.C = self._build_C(self.F)  # F x 2
+        self.P = self._build_P(self.I_r_width, self.I_r_height)
+        ## for multi-gpu, you need register buffer
+        self.register_buffer("inv_delta_C", torch.tensor(self._build_inv_delta_C(self.F, self.C)).float())  # F+3 x F+3
+        self.register_buffer("P_hat", torch.tensor(self._build_P_hat(self.F, self.C, self.P)).float())  # n x F+3
+        ## for fine-tuning with different image width, you may use below instead of self.register_buffer
+        #self.inv_delta_C = torch.tensor(self._build_inv_delta_C(self.F, self.C)).float().cuda()  # F+3 x F+3
+        #self.P_hat = torch.tensor(self._build_P_hat(self.F, self.C, self.P)).float().cuda()  # n x F+3
+
+    def _build_C(self, F):
+        """ Return coordinates of fiducial points in I_r; C """
+        ctrl_pts_x = np.linspace(-1.0, 1.0, int(F / 2))
+        ctrl_pts_y_top = -1 * np.ones(int(F / 2))
+        ctrl_pts_y_bottom = np.ones(int(F / 2))
+        ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+        ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+        C = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0)
+        return C  # F x 2
+
+    def _build_inv_delta_C(self, F, C):
+        """ Return inv_delta_C which is needed to calculate T """
+        hat_C = np.zeros((F, F), dtype=float)  # F x F
+        for i in range(0, F):
+            for j in range(i, F):
+                r = np.linalg.norm(C[i] - C[j])
+                hat_C[i, j] = r
+                hat_C[j, i] = r
+        np.fill_diagonal(hat_C, 1)
+        hat_C = (hat_C ** 2) * np.log(hat_C)
+        # print(C.shape, hat_C.shape)
+        delta_C = np.concatenate(  # F+3 x F+3
+            [
+                np.concatenate([np.ones((F, 1)), C, hat_C], axis=1),  # F x F+3
+                np.concatenate([np.zeros((2, 3)), np.transpose(C)], axis=1),  # 2 x F+3
+                np.concatenate([np.zeros((1, 3)), np.ones((1, F))], axis=1)  # 1 x F+3
+            ],
+            axis=0
+        )
+        inv_delta_C = np.linalg.inv(delta_C)
+        return inv_delta_C  # F+3 x F+3
+
+    def _build_P(self, I_r_width, I_r_height):
+        I_r_grid_x = (np.arange(-I_r_width, I_r_width, 2) + 1.0) / I_r_width  # self.I_r_width
+        I_r_grid_y = (np.arange(-I_r_height, I_r_height, 2) + 1.0) / I_r_height  # self.I_r_height
+        P = np.stack(  # self.I_r_width x self.I_r_height x 2
+            np.meshgrid(I_r_grid_x, I_r_grid_y),
+            axis=2
+        )
+        return P.reshape([-1, 2])  # n (= self.I_r_width x self.I_r_height) x 2
+
+    def _build_P_hat(self, F, C, P):
+        n = P.shape[0]  # n (= self.I_r_width x self.I_r_height)
+        P_tile = np.tile(np.expand_dims(P, axis=1), (1, F, 1))  # n x 2 -> n x 1 x 2 -> n x F x 2
+        C_tile = np.expand_dims(C, axis=0)  # 1 x F x 2
+        P_diff = P_tile - C_tile  # n x F x 2
+        rbf_norm = np.linalg.norm(P_diff, ord=2, axis=2, keepdims=False)  # n x F
+        rbf = np.multiply(np.square(rbf_norm), np.log(rbf_norm + self.eps))  # n x F
+        P_hat = np.concatenate([np.ones((n, 1)), P, rbf], axis=1)
+        return P_hat  # n x F+3
+
+    def build_P_prime(self, batch_C_prime):
+        """ Generate Grid from batch_C_prime [batch_size x F x 2] """
+        batch_size = batch_C_prime.size(0)
+        batch_inv_delta_C = self.inv_delta_C.repeat(batch_size, 1, 1)
+        batch_P_hat = self.P_hat.repeat(batch_size, 1, 1)
+        batch_C_prime_with_zeros = torch.cat((batch_C_prime, torch.zeros(
+            batch_size, 3, 2).float().to(device)), dim=1)  # batch_size x F+3 x 2
+        batch_T = torch.bmm(batch_inv_delta_C, batch_C_prime_with_zeros)  # batch_size x F+3 x 2
+        batch_P_prime = torch.bmm(batch_P_hat, batch_T)  # batch_size x n x 2
+        return batch_P_prime  # batch_size x n x 2

trainer/saved_models/folder.txt

@@ -0,0 +1 @@
+trained model will be saved here

trainer/test.py

@@ -0,0 +1,112 @@
+import os
+import time
+import string
+import argparse
+
+import torch
+import torch.backends.cudnn as cudnn
+import torch.utils.data
+import torch.nn.functional as F
+import numpy as np
+from nltk.metrics.distance import edit_distance
+
+from utils import CTCLabelConverter, AttnLabelConverter, Averager
+from dataset import hierarchical_dataset, AlignCollate
+from model import Model
+
+def validation(model, criterion, evaluation_loader, converter, opt, device):
+    """ validation or evaluation """
+    n_correct = 0
+    norm_ED = 0
+    length_of_data = 0
+    infer_time = 0
+    valid_loss_avg = Averager()
+
+    for i, (image_tensors, labels) in enumerate(evaluation_loader):
+        batch_size = image_tensors.size(0)
+        length_of_data = length_of_data + batch_size
+        image = image_tensors.to(device)
+        # For max length prediction
+        length_for_pred = torch.IntTensor([opt.batch_max_length] * batch_size).to(device)
+        text_for_pred = torch.LongTensor(batch_size, opt.batch_max_length + 1).fill_(0).to(device)
+
+        text_for_loss, length_for_loss = converter.encode(labels, batch_max_length=opt.batch_max_length)
+        
+        start_time = time.time()
+        if 'CTC' in opt.Prediction:
+            preds = model(image, text_for_pred)
+            forward_time = time.time() - start_time
+
+            # Calculate evaluation loss for CTC decoder.
+            preds_size = torch.IntTensor([preds.size(1)] * batch_size)
+            # permute 'preds' to use CTCloss format
+            cost = criterion(preds.log_softmax(2).permute(1, 0, 2), text_for_loss, preds_size, length_for_loss)
+
+            if opt.decode == 'greedy':
+                # Select max probabilty (greedy decoding) then decode index to character
+                _, preds_index = preds.max(2)
+                preds_index = preds_index.view(-1)
+                preds_str = converter.decode_greedy(preds_index.data, preds_size.data)
+            elif opt.decode == 'beamsearch':
+                preds_str = converter.decode_beamsearch(preds, beamWidth=2)
+
+        else:
+            preds = model(image, text_for_pred, is_train=False)
+            forward_time = time.time() - start_time
+
+            preds = preds[:, :text_for_loss.shape[1] - 1, :]
+            target = text_for_loss[:, 1:]  # without [GO] Symbol
+            cost = criterion(preds.contiguous().view(-1, preds.shape[-1]), target.contiguous().view(-1))
+
+            # select max probabilty (greedy decoding) then decode index to character
+            _, preds_index = preds.max(2)
+            preds_str = converter.decode(preds_index, length_for_pred)
+            labels = converter.decode(text_for_loss[:, 1:], length_for_loss)
+
+        infer_time += forward_time
+        valid_loss_avg.add(cost)
+
+        # calculate accuracy & confidence score
+        preds_prob = F.softmax(preds, dim=2)
+        preds_max_prob, _ = preds_prob.max(dim=2)
+        confidence_score_list = []
+        
+        for gt, pred, pred_max_prob in zip(labels, preds_str, preds_max_prob):
+            if 'Attn' in opt.Prediction:
+                gt = gt[:gt.find('[s]')]
+                pred_EOS = pred.find('[s]')
+                pred = pred[:pred_EOS]  # prune after "end of sentence" token ([s])
+                pred_max_prob = pred_max_prob[:pred_EOS]
+
+            if pred == gt:
+                n_correct += 1
+
+            '''
+            (old version) ICDAR2017 DOST Normalized Edit Distance https://rrc.cvc.uab.es/?ch=7&com=tasks
+            "For each word we calculate the normalized edit distance to the length of the ground truth transcription." 
+            if len(gt) == 0:
+                norm_ED += 1
+            else:
+                norm_ED += edit_distance(pred, gt) / len(gt)
+            '''
+            
+            # ICDAR2019 Normalized Edit Distance 
+            if len(gt) == 0 or len(pred) ==0:
+                norm_ED += 0
+            elif len(gt) > len(pred):
+                norm_ED += 1 - edit_distance(pred, gt) / len(gt)
+            else:
+                norm_ED += 1 - edit_distance(pred, gt) / len(pred)
+
+            # calculate confidence score (= multiply of pred_max_prob)
+            try:
+                confidence_score = pred_max_prob.cumprod(dim=0)[-1]
+            except:
+                confidence_score = 0  # for empty pred case, when prune after "end of sentence" token ([s])
+            confidence_score_list.append(confidence_score)
+            # print(pred, gt, pred==gt, confidence_score)
+
+    accuracy = n_correct / float(length_of_data) * 100
+    norm_ED = norm_ED / float(length_of_data) # ICDAR2019 Normalized Edit Distance
+
+    return valid_loss_avg.val(), accuracy, norm_ED, preds_str, confidence_score_list, labels, infer_time, length_of_data

trainer/train.py

@@ -0,0 +1,282 @@
+import os
+import sys
+import time
+import random
+import torch
+import torch.backends.cudnn as cudnn
+import torch.nn as nn
+import torch.nn.init as init
+import torch.optim as optim
+import torch.utils.data
+from torch.cuda.amp import autocast, GradScaler
+import numpy as np
+
+from utils import CTCLabelConverter, AttnLabelConverter, Averager
+from dataset import hierarchical_dataset, AlignCollate, Batch_Balanced_Dataset
+from model import Model
+from test import validation
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+def count_parameters(model):
+    print("Modules, Parameters")
+    total_params = 0
+    for name, parameter in model.named_parameters():
+        if not parameter.requires_grad: continue
+        param = parameter.numel()
+        #table.add_row([name, param])
+        total_params+=param
+        print(name, param)
+    print(f"Total Trainable Params: {total_params}")
+    return total_params
+
+def train(opt, show_number = 2, amp=False):
+    """ dataset preparation """
+    if not opt.data_filtering_off:
+        print('Filtering the images containing characters which are not in opt.character')
+        print('Filtering the images whose label is longer than opt.batch_max_length')
+
+    opt.select_data = opt.select_data.split('-')
+    opt.batch_ratio = opt.batch_ratio.split('-')
+    train_dataset = Batch_Balanced_Dataset(opt)
+
+    log = open(f'./saved_models/{opt.experiment_name}/log_dataset.txt', 'a', encoding="utf8")
+    AlignCollate_valid = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD, contrast_adjust=opt.contrast_adjust)
+    valid_dataset, valid_dataset_log = hierarchical_dataset(root=opt.valid_data, opt=opt)
+    valid_loader = torch.utils.data.DataLoader(
+        valid_dataset, batch_size=min(32, opt.batch_size),
+        shuffle=True,  # 'True' to check training progress with validation function.
+        num_workers=int(opt.workers), prefetch_factor=512,
+        collate_fn=AlignCollate_valid, pin_memory=True)
+    log.write(valid_dataset_log)
+    print('-' * 80)
+    log.write('-' * 80 + '\n')
+    log.close()
+    
+    """ model configuration """
+    if 'CTC' in opt.Prediction:
+        converter = CTCLabelConverter(opt.character)
+    else:
+        converter = AttnLabelConverter(opt.character)
+    opt.num_class = len(converter.character)
+
+    if opt.rgb:
+        opt.input_channel = 3
+    model = Model(opt)
+    print('model input parameters', opt.imgH, opt.imgW, opt.num_fiducial, opt.input_channel, opt.output_channel,
+          opt.hidden_size, opt.num_class, opt.batch_max_length, opt.Transformation, opt.FeatureExtraction,
+          opt.SequenceModeling, opt.Prediction)
+
+    if opt.saved_model != '':
+        pretrained_dict = torch.load(opt.saved_model)
+        if opt.new_prediction:
+            model.Prediction = nn.Linear(model.SequenceModeling_output, len(pretrained_dict['module.Prediction.weight']))  
+        
+        model = torch.nn.DataParallel(model).to(device) 
+        print(f'loading pretrained model from {opt.saved_model}')
+        if opt.FT:
+            model.load_state_dict(pretrained_dict, strict=False)
+        else:
+            model.load_state_dict(pretrained_dict)
+        if opt.new_prediction:
+            model.module.Prediction = nn.Linear(model.module.SequenceModeling_output, opt.num_class)  
+            for name, param in model.module.Prediction.named_parameters():
+                if 'bias' in name:
+                    init.constant_(param, 0.0)
+                elif 'weight' in name:
+                    init.kaiming_normal_(param)
+            model = model.to(device) 
+    else:
+        # weight initialization
+        for name, param in model.named_parameters():
+            if 'localization_fc2' in name:
+                print(f'Skip {name} as it is already initialized')
+                continue
+            try:
+                if 'bias' in name:
+                    init.constant_(param, 0.0)
+                elif 'weight' in name:
+                    init.kaiming_normal_(param)
+            except Exception as e:  # for batchnorm.
+                if 'weight' in name:
+                    param.data.fill_(1)
+                continue
+        model = torch.nn.DataParallel(model).to(device)
+    
+    model.train() 
+    print("Model:")
+    print(model)
+    count_parameters(model)
+    
+    """ setup loss """
+    if 'CTC' in opt.Prediction:
+        criterion = torch.nn.CTCLoss(zero_infinity=True).to(device)
+    else:
+        criterion = torch.nn.CrossEntropyLoss(ignore_index=0).to(device)  # ignore [GO] token = ignore index 0
+    # loss averager
+    loss_avg = Averager()
+
+    # freeze some layers
+    try:
+        if opt.freeze_FeatureFxtraction:
+            for param in model.module.FeatureExtraction.parameters():
+                param.requires_grad = False
+        if opt.freeze_SequenceModeling:
+            for param in model.module.SequenceModeling.parameters():
+                param.requires_grad = False
+    except:
+        pass
+    
+    # filter that only require gradient decent
+    filtered_parameters = []
+    params_num = []
+    for p in filter(lambda p: p.requires_grad, model.parameters()):
+        filtered_parameters.append(p)
+        params_num.append(np.prod(p.size()))
+    print('Trainable params num : ', sum(params_num))
+    # [print(name, p.numel()) for name, p in filter(lambda p: p[1].requires_grad, model.named_parameters())]
+
+    # setup optimizer
+    if opt.optim=='adam':
+        #optimizer = optim.Adam(filtered_parameters, lr=opt.lr, betas=(opt.beta1, 0.999))
+        optimizer = optim.Adam(filtered_parameters)
+    else:
+        optimizer = optim.Adadelta(filtered_parameters, lr=opt.lr, rho=opt.rho, eps=opt.eps)
+    print("Optimizer:")
+    print(optimizer)
+
+    """ final options """
+    # print(opt)
+    with open(f'./saved_models/{opt.experiment_name}/opt.txt', 'a', encoding="utf8") as opt_file:
+        opt_log = '------------ Options -------------\n'
+        args = vars(opt)
+        for k, v in args.items():
+            opt_log += f'{str(k)}: {str(v)}\n'
+        opt_log += '---------------------------------------\n'
+        print(opt_log)
+        opt_file.write(opt_log)
+
+    """ start training """
+    start_iter = 0
+    if opt.saved_model != '':
+        try:
+            start_iter = int(opt.saved_model.split('_')[-1].split('.')[0])
+            print(f'continue to train, start_iter: {start_iter}')
+        except:
+            pass
+
+    start_time = time.time()
+    best_accuracy = -1
+    best_norm_ED = -1
+    i = start_iter
+
+    scaler = GradScaler()
+    t1= time.time()
+        
+    while(True):
+        # train part
+        optimizer.zero_grad(set_to_none=True)
+        
+        if amp:
+            with autocast():
+                image_tensors, labels = train_dataset.get_batch()
+                image = image_tensors.to(device)
+                text, length = converter.encode(labels, batch_max_length=opt.batch_max_length)
+                batch_size = image.size(0)
+
+                if 'CTC' in opt.Prediction:
+                    preds = model(image, text).log_softmax(2)
+                    preds_size = torch.IntTensor([preds.size(1)] * batch_size)
+                    preds = preds.permute(1, 0, 2)
+                    torch.backends.cudnn.enabled = False
+                    cost = criterion(preds, text.to(device), preds_size.to(device), length.to(device))
+                    torch.backends.cudnn.enabled = True
+                else:
+                    preds = model(image, text[:, :-1])  # align with Attention.forward
+                    target = text[:, 1:]  # without [GO] Symbol
+                    cost = criterion(preds.view(-1, preds.shape[-1]), target.contiguous().view(-1))
+            scaler.scale(cost).backward()
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), opt.grad_clip)
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            image_tensors, labels = train_dataset.get_batch()
+            image = image_tensors.to(device)
+            text, length = converter.encode(labels, batch_max_length=opt.batch_max_length)
+            batch_size = image.size(0)
+            if 'CTC' in opt.Prediction:
+                preds = model(image, text).log_softmax(2)
+                preds_size = torch.IntTensor([preds.size(1)] * batch_size)
+                preds = preds.permute(1, 0, 2)
+                torch.backends.cudnn.enabled = False
+                cost = criterion(preds, text.to(device), preds_size.to(device), length.to(device))
+                torch.backends.cudnn.enabled = True
+            else:
+                preds = model(image, text[:, :-1])  # align with Attention.forward
+                target = text[:, 1:]  # without [GO] Symbol
+                cost = criterion(preds.view(-1, preds.shape[-1]), target.contiguous().view(-1))
+            cost.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), opt.grad_clip) 
+            optimizer.step()
+        loss_avg.add(cost)
+
+        # validation part
+        if (i % opt.valInterval == 0) and (i!=0):
+            print('training time: ', time.time()-t1)
+            t1=time.time()
+            elapsed_time = time.time() - start_time
+            # for log
+            with open(f'./saved_models/{opt.experiment_name}/log_train.txt', 'a', encoding="utf8") as log:
+                model.eval()
+                with torch.no_grad():
+                    valid_loss, current_accuracy, current_norm_ED, preds, confidence_score, labels,\
+                    infer_time, length_of_data = validation(model, criterion, valid_loader, converter, opt, device)
+                model.train()
+
+                # training loss and validation loss
+                loss_log = f'[{i}/{opt.num_iter}] Train loss: {loss_avg.val():0.5f}, Valid loss: {valid_loss:0.5f}, Elapsed_time: {elapsed_time:0.5f}'
+                loss_avg.reset()
+
+                current_model_log = f'{"Current_accuracy":17s}: {current_accuracy:0.3f}, {"Current_norm_ED":17s}: {current_norm_ED:0.4f}'
+
+                # keep best accuracy model (on valid dataset)
+                if current_accuracy > best_accuracy:
+                    best_accuracy = current_accuracy
+                    torch.save(model.state_dict(), f'./saved_models/{opt.experiment_name}/best_accuracy.pth')
+                if current_norm_ED > best_norm_ED:
+                    best_norm_ED = current_norm_ED
+                    torch.save(model.state_dict(), f'./saved_models/{opt.experiment_name}/best_norm_ED.pth')
+                best_model_log = f'{"Best_accuracy":17s}: {best_accuracy:0.3f}, {"Best_norm_ED":17s}: {best_norm_ED:0.4f}'
+
+                loss_model_log = f'{loss_log}\n{current_model_log}\n{best_model_log}'
+                print(loss_model_log)
+                log.write(loss_model_log + '\n')
+
+                # show some predicted results
+                dashed_line = '-' * 80
+                head = f'{"Ground Truth":25s} | {"Prediction":25s} | Confidence Score & T/F'
+                predicted_result_log = f'{dashed_line}\n{head}\n{dashed_line}\n'
+                
+                #show_number = min(show_number, len(labels))
+                
+                start = random.randint(0,len(labels) - show_number )    
+                for gt, pred, confidence in zip(labels[start:start+show_number], preds[start:start+show_number], confidence_score[start:start+show_number]):
+                    if 'Attn' in opt.Prediction:
+                        gt = gt[:gt.find('[s]')]
+                        pred = pred[:pred.find('[s]')]
+
+                    predicted_result_log += f'{gt:25s} | {pred:25s} | {confidence:0.4f}\t{str(pred == gt)}\n'
+                predicted_result_log += f'{dashed_line}'
+                print(predicted_result_log)
+                log.write(predicted_result_log + '\n')
+                print('validation time: ', time.time()-t1)
+                t1=time.time()
+        # save model per 1e+4 iter.
+        if (i + 1) % 1e+4 == 0:
+            torch.save(
+                model.state_dict(), f'./saved_models/{opt.experiment_name}/iter_{i+1}.pth')
+
+        if i == opt.num_iter:
+            print('end the training')
+            sys.exit()
+        i += 1