OpenOCR: A general OCR system with accuracy and efficiency
+ +If you find this project useful, please give us a star🌟.
+ +
+
+
+
+
+
+
+
+
+
+ 🚀 Quick Start | English | [简体中文](./README_ch.md)
+
+
+
+
+
+
+
+
+
+
+
+ 🚀 Quick Start | English | [简体中文](./README_ch.md)
+
+
+
+
+
+
+
+
+
+
+
+ 🚀 快速开始 | 简体中文 | [English](./README.md)
+
+准确高效的通用 OCR 系统 （由FVL实验室 OCR Team 创建） [本地快速部署]
""" + ) + with gr.Row(): + with gr.Column(scale=1): + input_image = gr.Image(label='Input image', + elem_classes=['image-container']) + + examples = gr.Examples(examples=e2e_img_example, + inputs=input_image, + label='Examples') + downstream = gr.Button('Run') + + # 添加参数调节组件 + with gr.Column(): + with gr.Row(): + det_input_size_textbox = gr.Number( + label='Detection Input Size', + value=960, + info='检测网络输入尺寸的最长边，默认为960。') + det_score_mode_dropdown = gr.Dropdown( + ['slow', 'fast'], + value='slow', + label='Detection Score Mode', + info='文本框的置信度计算模式，默认为 slow。slow 模式计算速度较慢，但准确度较高。fast 模式计算速度较快，但准确度较低。' + ) + with gr.Row(): + rec_drop_score_slider = gr.Slider( + 0.0, + 1.0, + value=0.4, + step=0.01, + label='Recognition Drop Score', + info='识别置信度阈值，默认值为0.4。低于该阈值的识别结果和对应的文本框被丢弃。') + mask_thresh_slider = gr.Slider( + 0.0, + 1.0, + value=0.3, + step=0.01, + label='Mask Threshold', + info='Mask 阈值，用于二值化 mask，默认值为0.3。如果存在文本截断时，请调低该值。') + with gr.Row(): + box_thresh_slider = gr.Slider( + 0.0, + 1.0, + value=0.6, + step=0.01, + label='Box Threshold', + info='文本框置信度阈值，默认值为0.6。如果存在文本被漏检时，请调低该值。') + unclip_ratio_slider = gr.Slider( + 1.5, + 2.0, + value=1.5, + step=0.05, + label='Unclip Ratio', + info='文本框解析时的膨胀系数，默认值为1.5。值越大文本框越大。') + + # 模型选择组件 + model_type_dropdown = gr.Dropdown( + ['mobile', 'server'], + value='mobile', + label='Model Type', + info='选择 OCR 模型类型：高效率模型mobile，高精度模型server。') + + with gr.Column(scale=1): + img_mask = gr.Image(label='mask', + interactive=False, + elem_classes=['image-container']) + img_output = gr.Image(label=' ', + interactive=False, + elem_classes=['image-container']) + + output = gr.Textbox(label='Result') + confidence = gr.Textbox(label='Latency') + + downstream.click(fn=main, + inputs=[ + input_image, model_type_dropdown, + det_input_size_textbox, rec_drop_score_slider, + mask_thresh_slider, box_thresh_slider, + unclip_ratio_slider, det_score_mode_dropdown + ], + outputs=[ + output, + confidence, + img_output, + img_mask, + ]) + + demo.launch(share=True) diff --git a/demo_rec.py b/demo_rec.py deleted file mode 100644 index a1c604a..0000000 --- a/demo_rec.py +++ /dev/null @@ -1,131 +0,0 @@ -import os - -import gradio as gr -import numpy as np -import torch - -from openrec.modeling import build_model -from openrec.postprocess import build_post_process -from openrec.preprocess import create_operators, transform -from tools.engine import Config -from tools.utils.ckpt import load_ckpt - - -def build_rec_process(cfg): - transforms = [] - for op in cfg['Eval']['dataset']['transforms']: - op_name = list(op)[0] - if 'Label' in op_name: - continue - # TODO - elif op_name in ['DecodeImage']: - op[op_name]['gradio_infer_mode'] = True - - elif op_name in ['RecResizeImg']: - op[op_name]['infer_mode'] = True - elif op_name == 'KeepKeys': - if cfg['Architecture']['algorithm'] == 'SRN': - op[op_name]['keep_keys'] = [ - 'image', - 'encoder_word_pos', - 'gsrm_word_pos', - 'gsrm_slf_attn_bias1', - 'gsrm_slf_attn_bias2', - ] - elif cfg['Architecture']['algorithm'] == 'SAR': - op[op_name]['keep_keys'] = ['image', 'valid_ratio'] - elif cfg['Architecture']['algorithm'] == 'RobustScanner': - op[op_name]['keep_keys'] = [ - 'image', 'valid_ratio', 'word_positons' - ] - else: - op[op_name]['keep_keys'] = ['image'] - transforms.append(op) - return transforms - - -def get_all_file_names_including_subdirs(dir_path): - all_file_names = [] - - for root, dirs, files in os.walk(dir_path): - for file_name in files: - all_file_names.append(os.path.join(root, file_name)) - - file_names_only = [os.path.basename(file) for file in all_file_names] - return file_names_only - - -root_directory = './configs/rec' -yml_Config = get_all_file_names_including_subdirs(root_directory) - - -def find_file_in_current_dir_and_subdirs(file_name): - for root, dirs, files in os.walk('.'): - if file_name in files: - relative_path = os.path.join(root, file_name) - return relative_path - - -def predict(input_image, Model_type, OCR_type): - - path = find_file_in_current_dir_and_subdirs(Model_type) - - cfg = Config(path).cfg - post_process_class = build_post_process(cfg['PostProcess']) - global_config = cfg['Global'] - char_num = len(getattr(post_process_class, 'character')) - cfg['Architecture']['Decoder']['out_channels'] = char_num - model = build_model(cfg['Architecture']) - load_ckpt(model, cfg) - model.eval() - - transforms = build_rec_process(cfg) - global_config['infer_mode'] = True - ops = create_operators(transforms, global_config) - data = {'image': input_image} - batch = transform(data, ops) - others = None - images = np.expand_dims(batch[0], axis=0) - images = torch.from_numpy(images) - with torch.no_grad(): - preds = model(images, others) - post_result = post_process_class(preds) - return post_result[0][0], post_result[0][1] - - -if __name__ == '__main__': - - with gr.Blocks() as demo: - with gr.Row(): - with gr.Column(scale=1): - input_image = gr.Image(label='Input Image') - - # TODO - OCR_type = gr.Radio(['STR', 'STD', 'E2E'], label='模型类别') - - Model_type = gr.Dropdown(choices=yml_Config, label='现有模型配置文件') - - downstream = gr.Button('识别结果') - - with gr.Column(scale=1): - - # TODO - img_output = gr.Image(label='图片识别结果') - - output = gr.Textbox(label='文字识别结果') - confidence = gr.Textbox(label='置信度') - - downstream.click( - fn=predict, - inputs=[ - input_image, - Model_type, - OCR_type, - ], - outputs=[ - output, - confidence, - # TODO img_output, - ]) - - demo.launch(debug=True) diff --git a/docs/finetune_det.md b/docs/finetune_det.md new file mode 100644 index 0000000..2cacffd --- /dev/null +++ b/docs/finetune_det.md @@ -0,0 +1,159 @@ +# Fine-tuning Text Detection Model of OpenOCR System + +1. [Data and Weights Preparation](#1-data-and-weights-preparation) + - [1.1 Data Preparation](#11-data-preparation) + - [1.2 Download Pre-trained Model](#12-download-pre-trained-model) +2. [Training](#2-training) + - [2.1 Start Training](#21-start-training) + - [2.2 Load Trained Model and Continue Training](#22-load-trained-model-and-continue-training) +3. [Evaluation and Test](#3-evaluation-and-test) + - [3.1 Evaluation](#31-evaluation) + - [3.2 Test](#32-test) +4. [ONNX Inference](#4-onnx-inference) + +______________________________________________________________________ + +## Installation + +#### Dependencies: + +- [PyTorch](http://pytorch.org/) version >= 1.13.0 +- Python version >= 3.7 + +```shell +conda create -n openocr python==3.8 +conda activate openocr +# install gpu version torch +conda install pytorch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 pytorch-cuda=11.8 -c pytorch -c nvidia +# or cpu version +conda install pytorch torchvision torchaudio cpuonly -c pytorch +``` + +#### Clone this repository: + +```shell +git clone https://github.com/Topdu/OpenOCR.git +cd OpenOCR +pip install albumentations +pip install -r requirements.txt +``` + +This section uses the icdar2015 dataset as an example to introduce the training, evaluation, and testing of the detection model in OpenOCR. + +## 1. Data and Weights Preparation + +### 1.1 Data Preparation + +**Note:** If you want to use your own dataset, please following the format of [icdar2015 dataset](https://aistudio.baidu.com/datasetdetail/46088). + +Downloading datasets from [icdar2015 dataset](https://aistudio.baidu.com/datasetdetail/46088)/[Google Drive](https://drive.google.com/file/d/1nfsYj-JzAqVouZPBDqmuP0Rkj6J6XFUJ/view?usp=sharing). + +#### File Directory + +``` +OpenOCR/ +icdar2015/text_localization/ + └─ icdar_c4_train_imgs/ Training data of the icdar dataset + └─ ch4_test_images/ Testing data of the icdar dataset + └─ train_icdar2015_label.txt Training annotations of the icdar dataset + └─ test_icdar2015_label.txt Testing annotations of the icdar dataset +``` + +The provided annotation file format is as follows, where the fields are separated by "\\t": + +``` +"Image file name json.dumps encoded image annotation information" +ch4_test_images/img_61.jpg [{"transcription": "MASA", "points": [[310, 104], [416, 141], [418, 216], [312, 179]], ...}] +``` + +Before being encoded with `json.dumps`, the image annotation information is a list containing multiple dictionaries. In each dictionary, the field `points` represents the coordinates (x, y) of the four corners of the text bounding box, arranged in a clockwise order starting from the top-left corner. The field `transcription` indicates the text content within the current bounding box. + +To modify the training and evaluation dataset paths in the configuration file `./configs/det/dbnet/repvit_db.yml` to your own dataset paths, for example: + +```yaml +Train: + dataset: + name: SimpleDataSet + data_dir: ../icdar2015/text_localization/ # Root directory of the training dataset + label_file_list: ["../icdar2015/text_localization/train_icdar2015_label.txt"] # Path to the training label file + ...... +Eval: + dataset: + name: SimpleDataSet + data_dir: ../icdar2015/text_localization/ # Root directory of the evaluation dataset + label_file_list: ["../icdar2015/text_localization/test_icdar2015_label.txt"] # Path to the evaluation label file +``` + +### 1.2 Download Pre-trained Model + +First download the pre-trained model. + +```bash +cd OpenOCR/ +wget https://github.com/Topdu/OpenOCR/releases/download/develop0.0.1/openocr_det_repvit_ch.pth +``` + +______________________________________________________________________ + +## 2. Training + +### 2.1 Start Training + +```bash +# multi-GPU training +CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 tools/train_det.py --c configs/det/dbnet/repvit_db.yml --o Global.pretrained_model=./openocr_det_repvit_ch.pth +# single GPU training +CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node=1 tools/train_det.py --c configs/det/dbnet/repvit_db.yml --o Global.pretrained_model=./openocr_det_repvit_ch.pth +``` + +### 2.2 Load Trained Model and Continue Training + +If you expect to load trained model and continue the training again, you can specify the parameter `Global.checkpoints` as the model path to be loaded. + +For example: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 tools/train_det.py --c configs/det/dbnet/repvit_db.yml --o Global.checkpoints=./your/trained/model +``` + +**Note**: The priority of `Global.checkpoints` is higher than that of `Global.pretrained_model`, that is, when two parameters are specified at the same time, the model specified by `Global.checkpoints` will be loaded first. If the model path specified by `Global.checkpoints` is wrong, the one specified by `Global.pretrained_model` will be loaded. + +______________________________________________________________________ + +## 3. Evaluation and Test + +### 3.1 Evaluation + +OpenOCR calculates three indicators for evaluating performance of OCR detection task: Precision, Recall, and Hmean(F-Score). + +```bash +python tools/eval_det.py --c configs/det/dbnet/repvit_db.yml --o Global.pretrained_model="{path/to/weights}/best.pth" +``` + +### 3.2 Test + +Test the detection result on all images in the folder or a single image: + +```bash +python tools/infer_det.py --c ./configs/det/dbnet/repvit_db.yml --o Global.infer_img=/path/img_fold or /path/img_file Global.pretrained_model={path/to/weights}/best.pth +``` + +______________________________________________________________________ + +## 4. ONNX Inference + +Firstly, we can convert Detection model to onnx model: + +```bash +pip install onnx +python tools/toonnx.py --c ./configs/det/dbnet/repvit_db.yml --o Global.device=cpu Global.pretrained_model={path/to/weights}/best.pth +``` + +The onnx model is saved in `./output/det_repsvtr_db/export_det/det_model.onnx`. + +The detection onnx model inference: + +```bash +pip install onnxruntime +python tools/infer_det.py --c ./configs/det/dbnet/repvit_db.yml --o Global.backend=onnx Global.device=cpu Global.infer_img=/path/img_fold or /path/img_file Global.onnx_model_path=./output/det_repsvtr_db/export_det/det_model.onnx +``` diff --git a/docs/finetune_rec.md b/docs/finetune_rec.md new file mode 100644 index 0000000..27c162e --- /dev/null +++ b/docs/finetune_rec.md @@ -0,0 +1,150 @@ +# Fine-tuning Text Recognition Model of OpenOCR system + +1. [Data and Weights Preparation](#1-data-and-weights-preparation) + - [1.1 Data Preparation](#11-data-preparation) + - [1.2 Download Pre-trained Model](#12-download-pre-trained-model) +2. [Training](#2-training) + - [2.1 Start Training](#21-start-training) + - [2.2 Load Trained Model and Continue Training](#22-load-trained-model-and-continue-training) +3. [Evaluation and Test](#3-evaluation-and-test) + - [3.1 Evaluation](#31-evaluation) + - [3.2 Test](#32-test) +4. [ONNX Inference](#4-onnx-inference) + +## Installation + +#### Dependencies: + +- [PyTorch](http://pytorch.org/) version >= 1.13.0 +- Python version >= 3.7 + +```shell +conda create -n openocr python==3.8 +conda activate openocr +# install gpu version torch +conda install pytorch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 pytorch-cuda=11.8 -c pytorch -c nvidia +# or cpu version +conda install pytorch torchvision torchaudio cpuonly -c pytorch +``` + +#### Clone this repository: + +```shell +git clone https://github.com/Topdu/OpenOCR.git +cd OpenOCR +pip install -r requirements.txt +``` + +This section uses the icdar2015 recognition dataset as an example to introduce the training, evaluation, and testing of the recognition model in OpenOCR. + +## 1. Data and Weights Preparation + +### 1.1 Data Preparation + +**Note:** If you want to use your own dataset, please following the following data format. + +Downloading datasets from [icdar2015 recognition dataset](https://aistudio.baidu.com/datasetdetail/75418)/[Google Drive](https://drive.google.com/file/d/1YviGN_f7xrRrMOSR4OGwv7uhKFjnxuUP/view?usp=sharing). + +#### File Directory + +``` +OpenOCR/ +ic15_data/ + └─ test/ Training data of the icdar dataset + └─ train/ Testing data of the icdar dataset + └─ rec_gt_test.txt Training annotations of the icdar dataset + └─ rec_gt_train.txt Testing annotations of the icdar dataset +``` + +The provided annotation file format is as follows, where the fields are separated by "\\t": + +``` +"Image file name label" +test/word_2077.png Underpass +``` + +To modify the training and evaluation dataset paths in the configuration file `./configs/rec/svtrv2/repsvtr_ch.yml` to your own dataset paths, for example: + +```yaml +Train: + dataset: + name: SimpleDataSet + data_dir: ../ic15_data/ # Root directory of the training dataset + label_file_list: ["../ic15_data/rec_gt_train.txt"] # Path to the training label file + ...... +Eval: + dataset: + name: SimpleDataSet + data_dir: ../ic15_data # Root directory of the evaluation dataset + label_file_list: ["../ic15_data/rec_gt_test.txt"] # Path to the evaluation label file +``` + +### 1.2 Download Pre-trained Model + +First download the pre-trained model. + +```bash +cd OpenOCR/ +wget https://github.com/Topdu/OpenOCR/releases/download/develop0.0.1/openocr_repsvtr_ch.pth +# Rec Server model +# wget https://github.com/Topdu/OpenOCR/releases/download/develop0.0.1/openocr_svtrv2_ch.pth +``` + +## 2. Training + +### 2.1 Start Training + +```bash +# multi-GPU training +CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 tools/train_rec.py --c ./configs/rec/svtrv2/repsvtr_ch.yml --o Global.pretrained_model=./openocr_repsvtr_ch.pth +# single GPU training +CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node=1 tools/train_rec.py --c ./configs/rec/svtrv2/repsvtr_ch.yml --o Global.pretrained_model=./openocr_repsvtr_ch.pth +``` + +### 2.2 Load Trained Model and Continue Training + +If you expect to load trained model and continue the training again, you can specify the parameter `Global.checkpoints` as the model path to be loaded. + +For example: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 tools/train_rec.py --c ./configs/rec/svtrv2/repsvtr_ch.yml --o Global.checkpoints=./your/trained/model +``` + +**Note**: The priority of `Global.checkpoints` is higher than that of `Global.pretrained_model`, that is, when two parameters are specified at the same time, the model specified by `Global.checkpoints` will be loaded first. If the model path specified by `Global.checkpoints` is wrong, the one specified by `Global.pretrained_model` will be loaded. + +## 3. Evaluation and Test + +### 3.1 Evaluation + +OpenOCR calculates the word accuracy for evaluating performance of OCR recognition task. + +```bash +python tools/eval_rec.py --c ./configs/rec/svtrv2/repsvtr_ch.yml --o Global.pretrained_model="{path/to/weights}/best.pth" +``` + +### 3.2 Test + +Test the recognition result on all images in the folder or a single image: + +```bash +python tools/infer_rec.py --c ./configs/rec/svtrv2/repsvtr_ch.yml --o Global.infer_img=/path/img_fold or /path/img_file Global.pretrained_model={path/to/weights}/best.pth +``` + +## 4. ONNX Inference + +Firstly, we can convert recognition model to onnx model: + +```bash +pip install onnx +python tools/toonnx.py --c ./configs/rec/svtrv2/repsvtr_ch.yml --o Global.device=cpu Global.pretrained_model={path/to/weights}/best.pth +``` + +The onnx model is saved in `./output/rec/repsvtr_ch/export_rec/rec_model.onnx`. + +The recognition onnx model infernce: + +```bash +pip install onnxruntime +python tools/infer_rec.py --c ./configs/rec/svtrv2/repsvtr_ch.yml --o Global.backend=onnx Global.device=cpu Global.infer_img=/path/img_fold or /path/img_file Global.onnx_model_path=./output/rec/repsvtr_ch/export_rec/rec_model.onnx +``` diff --git a/docs/openocr.md b/docs/openocr.md index 4129747..137198b 100644 --- a/docs/openocr.md +++ b/docs/openocr.md @@ -1,29 +1,58 @@ -# OpenOCR: A general OCR system for accuracy and efficiency +# OpenOCR: A general OCR system with accuracy and efficiency -We proposed strategies to comprehensively enhance CTC-based STR models and developed a novel CTC-based method, [SVTRv2](../configs/rec/svtrv2/). SVTRv2 can outperform previous attention-based STR methods in terms of accuracy while maintaining the advantages of CTC, such as fast inference and robust recognition of long text sequences. These features make SVTRv2 particularly well-suited for commercial applications. To this end, building on SVTRv2, we develop a practical version of the model from scratch on publicly available Chinese and English datasets. Combined with a detection model, this forms an accurate and efficient general OCR system, OpenOCR. Comparing with PP-OCRv4 released by PaddleOCR, OpenOCR achieve a 4.5% improvement on the [OCR competition leaderboard](https://aistudio.baidu.com/competition/detail/1131/0/leaderboard). +⚡\[[Quick Start](#quick-start)\] \[[Model](https://github.com/Topdu/OpenOCR/releases/tag/develop0.0.1)\] \[[ModelScope Demo](https://modelscope.cn/studios/topdktu/OpenOCR-Demo)\] \[[Hugging Face Demo](https://huggingface.co/spaces/topdu/OpenOCR-Demo)\] \[[Local Demo](#local-demo)\] \[[PaddleOCR Implementation](https://paddlepaddle.github.io/PaddleOCR/latest/algorithm/text_recognition/algorithm_rec_svtrv2.html)\] + +We proposed strategies to comprehensively enhance CTC-based STR models and developed a novel CTC-based method, [SVTRv2](../configs/rec/svtrv2/). SVTRv2 can outperform previous attention-based STR methods in terms of accuracy while maintaining the advantages of CTC, such as fast inference and robust recognition of long text. These features make SVTRv2 particularly well-suited for practical applications. To this end, building on SVTRv2, we develop a practical version of the model from scratch on publicly available Chinese and English datasets. Combined with a detection model, this forms a general OCR system with accuracy and efficiency, **OpenOCR**. Comparing with [PP-OCRv4](https://paddlepaddle.github.io/PaddleOCR/latest/ppocr/model_list.html) baseline in the [OCR competition leaderboard](https://aistudio.baidu.com/competition/detail/1131/0/leaderboard), OpenOCR (mobile) achieve a 4.5% improvement in terms of accuracy, while preserving quite similar inference speed on NVIDIA 1080Ti GPU. | Model | Config | E2E Metric | Downloading | | ------------------- | ----------------------------------------------------------------------------------- | ---------- | ---------------------------------------------------------------------------------------- | | PP-OCRv4 | | 62.77% | [PaddleOCR Model List](../../ppocr/model_list.md) | -| SVTRv2 (Rec Server) | [configs/rec/svtrv2/svtrv2_ch.yml](../configs/rec/svtrv2/svtrv2_ch.yml) | 68.81% | [Google Dirve ](https://drive.google.com/file/d/13LXbIVEyx2Aat3X_vVte4JQgQ7yJWdxH/view?usp=drive_link) | -| RepSVTR (Mobile) | [Rec: configs/rec/svtrv2/repsvtr_ch.yml](../configs/rec/svtrv2/repsvtr_ch.yml)| Method | @@ -1678,10 +1647,17 @@ Downloading all model checkpoints from [Google Drive](<>) and [Baidu Yun](<>). ## Citation +If you find our method useful for your reserach, please cite: + ```bibtex -@article{Du2024SVTRv4, - title = {SVTRv2: Scene Text Recognition with a Single Visual Model}, - author = {Yongkun Du, Zhineng Chen\*, Hongtao Xie, Caiyan Jia, Yu-Gang Jiang}, - year = {2024} +@article{Du2024SVTRv2, + title={SVTRv2: CTC Beats Encoder-Decoder Models in Scene Text Recognition}, + author={Yongkun Du and Zhineng Chen and Hongtao Xie and Caiyan Jia and Yu-Gang Jiang}, + journal={CoRR}, + volume={abs/2411.15858}, + eprinttype={arXiv}, + year={2024}, + primaryClass={cs.CV}, + url={https://arxiv.org/abs/2411.15858} } ``` diff --git a/opendet/losses/__init__.py b/opendet/losses/__init__.py new file mode 100644 index 0000000..f3dd72c --- /dev/null +++ b/opendet/losses/__init__.py @@ -0,0 +1,20 @@ +import copy +from importlib import import_module + +name_to_module = { + 'DBLoss': '.db_loss', +} + + +def build_loss(config): + config = copy.deepcopy(config) + module_name = config.pop('name') + assert module_name in name_to_module, Exception( + '{} is not supported. The losses in {} are supportes'.format( + module_name, list(name_to_module.keys()))) + + module_path = name_to_module[module_name] + module = import_module(module_path, package=__package__) + module_class = getattr(module, module_name) + + return module_class(**config) diff --git a/opendet/losses/db_loss.py b/opendet/losses/db_loss.py new file mode 100644 index 0000000..7563dd6 --- /dev/null +++ b/opendet/losses/db_loss.py @@ -0,0 +1,87 @@ +# copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve. +# +# 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. +""" +This code is refer from: +https://github.com/WenmuZhou/DBNet.pytorch/blob/master/models/losses/DB_loss.py +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import torch +from torch import nn + +from .det_basic_loss import BalanceLoss, MaskL1Loss, DiceLoss + + +class DBLoss(nn.Module): + """ + Differentiable Binarization (DB) Loss Function + args: + param (dict): the super paramter for DB Loss + """ + + def __init__(self, + balance_loss=True, + main_loss_type='DiceLoss', + alpha=5, + beta=10, + ohem_ratio=3, + eps=1e-6, + **kwargs): + super(DBLoss, self).__init__() + self.alpha = alpha + self.beta = beta + self.dice_loss = DiceLoss(eps=eps) + self.l1_loss = MaskL1Loss(eps=eps) + self.bce_loss = BalanceLoss(balance_loss=balance_loss, + main_loss_type=main_loss_type, + negative_ratio=ohem_ratio) + + def forward(self, predicts, labels): + predict_maps = predicts['maps'] + label_threshold_map, label_threshold_mask, label_shrink_map, label_shrink_mask = labels[ + 1:] + shrink_maps = predict_maps[:, 0, :, :] + threshold_maps = predict_maps[:, 1, :, :] + binary_maps = predict_maps[:, 2, :, :] + + loss_shrink_maps = self.bce_loss(shrink_maps, label_shrink_map, + label_shrink_mask) + loss_threshold_maps = self.l1_loss(threshold_maps, label_threshold_map, + label_threshold_mask) + loss_binary_maps = self.dice_loss(binary_maps, label_shrink_map, + label_shrink_mask) + loss_shrink_maps = self.alpha * loss_shrink_maps + loss_threshold_maps = self.beta * loss_threshold_maps + # CBN loss + if 'distance_maps' in predicts.keys(): + # distance_maps = predicts['distance_maps'] + cbn_maps = predicts['cbn_maps'] + cbn_loss = self.bce_loss(cbn_maps[:, 0, :, :], label_shrink_map, + label_shrink_mask) + else: + # dis_loss = torch.tensor([0.]) + cbn_loss = torch.tensor([0.], device=predict_maps.device) + + loss_all = loss_shrink_maps + loss_threshold_maps + loss_binary_maps + losses = { + 'loss': loss_all + cbn_loss, + 'loss_shrink_maps': loss_shrink_maps, + 'loss_threshold_maps': loss_threshold_maps, + 'loss_binary_maps': loss_binary_maps, + 'loss_cbn': cbn_loss + } + return losses diff --git a/opendet/losses/det_basic_loss.py b/opendet/losses/det_basic_loss.py new file mode 100644 index 0000000..66a2f6f --- /dev/null +++ b/opendet/losses/det_basic_loss.py @@ -0,0 +1,159 @@ +# copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve. +# +# 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. +""" +This code is refer from: +https://github.com/WenmuZhou/DBNet.pytorch/blob/master/models/losses/basic_loss.py +""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import torch +from torch import nn +import torch.nn.functional as F + + +class BalanceLoss(nn.Module): + + def __init__( + self, + balance_loss=True, + main_loss_type='DiceLoss', + negative_ratio=3, + return_origin=False, + eps=1e-6, + **kwargs, + ): + """ + The BalanceLoss for Differentiable Binarization text detection + args: + balance_loss (bool): whether balance loss or not, default is True + main_loss_type (str): can only be one of ['CrossEntropy','DiceLoss', + 'Euclidean','BCELoss', 'MaskL1Loss'], default is 'DiceLoss'. + negative_ratio (int|float): float, default is 3. + return_origin (bool): whether return unbalanced loss or not, default is False. + eps (float): default is 1e-6. + """ + super(BalanceLoss, self).__init__() + self.balance_loss = balance_loss + self.main_loss_type = main_loss_type + self.negative_ratio = negative_ratio + self.return_origin = return_origin + self.eps = eps + + if self.main_loss_type == 'CrossEntropy': + self.loss = nn.CrossEntropyLoss() + elif self.main_loss_type == 'Euclidean': + self.loss = nn.MSELoss() + elif self.main_loss_type == 'DiceLoss': + self.loss = DiceLoss(self.eps) + elif self.main_loss_type == 'BCELoss': + self.loss = BCELoss(reduction='none') + elif self.main_loss_type == 'MaskL1Loss': + self.loss = MaskL1Loss(self.eps) + else: + loss_type = [ + 'CrossEntropy', + 'DiceLoss', + 'Euclidean', + 'BCELoss', + 'MaskL1Loss', + ] + raise Exception( + 'main_loss_type in BalanceLoss() can only be one of {}'.format( + loss_type)) + + def forward(self, pred, gt, mask=None): + """ + The BalanceLoss for Differentiable Binarization text detection + args: + pred (variable): predicted feature maps. + gt (variable): ground truth feature maps. + mask (variable): masked maps. + return: (variable) balanced loss + """ + positive = gt * mask + negative = (1 - gt) * mask + + positive_count = int(positive.sum()) + negative_count = int( + min(negative.sum(), positive_count * self.negative_ratio)) + loss = self.loss(pred, gt, mask=mask) + + if not self.balance_loss: + return loss + + positive_loss = positive * loss + negative_loss = negative * loss + negative_loss = negative_loss.reshape(-1) + if negative_count > 0: + negative_loss, _ = negative_loss.topk(negative_count) + balance_loss = (positive_loss.sum() + negative_loss.sum()) / ( + positive_count + negative_count + self.eps) + else: + balance_loss = positive_loss.sum() / (positive_count + self.eps) + if self.return_origin: + return balance_loss, loss + + return balance_loss + + +class DiceLoss(nn.Module): + + def __init__(self, eps=1e-6): + super(DiceLoss, self).__init__() + self.eps = eps + + def forward(self, pred, gt, mask, weights=None): + """ + DiceLoss function. + """ + + assert pred.shape == gt.shape + assert pred.shape == mask.shape + if weights is not None: + assert weights.shape == mask.shape + mask = weights * mask + intersection = torch.sum(pred * gt * mask) + + union = torch.sum(pred * mask) + torch.sum(gt * mask) + self.eps + loss = 1 - 2.0 * intersection / union + assert loss <= 1 + return loss + + +class MaskL1Loss(nn.Module): + + def __init__(self, eps=1e-6): + super(MaskL1Loss, self).__init__() + self.eps = eps + + def forward(self, pred, gt, mask): + """ + Mask L1 Loss + """ + loss = (torch.abs(pred - gt) * mask).sum() / (mask.sum() + self.eps) + loss = torch.mean(loss) + return loss + + +class BCELoss(nn.Module): + + def __init__(self, reduction='mean'): + super(BCELoss, self).__init__() + self.reduction = reduction + + def forward(self, input, label, mask=None, weight=None, name=None): + loss = F.binary_cross_entropy(input, label, reduction=self.reduction) + return loss diff --git a/opendet/metrics/__init__.py b/opendet/metrics/__init__.py new file mode 100644 index 0000000..6a09ae5 --- /dev/null +++ b/opendet/metrics/__init__.py @@ -0,0 +1,16 @@ +import copy + +__all__ = ['build_metric'] + +from .det_metric import DetMetric + +support_dict = ['DetMetric'] + + +def build_metric(config): + config = copy.deepcopy(config) + module_name = config.pop('name') + assert module_name in support_dict, Exception( + 'metric only support {}'.format(support_dict)) + module_class = eval(module_name)(**config) + return module_class diff --git a/opendet/metrics/det_metric.py b/opendet/metrics/det_metric.py new file mode 100644 index 0000000..276fb8c --- /dev/null +++ b/opendet/metrics/det_metric.py @@ -0,0 +1,156 @@ +# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. +# +# 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. + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +__all__ = ['DetMetric', 'DetFCEMetric'] + +from .eval_det_iou import DetectionIoUEvaluator + + +class DetMetric(object): + + def __init__(self, main_indicator='hmean', **kwargs): + self.evaluator = DetectionIoUEvaluator() + self.main_indicator = main_indicator + self.reset() + + def __call__(self, preds, batch, **kwargs): + """ + batch: a list produced by dataloaders. + image: np.ndarray of shape (N, C, H, W). + ratio_list: np.ndarray of shape(N,2) + polygons: np.ndarray of shape (N, K, 4, 2), the polygons of objective regions. + ignore_tags: np.ndarray of shape (N, K), indicates whether a region is ignorable or not. + preds: a list of dict produced by post process + points: np.ndarray of shape (N, K, 4, 2), the polygons of objective regions. + """ + gt_polyons_batch = batch[2] + ignore_tags_batch = batch[3] + for pred, gt_polyons, ignore_tags in zip(preds, gt_polyons_batch, + ignore_tags_batch): + # prepare gt + gt_info_list = [{ + 'points': gt_polyon, + 'text': '', + 'ignore': ignore_tag + } for gt_polyon, ignore_tag in zip(gt_polyons, ignore_tags)] + # prepare det + det_info_list = [{ + 'points': det_polyon, + 'text': '' + } for det_polyon in pred['points']] + result = self.evaluator.evaluate_image(gt_info_list, det_info_list) + self.results.append(result) + + def get_metric(self): + """ + return metrics { + 'precision': 0, + 'recall': 0, + 'hmean': 0 + } + """ + + metrics = self.evaluator.combine_results(self.results) + self.reset() + return metrics + + def reset(self): + self.results = [] # clear results + + +class DetFCEMetric(object): + + def __init__(self, main_indicator='hmean', **kwargs): + self.evaluator = DetectionIoUEvaluator() + self.main_indicator = main_indicator + self.reset() + + def __call__(self, preds, batch, **kwargs): + """ + batch: a list produced by dataloaders. + image: np.ndarray of shape (N, C, H, W). + ratio_list: np.ndarray of shape(N,2) + polygons: np.ndarray of shape (N, K, 4, 2), the polygons of objective regions. + ignore_tags: np.ndarray of shape (N, K), indicates whether a region is ignorable or not. + preds: a list of dict produced by post process + points: np.ndarray of shape (N, K, 4, 2), the polygons of objective regions. + """ + gt_polyons_batch = batch[2] + ignore_tags_batch = batch[3] + + for pred, gt_polyons, ignore_tags in zip(preds, gt_polyons_batch, + ignore_tags_batch): + # prepare gt + gt_info_list = [{ + 'points': gt_polyon, + 'text': '', + 'ignore': ignore_tag + } for gt_polyon, ignore_tag in zip(gt_polyons, ignore_tags)] + # prepare det + det_info_list = [{ + 'points': det_polyon, + 'text': '', + 'score': score + } for det_polyon, score in zip(pred['points'], pred['scores'])] + + for score_thr in self.results.keys(): + det_info_list_thr = [ + det_info for det_info in det_info_list + if det_info['score'] >= score_thr + ] + result = self.evaluator.evaluate_image(gt_info_list, + det_info_list_thr) + self.results[score_thr].append(result) + + def get_metric(self): + """ + return metrics {'heman':0, + 'thr 0.3':'precision: 0 recall: 0 hmean: 0', + 'thr 0.4':'precision: 0 recall: 0 hmean: 0', + 'thr 0.5':'precision: 0 recall: 0 hmean: 0', + 'thr 0.6':'precision: 0 recall: 0 hmean: 0', + 'thr 0.7':'precision: 0 recall: 0 hmean: 0', + 'thr 0.8':'precision: 0 recall: 0 hmean: 0', + 'thr 0.9':'precision: 0 recall: 0 hmean: 0', + } + """ + metrics = {} + hmean = 0 + for score_thr in self.results.keys(): + metric = self.evaluator.combine_results(self.results[score_thr]) + # for key, value in metric.items(): + # metrics['{}_{}'.format(key, score_thr)] = value + metric_str = 'precision:{:.5f} recall:{:.5f} hmean:{:.5f}'.format( + metric['precision'], metric['recall'], metric['hmean']) + metrics['thr {}'.format(score_thr)] = metric_str + hmean = max(hmean, metric['hmean']) + metrics['hmean'] = hmean + + self.reset() + return metrics + + def reset(self): + self.results = { + 0.3: [], + 0.4: [], + 0.5: [], + 0.6: [], + 0.7: [], + 0.8: [], + 0.9: [], + } # clear results diff --git a/opendet/metrics/eval_det_iou.py b/opendet/metrics/eval_det_iou.py new file mode 100644 index 0000000..8cfb8c1 --- /dev/null +++ b/opendet/metrics/eval_det_iou.py @@ -0,0 +1,211 @@ +#!/usr/bin/env python +import numpy as np +from shapely.geometry import Polygon +""" +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 + + numGlobalCareGt = 0 + numGlobalCareDet = 0 + + precision = 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 = [] + + evaluationLog = '' + + for n in range(len(gt)): + points = gt[n]['points'] + dontCare = gt[n]['ignore'] + if not Polygon(points).is_valid: + continue + + 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'] + if not Polygon(points).is_valid: + 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 matrixs + 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: + precision = float(0) if numDetCare > 0 else float(1) + else: + precision = 0 if numDetCare == 0 else float( + detMatched) / numDetCare + + matchedSum += detMatched + numGlobalCareGt += numGtCare + numGlobalCareDet += numDetCare + + perSampleMetrics = { + 'gtCare': numGtCare, + 'detCare': numDetCare, + 'detMatched': detMatched, + } + 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)) + 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) diff --git a/opendet/postprocess/db_postprocess.py b/opendet/postprocess/db_postprocess.py index dd6b199..76c9f14 100644 --- a/opendet/postprocess/db_postprocess.py +++ b/opendet/postprocess/db_postprocess.py @@ -1,6 +1,5 @@ import numpy as np import cv2 -import torch from shapely.geometry import Polygon import pyclipper """ @@ -208,16 +207,21 @@ def box_score_slow(self, bitmap, contour): cv2.fillPoly(mask, contour.reshape(1, -1, 2).astype('int32'), 1) return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0] - def __call__(self, outs_dict, shape_list): + def __call__(self, outs_dict, batch, **kwargs): + self.thresh = kwargs.get('mask_thresh', self.thresh) + self.box_thresh = kwargs.get('box_thresh', self.box_thresh) + self.unclip_ratio = kwargs.get('unclip_ratio', self.unclip_ratio) + self.box_type = kwargs.get('box_type', self.box_type) + self.score_mode = kwargs.get('score_mode', self.score_mode) pred = outs_dict['maps'] - if isinstance(pred, torch.Tensor): + if kwargs.get('torch_tensor', True): pred = pred.detach().cpu().numpy() pred = pred[:, 0, :, :] segmentation = pred > self.thresh boxes_batch = [] for batch_index in range(pred.shape[0]): - src_h, src_w, ratio_h, ratio_w = shape_list[batch_index] + src_h, src_w, ratio_h, ratio_w = batch[1][batch_index] if self.dilation_kernel is not None: mask = cv2.dilate( np.array(segmentation[batch_index]).astype(np.uint8), diff --git a/opendet/preprocess/__init__.py b/opendet/preprocess/__init__.py index 851cb65..ac77d4f 100644 --- a/opendet/preprocess/__init__.py +++ b/opendet/preprocess/__init__.py @@ -1,10 +1,19 @@ +import copy import io - import cv2 import numpy as np from PIL import Image +from importlib import import_module -from .db_resize_for_test import DetResizeForTest +MODULE_MAPPING = { + 'DetResizeForTest': '.db_resize_for_test', + 'CopyPaste': '.crop_paste', + 'IaaAugment': '.iaa_augment', + 'EastRandomCropData': '.crop_resize', + 'DetLabelEncode': '.db_label_encode', + 'MakeBorderMap': '.db_label_encode', + 'MakeShrinkMap': '.db_label_encode', +} class NormalizeImage(object): @@ -134,21 +143,28 @@ def __call__(self, data): return data -def create_operators(op_param_list, global_config=None): - """create operators based on the config. +def dynamic_import(class_name): + module_path = MODULE_MAPPING.get(class_name) + if not module_path: + raise ValueError(f'Unsupported class: {class_name}') + + module = import_module(module_path, package=__package__) + return getattr(module, class_name) - Args: - params(list): a dict list, used to create some operators - """ - assert isinstance(op_param_list, list), 'operator config should be a list' + +def create_operators(op_param_list, global_config=None): ops = [] - for operator in op_param_list: - assert isinstance(operator, - dict) and len(operator) == 1, 'yaml format error' - op_name = list(operator)[0] - param = {} if operator[op_name] is None else operator[op_name] - if global_config is not None: + for op_info in op_param_list: + op_name = list(op_info.keys())[0] + param = copy.deepcopy(op_info[op_name]) or {} + + if global_config: param.update(global_config) - op = eval(op_name)(**param) - ops.append(op) + + if op_name in globals(): + op_class = globals()[op_name] + else: + op_class = dynamic_import(op_name) + + ops.append(op_class(**param)) return ops diff --git a/opendet/preprocess/crop_paste.py b/opendet/preprocess/crop_paste.py new file mode 100644 index 0000000..e204a72 --- /dev/null +++ b/opendet/preprocess/crop_paste.py @@ -0,0 +1,178 @@ +# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve. +# +# 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. +import cv2 +import random +import numpy as np +from PIL import Image +from shapely.geometry import Polygon + +from .iaa_augment import IaaAugment +from .crop_resize import is_poly_outside_rect +from tools.infer.utility import get_rotate_crop_image + + +class CopyPaste(object): + + def __init__(self, objects_paste_ratio=0.2, limit_paste=True, **kwargs): + self.ext_data_num = 1 + self.objects_paste_ratio = objects_paste_ratio + self.limit_paste = limit_paste + augmenter_args = [{'type': 'Resize', 'args': {'size': [0.5, 3]}}] + self.aug = IaaAugment(augmenter_args) + + def __call__(self, data): + point_num = data['polys'].shape[1] + src_img = data['image'] + src_polys = data['polys'].tolist() + src_texts = data['texts'] + src_ignores = data['ignore_tags'].tolist() + ext_data = data['ext_data'][0] + ext_image = ext_data['image'] + ext_polys = ext_data['polys'] + ext_texts = ext_data['texts'] + ext_ignores = ext_data['ignore_tags'] + + indexes = [i for i in range(len(ext_ignores)) if not ext_ignores[i]] + select_num = max( + 1, min(int(self.objects_paste_ratio * len(ext_polys)), 30)) + + random.shuffle(indexes) + select_idxs = indexes[:select_num] + select_polys = ext_polys[select_idxs] + select_ignores = ext_ignores[select_idxs] + + src_img = cv2.cvtColor(src_img, cv2.COLOR_BGR2RGB) + ext_image = cv2.cvtColor(ext_image, cv2.COLOR_BGR2RGB) + src_img = Image.fromarray(src_img).convert('RGBA') + for idx, poly, tag in zip(select_idxs, select_polys, select_ignores): + box_img = get_rotate_crop_image(ext_image, poly) + + src_img, box = self.paste_img(src_img, box_img, src_polys) + if box is not None: + box = box.tolist() + for _ in range(len(box), point_num): + box.append(box[-1]) + src_polys.append(box) + src_texts.append(ext_texts[idx]) + src_ignores.append(tag) + src_img = cv2.cvtColor(np.array(src_img), cv2.COLOR_RGB2BGR) + h, w = src_img.shape[:2] + src_polys = np.array(src_polys) + src_polys[:, :, 0] = np.clip(src_polys[:, :, 0], 0, w) + src_polys[:, :, 1] = np.clip(src_polys[:, :, 1], 0, h) + data['image'] = src_img + data['polys'] = src_polys + data['texts'] = src_texts + data['ignore_tags'] = np.array(src_ignores) + return data + + def paste_img(self, src_img, box_img, src_polys): + box_img_pil = Image.fromarray(box_img).convert('RGBA') + src_w, src_h = src_img.size + box_w, box_h = box_img_pil.size + + angle = np.random.randint(0, 360) + box = np.array([[[0, 0], [box_w, 0], [box_w, box_h], [0, box_h]]]) + box = rotate_bbox(box_img, box, angle)[0] + box_img_pil = box_img_pil.rotate(angle, expand=1) + box_w, box_h = box_img_pil.width, box_img_pil.height + if src_w - box_w < 0 or src_h - box_h < 0: + return src_img, None + + paste_x, paste_y = self.select_coord(src_polys, box, src_w - box_w, + src_h - box_h) + if paste_x is None: + return src_img, None + box[:, 0] += paste_x + box[:, 1] += paste_y + r, g, b, A = box_img_pil.split() + src_img.paste(box_img_pil, (paste_x, paste_y), mask=A) + + return src_img, box + + def select_coord(self, src_polys, box, endx, endy): + if self.limit_paste: + xmin, ymin, xmax, ymax = ( + box[:, 0].min(), + box[:, 1].min(), + box[:, 0].max(), + box[:, 1].max(), + ) + for _ in range(50): + paste_x = random.randint(0, endx) + paste_y = random.randint(0, endy) + xmin1 = xmin + paste_x + xmax1 = xmax + paste_x + ymin1 = ymin + paste_y + ymax1 = ymax + paste_y + + num_poly_in_rect = 0 + for poly in src_polys: + if not is_poly_outside_rect(poly, xmin1, ymin1, + xmax1 - xmin1, ymax1 - ymin1): + num_poly_in_rect += 1 + break + if num_poly_in_rect == 0: + return paste_x, paste_y + return None, None + else: + paste_x = random.randint(0, endx) + paste_y = random.randint(0, endy) + return paste_x, paste_y + + +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 rotate_bbox(img, text_polys, angle, scale=1): + """ + from https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/augment.py + Args: + img: np.ndarray + text_polys: np.ndarray N*4*2 + angle: int + scale: int + + Returns: + + """ + w = img.shape[1] + h = img.shape[0] + + rangle = np.deg2rad(angle) + nw = abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w) + nh = abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w) + rot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale) + rot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0])) + rot_mat[0, 2] += rot_move[0] + rot_mat[1, 2] += rot_move[1] + + # ---------------------- rotate box ---------------------- + rot_text_polys = list() + for bbox in text_polys: + point1 = np.dot(rot_mat, np.array([bbox[0, 0], bbox[0, 1], 1])) + point2 = np.dot(rot_mat, np.array([bbox[1, 0], bbox[1, 1], 1])) + point3 = np.dot(rot_mat, np.array([bbox[2, 0], bbox[2, 1], 1])) + point4 = np.dot(rot_mat, np.array([bbox[3, 0], bbox[3, 1], 1])) + rot_text_polys.append([point1, point2, point3, point4]) + return np.array(rot_text_polys, dtype=np.float32) diff --git a/opendet/preprocess/crop_resize.py b/opendet/preprocess/crop_resize.py index fa67a47..4604b4d 100644 --- a/opendet/preprocess/crop_resize.py +++ b/opendet/preprocess/crop_resize.py @@ -1,4 +1,5 @@ import cv2 +import numpy as np def padding_image(img, size=(640, 640)): @@ -37,6 +38,156 @@ def padding_image(img, size=(640, 640)): return padded_img +def is_poly_outside_rect(poly, x, y, w, h): + poly = np.array(poly) + if poly[:, 0].max() < x or poly[:, 0].min() > x + w: + return True + if poly[:, 1].max() < y or poly[:, 1].min() > y + h: + return True + return False + + +def split_regions(axis): + regions = [] + min_axis = 0 + for i in range(1, axis.shape[0]): + if axis[i] != axis[i - 1] + 1: + region = axis[min_axis:i] + min_axis = i + regions.append(region) + return regions + + +def random_select(axis, max_size): + xx = np.random.choice(axis, size=2) + xmin = np.min(xx) + xmax = np.max(xx) + xmin = np.clip(xmin, 0, max_size - 1) + xmax = np.clip(xmax, 0, max_size - 1) + return xmin, xmax + + +def region_wise_random_select(regions, max_size): + selected_index = list(np.random.choice(len(regions), 2)) + selected_values = [] + for index in selected_index: + axis = regions[index] + xx = int(np.random.choice(axis, size=1)) + selected_values.append(xx) + xmin = min(selected_values) + xmax = max(selected_values) + return xmin, xmax + + +def crop_area(im, text_polys, min_crop_side_ratio, max_tries): + h, w, _ = im.shape + h_array = np.zeros(h, dtype=np.int32) + w_array = np.zeros(w, dtype=np.int32) + for points in text_polys: + points = np.round(points, decimals=0).astype(np.int32) + minx = np.min(points[:, 0]) + maxx = np.max(points[:, 0]) + w_array[minx:maxx] = 1 + miny = np.min(points[:, 1]) + maxy = np.max(points[:, 1]) + h_array[miny:maxy] = 1 + # ensure the cropped area not across a text + h_axis = np.where(h_array == 0)[0] + w_axis = np.where(w_array == 0)[0] + + if len(h_axis) == 0 or len(w_axis) == 0: + return 0, 0, w, h + + h_regions = split_regions(h_axis) + w_regions = split_regions(w_axis) + + for i in range(max_tries): + if len(w_regions) > 1: + xmin, xmax = region_wise_random_select(w_regions, w) + else: + xmin, xmax = random_select(w_axis, w) + if len(h_regions) > 1: + ymin, ymax = region_wise_random_select(h_regions, h) + else: + ymin, ymax = random_select(h_axis, h) + + if (xmax - xmin < min_crop_side_ratio * w or ymax - ymin < min_crop_side_ratio * h): + # area too small + continue + num_poly_in_rect = 0 + for poly in text_polys: + if not is_poly_outside_rect(poly, xmin, ymin, xmax - xmin, + ymax - ymin): + num_poly_in_rect += 1 + break + + if num_poly_in_rect > 0: + return xmin, ymin, xmax - xmin, ymax - ymin + + return 0, 0, w, h + + +class EastRandomCropData(object): + + def __init__( + self, + size=(640, 640), + max_tries=10, + min_crop_side_ratio=0.1, + keep_ratio=True, + **kwargs, + ): + self.size = size + self.max_tries = max_tries + self.min_crop_side_ratio = min_crop_side_ratio + self.keep_ratio = keep_ratio + + def __call__(self, data): + img = data['image'] + text_polys = data['polys'] + ignore_tags = data['ignore_tags'] + texts = data['texts'] + all_care_polys = [ + text_polys[i] for i, tag in enumerate(ignore_tags) if not tag + ] + # 计算crop区域 + crop_x, crop_y, crop_w, crop_h = crop_area(img, all_care_polys, + self.min_crop_side_ratio, + self.max_tries) + # crop 图片 保持比例填充 + scale_w = self.size[0] / crop_w + scale_h = self.size[1] / crop_h + scale = min(scale_w, scale_h) + h = int(crop_h * scale) + w = int(crop_w * scale) + if self.keep_ratio: + padimg = np.zeros((self.size[1], self.size[0], img.shape[2]), + img.dtype) + padimg[:h, :w] = cv2.resize( + img[crop_y:crop_y + crop_h, crop_x:crop_x + crop_w], (w, h)) + img = padimg + else: + img = cv2.resize( + img[crop_y:crop_y + crop_h, crop_x:crop_x + crop_w], + tuple(self.size), + ) + # crop 文本框 + text_polys_crop = [] + ignore_tags_crop = [] + texts_crop = [] + for poly, text, tag in zip(text_polys, texts, ignore_tags): + poly = ((poly - (crop_x, crop_y)) * scale).tolist() + if not is_poly_outside_rect(poly, 0, 0, w, h): + text_polys_crop.append(poly) + ignore_tags_crop.append(tag) + texts_crop.append(text) + data['image'] = img + data['polys'] = np.array(text_polys_crop) + data['ignore_tags'] = ignore_tags_crop + data['texts'] = texts_crop + return data + + class CropResize(object): def __init__(self, size=(640, 640), interpolation=cv2.INTER_LINEAR): diff --git a/opendet/preprocess/db_label_encode.py b/opendet/preprocess/db_label_encode.py new file mode 100644 index 0000000..83a6cb2 --- /dev/null +++ b/opendet/preprocess/db_label_encode.py @@ -0,0 +1,313 @@ +import numpy as np +import json +import cv2 + +np.seterr(divide='ignore', invalid='ignore') +import pyclipper +from shapely.geometry import Polygon +import warnings + +warnings.simplefilter('ignore') + + +class DetLabelEncode(object): + + def __init__(self, **kwargs): + pass + + def __call__(self, data): + label = data['label'] + label = json.loads(label) + nBox = len(label) + boxes, txts, txt_tags = [], [], [] + for bno in range(0, nBox): + box = label[bno]['points'] + txt = label[bno]['transcription'] + boxes.append(box) + txts.append(txt) + if txt in ['*', '###']: + txt_tags.append(True) + else: + txt_tags.append(False) + if len(boxes) == 0: + return None + boxes = self.expand_points_num(boxes) + boxes = np.array(boxes, dtype=np.float32) + txt_tags = np.array(txt_tags, dtype=np.bool_) + + data['polys'] = boxes + data['texts'] = txts + data['ignore_tags'] = txt_tags + return data + + def order_points_clockwise(self, pts): + rect = np.zeros((4, 2), dtype='float32') + s = pts.sum(axis=1) + rect[0] = pts[np.argmin(s)] + rect[2] = pts[np.argmax(s)] + tmp = np.delete(pts, (np.argmin(s), np.argmax(s)), axis=0) + diff = np.diff(np.array(tmp), axis=1) + rect[1] = tmp[np.argmin(diff)] + rect[3] = tmp[np.argmax(diff)] + return rect + + def expand_points_num(self, boxes): + max_points_num = 0 + for box in boxes: + if len(box) > max_points_num: + max_points_num = len(box) + ex_boxes = [] + for box in boxes: + ex_box = box + [box[-1]] * (max_points_num - len(box)) + ex_boxes.append(ex_box) + return ex_boxes + + +class MakeBorderMap(object): + + def __init__(self, + shrink_ratio=0.4, + thresh_min=0.3, + thresh_max=0.7, + **kwargs): + self.shrink_ratio = shrink_ratio + self.thresh_min = thresh_min + self.thresh_max = thresh_max + if 'total_epoch' in kwargs and 'epoch' in kwargs and kwargs[ + 'epoch'] != 'None': + self.shrink_ratio = self.shrink_ratio + 0.2 * kwargs[ + 'epoch'] / float(kwargs['total_epoch']) + + def __call__(self, data): + img = data['image'] + text_polys = data['polys'] + ignore_tags = data['ignore_tags'] + + canvas = np.zeros(img.shape[:2], dtype=np.float32) + mask = np.zeros(img.shape[:2], dtype=np.float32) + + for i in range(len(text_polys)): + if ignore_tags[i]: + continue + self.draw_border_map(text_polys[i], canvas, mask=mask) + canvas = canvas * (self.thresh_max - self.thresh_min) + self.thresh_min + + data['threshold_map'] = canvas + data['threshold_mask'] = mask + return data + + def draw_border_map(self, polygon, canvas, mask): + polygon = np.array(polygon) + assert polygon.ndim == 2 + assert polygon.shape[1] == 2 + + polygon_shape = Polygon(polygon) + if polygon_shape.area <= 0: + return + distance = (polygon_shape.area * (1 - np.power(self.shrink_ratio, 2)) / + polygon_shape.length) + subject = [tuple(l) for l in polygon] + padding = pyclipper.PyclipperOffset() + padding.AddPath(subject, pyclipper.JT_ROUND, + pyclipper.ET_CLOSEDPOLYGON) + + padded_polygon = np.array(padding.Execute(distance)[0]) + cv2.fillPoly(mask, [padded_polygon.astype(np.int32)], 1.0) + + xmin = padded_polygon[:, 0].min() + xmax = padded_polygon[:, 0].max() + ymin = padded_polygon[:, 1].min() + ymax = padded_polygon[:, 1].max() + width = xmax - xmin + 1 + height = ymax - ymin + 1 + + polygon[:, 0] = polygon[:, 0] - xmin + polygon[:, 1] = polygon[:, 1] - ymin + + xs = np.broadcast_to( + np.linspace(0, width - 1, num=width).reshape(1, width), + (height, width)) + ys = np.broadcast_to( + np.linspace(0, height - 1, num=height).reshape(height, 1), + (height, width)) + + distance_map = np.zeros((polygon.shape[0], height, width), + dtype=np.float32) + for i in range(polygon.shape[0]): + j = (i + 1) % polygon.shape[0] + absolute_distance = self._distance(xs, ys, polygon[i], polygon[j]) + distance_map[i] = np.clip(absolute_distance / distance, 0, 1) + distance_map = distance_map.min(axis=0) + + xmin_valid = min(max(0, xmin), canvas.shape[1] - 1) + xmax_valid = min(max(0, xmax), canvas.shape[1] - 1) + ymin_valid = min(max(0, ymin), canvas.shape[0] - 1) + ymax_valid = min(max(0, ymax), canvas.shape[0] - 1) + canvas[ymin_valid:ymax_valid + 1, xmin_valid:xmax_valid + 1] = np.fmax( + 1 - distance_map[ymin_valid - ymin:ymax_valid - ymax + height, + xmin_valid - xmin:xmax_valid - xmax + width, ], + canvas[ymin_valid:ymax_valid + 1, xmin_valid:xmax_valid + 1], + ) + + def _distance(self, xs, ys, point_1, point_2): + """ + compute the distance from point to a line + ys: coordinates in the first axis + xs: coordinates in the second axis + point_1, point_2: (x, y), the end of the line + """ + height, width = xs.shape[:2] + square_distance_1 = np.square(xs - point_1[0]) + np.square(ys - + point_1[1]) + square_distance_2 = np.square(xs - point_2[0]) + np.square(ys - + point_2[1]) + square_distance = np.square(point_1[0] - + point_2[0]) + np.square(point_1[1] - + point_2[1]) + + cosin = (square_distance - square_distance_1 - square_distance_2) / ( + 2 * np.sqrt(square_distance_1 * square_distance_2)) + square_sin = 1 - np.square(cosin) + square_sin = np.nan_to_num(square_sin) + result = np.sqrt(square_distance_1 * square_distance_2 * square_sin / + square_distance) + + result[cosin < 0] = np.sqrt( + np.fmin(square_distance_1, square_distance_2))[cosin < 0] + # self.extend_line(point_1, point_2, result) + return result + + def extend_line(self, point_1, point_2, result, shrink_ratio): + ex_point_1 = ( + int( + round(point_1[0] + (point_1[0] - point_2[0]) * + (1 + shrink_ratio))), + int( + round(point_1[1] + (point_1[1] - point_2[1]) * + (1 + shrink_ratio))), + ) + cv2.line( + result, + tuple(ex_point_1), + tuple(point_1), + 4096.0, + 1, + lineType=cv2.LINE_AA, + shift=0, + ) + ex_point_2 = ( + int( + round(point_2[0] + (point_2[0] - point_1[0]) * + (1 + shrink_ratio))), + int( + round(point_2[1] + (point_2[1] - point_1[1]) * + (1 + shrink_ratio))), + ) + cv2.line( + result, + tuple(ex_point_2), + tuple(point_2), + 4096.0, + 1, + lineType=cv2.LINE_AA, + shift=0, + ) + return ex_point_1, ex_point_2 + + +class MakeShrinkMap(object): + r""" + Making binary mask from detection data with ICDAR format. + Typically following the process of class `MakeICDARData`. + """ + + def __init__(self, min_text_size=8, shrink_ratio=0.4, **kwargs): + self.min_text_size = min_text_size + self.shrink_ratio = shrink_ratio + if 'total_epoch' in kwargs and 'epoch' in kwargs and kwargs[ + 'epoch'] != 'None': + self.shrink_ratio = self.shrink_ratio + 0.2 * kwargs[ + 'epoch'] / float(kwargs['total_epoch']) + + def __call__(self, data): + image = data['image'] + text_polys = data['polys'] + ignore_tags = data['ignore_tags'] + + h, w = image.shape[:2] + text_polys, ignore_tags = self.validate_polygons( + text_polys, ignore_tags, h, w) + gt = np.zeros((h, w), dtype=np.float32) + mask = np.ones((h, w), dtype=np.float32) + for i in range(len(text_polys)): + polygon = text_polys[i] + height = max(polygon[:, 1]) - min(polygon[:, 1]) + width = max(polygon[:, 0]) - min(polygon[:, 0]) + if ignore_tags[i] or min(height, width) < self.min_text_size: + cv2.fillPoly(mask, + polygon.astype(np.int32)[np.newaxis, :, :], 0) + ignore_tags[i] = True + else: + polygon_shape = Polygon(polygon) + subject = [tuple(l) for l in polygon] + padding = pyclipper.PyclipperOffset() + padding.AddPath(subject, pyclipper.JT_ROUND, + pyclipper.ET_CLOSEDPOLYGON) + shrunk = [] + + # Increase the shrink ratio every time we get multiple polygon returned back + possible_ratios = np.arange(self.shrink_ratio, 1, + self.shrink_ratio) + np.append(possible_ratios, 1) + # print(possible_ratios) + for ratio in possible_ratios: + # print(f"Change shrink ratio to {ratio}") + distance = (polygon_shape.area * (1 - np.power(ratio, 2)) / + polygon_shape.length) + shrunk = padding.Execute(-distance) + if len(shrunk) == 1: + break + + if shrunk == []: + cv2.fillPoly(mask, + polygon.astype(np.int32)[np.newaxis, :, :], 0) + ignore_tags[i] = True + continue + + for each_shrink in shrunk: + shrink = np.array(each_shrink).reshape(-1, 2) + cv2.fillPoly(gt, [shrink.astype(np.int32)], 1) + + data['shrink_map'] = gt + data['shrink_mask'] = mask + return data + + def validate_polygons(self, polygons, ignore_tags, h, w): + """ + polygons (numpy.array, required): of shape (num_instances, num_points, 2) + """ + if len(polygons) == 0: + return polygons, ignore_tags + assert len(polygons) == len(ignore_tags) + for polygon in polygons: + polygon[:, 0] = np.clip(polygon[:, 0], 0, w - 1) + polygon[:, 1] = np.clip(polygon[:, 1], 0, h - 1) + + for i in range(len(polygons)): + area = self.polygon_area(polygons[i]) + if abs(area) < 1: + ignore_tags[i] = True + if area > 0: + polygons[i] = polygons[i][::-1, :] + return polygons, ignore_tags + + def polygon_area(self, polygon): + """ + compute polygon area + """ + area = 0 + q = polygon[-1] + for p in polygon: + area += p[0] * q[1] - p[1] * q[0] + q = p + return area / 2.0 diff --git a/opendet/preprocess/db_resize_for_test.py b/opendet/preprocess/db_resize_for_test.py index dbd95f0..e974884 100644 --- a/opendet/preprocess/db_resize_for_test.py +++ b/opendet/preprocess/db_resize_for_test.py @@ -27,6 +27,8 @@ def __init__(self, **kwargs): def __call__(self, data): img = data['image'] + if 'max_sile_len' in data: + self.limit_side_len = data['max_sile_len'] src_h, src_w, _ = img.shape if sum([src_h, src_w]) < 64: img = self.image_padding(img) diff --git a/opendet/preprocess/iaa_augment.py b/opendet/preprocess/iaa_augment.py new file mode 100644 index 0000000..1817faa --- /dev/null +++ b/opendet/preprocess/iaa_augment.py @@ -0,0 +1,230 @@ +# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve. +# +# 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. +""" +This code is refer from: +https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/iaa_augment.py +""" +import os + +# Prevent automatic updates in Albumentations for stability in augmentation behavior +os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1' + +import numpy as np +import albumentations as A +from albumentations.core.transforms_interface import DualTransform +from albumentations.augmentations.geometric import functional as fgeometric +from packaging import version + +ALBU_VERSION = version.parse(A.__version__) +IS_ALBU_NEW_VERSION = ALBU_VERSION >= version.parse('1.4.15') + + +# Custom resize transformation mimicking Imgaug's behavior with scaling +class ImgaugLikeResize(DualTransform): + + def __init__(self, scale_range=(0.5, 3.0), interpolation=1, p=1.0): + super(ImgaugLikeResize, self).__init__(p) + self.scale_range = scale_range + self.interpolation = interpolation + + # Resize the image based on a randomly chosen scale within the scale range + def apply(self, img, scale=1.0, **params): + height, width = img.shape[:2] + new_height = int(height * scale) + new_width = int(width * scale) + + if IS_ALBU_NEW_VERSION: + return fgeometric.resize(img, (new_height, new_width), + interpolation=self.interpolation) + return fgeometric.resize(img, + new_height, + new_width, + interpolation=self.interpolation) + + # Apply the same scaling transformation to keypoints (e.g., polygon points) + def apply_to_keypoints(self, keypoints, scale=1.0, **params): + return np.array([(x * scale, y * scale) + tuple(rest) + for x, y, *rest in keypoints]) + + # Get random scale parameter within the specified range + def get_params(self): + scale = np.random.uniform(self.scale_range[0], self.scale_range[1]) + return {'scale': scale} + + +# Builder class to translate custom augmenter arguments into Albumentations-compatible format +class AugmenterBuilder(object): + + def __init__(self): + # Map common Imgaug transformations to equivalent Albumentations transforms + self.imgaug_to_albu = { + 'Fliplr': 'HorizontalFlip', + 'Flipud': 'VerticalFlip', + 'Affine': 'Affine', + # Additional mappings can be added here if needed + } + + # Recursive method to construct augmentation pipeline based on provided arguments + def build(self, args, root=True): + if args is None or len(args) == 0: + return None + elif isinstance(args, list): + # Build the full augmentation sequence if it's a root-level call + if root: + sequence = [self.build(value, root=False) for value in args] + return A.Compose( + sequence, + keypoint_params=A.KeypointParams(format='xy', + remove_invisible=False), + ) + else: + # Build individual augmenters for nested arguments + augmenter_type = args[0] + augmenter_args = args[1] if len(args) > 1 else {} + augmenter_args_mapped = self.map_arguments( + augmenter_type, augmenter_args) + augmenter_type_mapped = self.imgaug_to_albu.get( + augmenter_type, augmenter_type) + if augmenter_type_mapped == 'Resize': + return ImgaugLikeResize(**augmenter_args_mapped) + else: + cls = getattr(A, augmenter_type_mapped) + return cls( + **{ + k: self.to_tuple_if_list(v) + for k, v in augmenter_args_mapped.items() + }) + elif isinstance(args, dict): + # Process individual transformation specified as dictionary + augmenter_type = args['type'] + augmenter_args = args.get('args', {}) + augmenter_args_mapped = self.map_arguments(augmenter_type, + augmenter_args) + augmenter_type_mapped = self.imgaug_to_albu.get( + augmenter_type, augmenter_type) + if augmenter_type_mapped == 'Resize': + return ImgaugLikeResize(**augmenter_args_mapped) + else: + cls = getattr(A, augmenter_type_mapped) + return cls( + **{ + k: self.to_tuple_if_list(v) + for k, v in augmenter_args_mapped.items() + }) + else: + raise RuntimeError('Unknown augmenter arg: ' + str(args)) + + # Map arguments to expected format for each augmenter type + def map_arguments(self, augmenter_type, augmenter_args): + augmenter_args = augmenter_args.copy( + ) # Avoid modifying the original arguments + if augmenter_type == 'Resize': + # Ensure size is a valid 2-element list or tuple + size = augmenter_args.get('size') + if size: + if not isinstance(size, (list, tuple)) or len(size) != 2: + raise ValueError( + f"'size' must be a list or tuple of two numbers, but got {size}" + ) + min_scale, max_scale = size + return { + 'scale_range': (min_scale, max_scale), + 'interpolation': 1, # Linear interpolation + 'p': 1.0, + } + else: + return { + 'scale_range': (1.0, 1.0), + 'interpolation': 1, + 'p': 1.0 + } + elif augmenter_type == 'Affine': + # Map rotation to a tuple and ensure p=1.0 to apply transformation + rotate = augmenter_args.get('rotate', 0) + if isinstance(rotate, list): + rotate = tuple(rotate) + elif isinstance(rotate, (int, float)): + rotate = (float(rotate), float(rotate)) + augmenter_args['rotate'] = rotate + augmenter_args['p'] = 1.0 + return augmenter_args + else: + # For other augmenters, ensure 'p' probability is specified + p = augmenter_args.get('p', 1.0) + augmenter_args['p'] = p + return augmenter_args + + # Convert lists to tuples for Albumentations compatibility + def to_tuple_if_list(self, obj): + if isinstance(obj, list): + return tuple(obj) + return obj + + +# Wrapper class for image and polygon transformations using Imgaug-style augmentation +class IaaAugment: + + def __init__(self, augmenter_args=None, **kwargs): + if augmenter_args is None: + # Default augmenters if none are specified + augmenter_args = [ + { + 'type': 'Fliplr', + 'args': { + 'p': 0.5 + } + }, + { + 'type': 'Affine', + 'args': { + 'rotate': [-10, 10] + } + }, + { + 'type': 'Resize', + 'args': { + 'size': [0.5, 3] + } + }, + ] + self.augmenter = AugmenterBuilder().build(augmenter_args) + + # Apply the augmentations to image and polygon data + def __call__(self, data): + image = data['image'] + + if self.augmenter: + # Flatten polygons to individual keypoints for transformation + keypoints = [] + keypoints_lengths = [] + for poly in data['polys']: + keypoints.extend([tuple(point) for point in poly]) + keypoints_lengths.append(len(poly)) + + # Apply the augmentation pipeline to image and keypoints + transformed = self.augmenter(image=image, keypoints=keypoints) + data['image'] = transformed['image'] + + # Extract transformed keypoints and reconstruct polygon structures + transformed_keypoints = transformed['keypoints'] + + # Reassemble polygons from transformed keypoints + new_polys = [] + idx = 0 + for length in keypoints_lengths: + new_poly = transformed_keypoints[idx:idx + length] + new_polys.append(np.array([kp[:2] for kp in new_poly])) + idx += length + data['polys'] = np.array(new_polys) + return data diff --git a/openrec/losses/__init__.py b/openrec/losses/__init__.py index 7af475d..25fc17d 100644 --- a/openrec/losses/__init__.py +++ b/openrec/losses/__init__.py @@ -1,40 +1,43 @@ import copy - +from importlib import import_module from torch import nn -from .abinet_loss import ABINetLoss -from .ar_loss import ARLoss -from .cdistnet_loss import CDistNetLoss -from .ce_loss import CELoss -from .cppd_loss import CPPDLoss -from .ctc_loss import CTCLoss -from .igtr_loss import IGTRLoss -from .lister_loss import LISTERLoss -from .lpv_loss import LPVLoss -from .mgp_loss import MGPLoss -from .parseq_loss import PARSeqLoss -from .robustscanner_loss import RobustScannerLoss -from .smtr_loss import SMTRLoss -from .srn_loss import SRNLoss -from .visionlan_loss import VisionLANLoss -from .cam_loss import CAMLoss -from .seed_loss import SEEDLoss - -support_dict = [ - 'CTCLoss', 'ARLoss', 'CELoss', 'CPPDLoss', 'ABINetLoss', 'CDistNetLoss', - 'VisionLANLoss', 'PARSeqLoss', 'IGTRLoss', 'SMTRLoss', 'LPVLoss', - 'RobustScannerLoss', 'SRNLoss', 'LISTERLoss', 'GTCLoss', 'MGPLoss', - 'CAMLoss', 'SEEDLoss' -] +name_to_module = { + 'ABINetLoss': '.abinet_loss', + 'ARLoss': '.ar_loss', + 'CDistNetLoss': '.cdistnet_loss', + 'CELoss': '.ce_loss', + 'CPPDLoss': '.cppd_loss', + 'CTCLoss': '.ctc_loss', + 'IGTRLoss': '.igtr_loss', + 'LISTERLoss': '.lister_loss', + 'LPVLoss': '.lpv_loss', + 'MGPLoss': '.mgp_loss', + 'PARSeqLoss': '.parseq_loss', + 'RobustScannerLoss': '.robustscanner_loss', + 'SEEDLoss': '.seed_loss', + 'SMTRLoss': '.smtr_loss', + 'SRNLoss': '.srn_loss', + 'VisionLANLoss': '.visionlan_loss', + 'CAMLoss': '.cam_loss', +} def build_loss(config): config = copy.deepcopy(config) module_name = config.pop('name') - assert module_name in support_dict, Exception( - 'loss only support {}'.format(support_dict)) - module_class = eval(module_name)(**config) - return module_class + + if module_name in globals(): + module_class = globals()[module_name] + else: + assert module_name in name_to_module, Exception( + '{} is not supported. The losses in {} are supportes'.format( + module_name, list(name_to_module.keys()))) + module_path = name_to_module[module_name] + module = import_module(module_path, package=__package__) + module_class = getattr(module, module_name) + + return module_class(**config) class GTCLoss(nn.Module): @@ -46,7 +49,10 @@ def __init__(self, zero_infinity=True, **kwargs): super(GTCLoss, self).__init__() - self.ctc_loss = CTCLoss(zero_infinity=zero_infinity) + # 动态构建CTCLoss + ctc_config = {'name': 'CTCLoss', 'zero_infinity': zero_infinity} + self.ctc_loss = build_loss(ctc_config) + # 构建GTC损失 self.gtc_loss = build_loss(gtc_loss) self.gtc_weight = gtc_weight self.ctc_weight = ctc_weight diff --git a/openrec/modeling/decoders/__init__.py b/openrec/modeling/decoders/__init__.py index f570fd0..91dc019 100644 --- a/openrec/modeling/decoders/__init__.py +++ b/openrec/modeling/decoders/__init__.py @@ -1,48 +1,52 @@ import torch.nn as nn +from importlib import import_module __all__ = ['build_decoder'] +class_to_module = { + 'ABINetDecoder': '.abinet_decoder', + 'ASTERDecoder': '.aster_decoder', + 'CDistNetDecoder': '.cdistnet_decoder', + 'CPPDDecoder': '.cppd_decoder', + 'RCTCDecoder': '.rctc_decoder', + 'CTCDecoder': '.ctc_decoder', + 'DANDecoder': '.dan_decoder', + 'IGTRDecoder': '.igtr_decoder', + 'LISTERDecoder': '.lister_decoder', + 'LPVDecoder': '.lpv_decoder', + 'MGPDecoder': '.mgp_decoder', + 'NRTRDecoder': '.nrtr_decoder', + 'PARSeqDecoder': '.parseq_decoder', + 'RobustScannerDecoder': '.robustscanner_decoder', + 'SARDecoder': '.sar_decoder', + 'SMTRDecoder': '.smtr_decoder', + 'SMTRDecoderNumAttn': '.smtr_decoder_nattn', + 'SRNDecoder': '.srn_decoder', + 'VisionLANDecoder': '.visionlan_decoder', + 'MATRNDecoder': '.matrn_decoder', + 'CAMDecoder': '.cam_decoder', + 'OTEDecoder': '.ote_decoder', + 'BUSDecoder': '.bus_decoder', + 'DptrParseq': '.dptr_parseq_clip_b_decoder', +} -def build_decoder(config): - # rec decoder - from .abinet_decoder import ABINetDecoder - from .aster_decoder import ASTERDecoder - from .cdistnet_decoder import CDistNetDecoder - from .cppd_decoder import CPPDDecoder - from .rctc_decoder import RCTCDecoder - from .ctc_decoder import CTCDecoder - from .dan_decoder import DANDecoder - from .igtr_decoder import IGTRDecoder - from .lister_decoder import LISTERDecoder - from .lpv_decoder import LPVDecoder - from .mgp_decoder import MGPDecoder - from .nrtr_decoder import NRTRDecoder - from .parseq_decoder import PARSeqDecoder - from .robustscanner_decoder import RobustScannerDecoder - from .sar_decoder import SARDecoder - from .smtr_decoder import SMTRDecoder - from .smtr_decoder_nattn import SMTRDecoderNumAttn - from .srn_decoder import SRNDecoder - from .visionlan_decoder import VisionLANDecoder - from .matrn_decoder import MATRNDecoder - from .cam_decoder import CAMDecoder - from .ote_decoder import OTEDecoder - from .bus_decoder import BUSDecoder - support_dict = [ - 'CTCDecoder', 'NRTRDecoder', 'CPPDDecoder', 'ABINetDecoder', - 'CDistNetDecoder', 'VisionLANDecoder', 'PARSeqDecoder', 'IGTRDecoder', - 'SMTRDecoder', 'LPVDecoder', 'SARDecoder', 'RobustScannerDecoder', - 'SRNDecoder', 'ASTERDecoder', 'RCTCDecoder', 'LISTERDecoder', - 'GTCDecoder', 'SMTRDecoderNumAttn', 'MATRNDecoder', 'MGPDecoder', - 'DANDecoder', 'CAMDecoder', 'OTEDecoder', 'BUSDecoder' - ] +def build_decoder(config): module_name = config.pop('name') - assert module_name in support_dict, Exception( - 'decoder only support {}'.format(support_dict)) - module_class = eval(module_name)(**config) - return module_class + + # Check if the class is defined in current module (e.g., GTCDecoder) + if module_name in globals(): + module_class = globals()[module_name] + else: + if module_name not in class_to_module: + raise ValueError(f'Unsupported decoder: {module_name}') + module_str = class_to_module[module_name] + # Dynamically import the module and get the class + module = import_module(module_str, package=__package__) + module_class = getattr(module, module_name) + + return module_class(**config) class GTCDecoder(nn.Module): diff --git a/openrec/modeling/decoders/ctc_decoder.py b/openrec/modeling/decoders/ctc_decoder.py index aca917b..aac753a 100644 --- a/openrec/modeling/decoders/ctc_decoder.py +++ b/openrec/modeling/decoders/ctc_decoder.py @@ -132,7 +132,7 @@ def forward(self, x): z = self.conv2(z) # SVTR global block B, C, H, W = z.shape - z = z.flatten(2).transpose(1, 2) + z = z.flatten(2).transpose(1, 2).contiguous() for blk in self.svtr_block: z = blk(z) z = self.norm(z) diff --git a/openrec/modeling/decoders/dptr_parseq_clip_b_decoder.py b/openrec/modeling/decoders/dptr_parseq_clip_b_decoder.py new file mode 100644 index 0000000..566b4a2 --- /dev/null +++ b/openrec/modeling/decoders/dptr_parseq_clip_b_decoder.py @@ -0,0 +1,1398 @@ +# Scene Text Recognition Model Hub +# Copyright 2022 Darwin Bautista +# +# 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 +# +# https://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. + +import math +from itertools import permutations +from collections import OrderedDict +import hashlib +import os +import gzip +import html +import urllib +import warnings +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.nn.modules import transformer +from typing import Any, Optional, Tuple, List, Union +from pkg_resources import packaging +from PIL import Image +from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize +from tqdm import tqdm +from functools import lru_cache + +import ftfy +import regex as re + +try: + from torchvision.transforms import InterpolationMode + BICUBIC = InterpolationMode.BICUBIC +except ImportError: + BICUBIC = Image.BICUBIC + + +@lru_cache() +def default_bpe(): + return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz") + + +@lru_cache() +def bytes_to_unicode(): + """ + Returns list of utf-8 byte and a corresponding list of unicode strings. + The reversible bpe codes work on unicode strings. + This means you need a large # of unicode characters in your vocab if you want to avoid UNKs. + When you're at something like a 10B token dataset you end up needing around 5K for decent coverage. + This is a signficant percentage of your normal, say, 32K bpe vocab. + To avoid that, we want lookup tables between utf-8 bytes and unicode strings. + And avoids mapping to whitespace/control characters the bpe code barfs on. + """ + bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1)) + cs = bs[:] + n = 0 + for b in range(2**8): + if b not in bs: + bs.append(b) + cs.append(2**8+n) + n += 1 + cs = [chr(n) for n in cs] + return dict(zip(bs, cs)) + + +def get_pairs(word): + """Return set of symbol pairs in a word. + Word is represented as tuple of symbols (symbols being variable-length strings). + """ + pairs = set() + prev_char = word[0] + for char in word[1:]: + pairs.add((prev_char, char)) + prev_char = char + return pairs + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r'\s+', ' ', text) + text = text.strip() + return text + + +class SimpleTokenizer(object): + def __init__(self, bpe_path: str = default_bpe()): + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + merges = gzip.open(bpe_path).read().decode("utf-8").split('\n') + merges = merges[1:49152-256-2+1] + merges = [tuple(merge.split()) for merge in merges] + vocab = list(bytes_to_unicode().values()) + vocab = vocab + [v+'' for v in vocab] + for merge in merges: + vocab.append(''.join(merge)) + vocab.extend(['<|startoftext|>', '<|endoftext|>']) + self.encoder = dict(zip(vocab, range(len(vocab)))) + self.decoder = {v: k for k, v in self.encoder.items()} + self.bpe_ranks = dict(zip(merges, range(len(merges)))) + self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'} + self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE) + + def bpe(self, token): + if token in self.cache: + return self.cache[token] + word = tuple(token[:-1]) + ( token[-1] + '',) + pairs = get_pairs(word) + + if not pairs: + return token+'' + + while True: + bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf'))) + if bigram not in self.bpe_ranks: + break + first, second = bigram + new_word = [] + i = 0 + while i < len(word): + try: + j = word.index(first, i) + new_word.extend(word[i:j]) + i = j + except: + new_word.extend(word[i:]) + break + + if word[i] == first and i < len(word)-1 and word[i+1] == second: + new_word.append(first+second) + i += 2 + else: + new_word.append(word[i]) + i += 1 + new_word = tuple(new_word) + word = new_word + if len(word) == 1: + break + else: + pairs = get_pairs(word) + word = ' '.join(word) + self.cache[token] = word + return word + + def encode(self, text): + bpe_tokens = [] + text = whitespace_clean(basic_clean(text)).lower() + for token in re.findall(self.pat, text): + token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8')) + bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' ')) + return bpe_tokens + + def decode(self, tokens): + text = ''.join([self.decoder[token] for token in tokens]) + text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('', ' ') + return text + + +if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"): + warnings.warn("PyTorch version 1.7.1 or higher is recommended") + + +__all__ = ["available_models", "load", "tokenize"] +_tokenizer = SimpleTokenizer() + +_MODELS = { + "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt", + "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt", + "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt", + "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt", + "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt", + "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt", + "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt", + "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt", + "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt", +} + + +def convert_weights(model: nn.Module): + """Convert applicable model parameters to fp16""" + + def _convert_weights_to_fp16(l): + if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)): + l.weight.data = l.weight.data.half() + if l.bias is not None: + l.bias.data = l.bias.data.half() + + if isinstance(l, nn.MultiheadAttention): + for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]: + tensor = getattr(l, attr) + if tensor is not None: + tensor.data = tensor.data.half() + + for name in ["text_projection", "proj"]: + if hasattr(l, name): + attr = getattr(l, name) + if attr is not None: + attr.data = attr.data.half() + + model.apply(_convert_weights_to_fp16) + + +def build_model(state_dict: dict): + vit = "visual.proj" in state_dict + + if vit: + vision_width = state_dict["visual.conv1.weight"].shape[0] + vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")]) + vision_patch_size = state_dict["visual.conv1.weight"].shape[-1] + grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5) + image_resolution = vision_patch_size * grid_size + else: + counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]] + vision_layers = tuple(counts) + vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0] + output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5) + vision_patch_size = None + assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0] + image_resolution = output_width * 32 + + embed_dim = state_dict["text_projection"].shape[1] + context_length = state_dict["positional_embedding"].shape[0] + vocab_size = state_dict["token_embedding.weight"].shape[0] + transformer_width = state_dict["ln_final.weight"].shape[0] + transformer_heads = transformer_width // 64 + transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks"))) + + model = CLIP( + embed_dim, + image_resolution, vision_layers, vision_width, vision_patch_size, + context_length, vocab_size, transformer_width, transformer_heads, transformer_layers + ) + + for key in ["input_resolution", "context_length", "vocab_size"]: + if key in state_dict: + del state_dict[key] + + convert_weights(model) + model.load_state_dict(state_dict) + return model.eval() + + +def _download(url: str, root: str): + os.makedirs(root, exist_ok=True) + filename = os.path.basename(url) + + expected_sha256 = url.split("/")[-2] + download_target = os.path.join(root, filename) + + if os.path.exists(download_target) and not os.path.isfile(download_target): + raise RuntimeError(f"{download_target} exists and is not a regular file") + + if os.path.isfile(download_target): + if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256: + return download_target + else: + warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file") + + with urllib.request.urlopen(url) as source, open(download_target, "wb") as output: + with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop: + while True: + buffer = source.read(8192) + if not buffer: + break + + output.write(buffer) + loop.update(len(buffer)) + + if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256: + raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match") + + return download_target + + +def _convert_image_to_rgb(image): + return image.convert("RGB") + + +def _transform(n_px): + return Compose([ + Resize(n_px, interpolation=BICUBIC), + CenterCrop(n_px), + _convert_image_to_rgb, + ToTensor(), + Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)), + ]) + + +def available_models() -> List[str]: + """Returns the names of available CLIP models""" + return list(_MODELS.keys()) + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1): + super().__init__() + + # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1 + self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.relu1 = nn.ReLU(inplace=True) + + self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + self.relu2 = nn.ReLU(inplace=True) + + self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity() + + self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.relu3 = nn.ReLU(inplace=True) + + self.downsample = None + self.stride = stride + + if stride > 1 or inplanes != planes * Bottleneck.expansion: + # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1 + self.downsample = nn.Sequential(OrderedDict([ + ("-1", nn.AvgPool2d(stride)), + ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)), + ("1", nn.BatchNorm2d(planes * self.expansion)) + ])) + + def forward(self, x: torch.Tensor): + identity = x + + out = self.relu1(self.bn1(self.conv1(x))) + out = self.relu2(self.bn2(self.conv2(out))) + out = self.avgpool(out) + out = self.bn3(self.conv3(out)) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.relu3(out) + return out + + +class AttentionPool2d(nn.Module): + def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None): + super().__init__() + self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5) + self.k_proj = nn.Linear(embed_dim, embed_dim) + self.q_proj = nn.Linear(embed_dim, embed_dim) + self.v_proj = nn.Linear(embed_dim, embed_dim) + self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim) + self.num_heads = num_heads + + def forward(self, x): + x = x.flatten(start_dim=2).permute(2, 0, 1) # NCHW -> (HW)NC + x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC + x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC + x, _ = F.multi_head_attention_forward( + query=x[:1], key=x, value=x, + embed_dim_to_check=x.shape[-1], + num_heads=self.num_heads, + q_proj_weight=self.q_proj.weight, + k_proj_weight=self.k_proj.weight, + v_proj_weight=self.v_proj.weight, + in_proj_weight=None, + in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]), + bias_k=None, + bias_v=None, + add_zero_attn=False, + dropout_p=0, + out_proj_weight=self.c_proj.weight, + out_proj_bias=self.c_proj.bias, + use_separate_proj_weight=True, + training=self.training, + need_weights=False + ) + return x.squeeze(0) + + +class ModifiedResNet(nn.Module): + """ + A ResNet class that is similar to torchvision's but contains the following changes: + - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool. + - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1 + - The final pooling layer is a QKV attention instead of an average pool + """ + + def __init__(self, layers, output_dim, heads, input_resolution=224, width=64): + super().__init__() + self.output_dim = output_dim + self.input_resolution = input_resolution + + # the 3-layer stem + self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(width // 2) + self.relu1 = nn.ReLU(inplace=True) + self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(width // 2) + self.relu2 = nn.ReLU(inplace=True) + self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False) + self.bn3 = nn.BatchNorm2d(width) + self.relu3 = nn.ReLU(inplace=True) + self.avgpool = nn.AvgPool2d(2) + + # residual layers + self._inplanes = width # this is a *mutable* variable used during construction + self.layer1 = self._make_layer(width, layers[0]) + self.layer2 = self._make_layer(width * 2, layers[1], stride=2) + self.layer3 = self._make_layer(width * 4, layers[2], stride=2) + self.layer4 = self._make_layer(width * 8, layers[3], stride=2) + + embed_dim = width * 32 # the ResNet feature dimension + self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim) + + def _make_layer(self, planes, blocks, stride=1): + layers = [Bottleneck(self._inplanes, planes, stride)] + + self._inplanes = planes * Bottleneck.expansion + for _ in range(1, blocks): + layers.append(Bottleneck(self._inplanes, planes)) + + return nn.Sequential(*layers) + + def forward(self, x): + def stem(x): + x = self.relu1(self.bn1(self.conv1(x))) + x = self.relu2(self.bn2(self.conv2(x))) + x = self.relu3(self.bn3(self.conv3(x))) + x = self.avgpool(x) + return x + + x = x.type(self.conv1.weight.dtype) + x = stem(x) + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + x = self.attnpool(x) + + return x + + +class LayerNorm(nn.LayerNorm): + """Subclass torch's LayerNorm to handle fp16.""" + + def forward(self, x: torch.Tensor): + orig_type = x.dtype + ret = super().forward(x.type(torch.float32)) + return ret.type(orig_type) + + +class QuickGELU(nn.Module): + def forward(self, x: torch.Tensor): + return x * torch.sigmoid(1.702 * x) + + +class ResidualAttentionBlock(nn.Module): + def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None): + super().__init__() + + self.attn = nn.MultiheadAttention(d_model, n_head) + self.ln_1 = LayerNorm(d_model) + self.mlp = nn.Sequential(OrderedDict([ + ("c_fc", nn.Linear(d_model, d_model * 4)), + ("gelu", QuickGELU()), + ("c_proj", nn.Linear(d_model * 4, d_model)) + ])) + self.ln_2 = LayerNorm(d_model) + self.attn_mask = attn_mask + + def attention(self, x: torch.Tensor): + self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None + return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0] + + def forward(self, x: torch.Tensor): + x = x + self.attention(self.ln_1(x)) + x = x + self.mlp(self.ln_2(x)) + return x + + +class Transformer(nn.Module): + def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None): + super().__init__() + self.width = width + self.layers = layers + self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)]) + + def forward(self, x: torch.Tensor): + return self.resblocks(x) + + +class VisionTransformer(nn.Module): + def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int): + super().__init__() + self.input_resolution = input_resolution + self.output_dim = output_dim + self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False) + + scale = width ** -0.5 + self.class_embedding = nn.Parameter(scale * torch.randn(width)) + self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width)) + self.ln_pre = LayerNorm(width) + + self.transformer = Transformer(width, layers, heads) + + self.ln_post = LayerNorm(width) + self.proj = nn.Parameter(scale * torch.randn(width, output_dim)) + + def forward(self, x: torch.Tensor): + x = self.conv1(x) # shape = [*, width, grid, grid] + x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [*, width, grid ** 2] + x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width] + x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1) # shape = [*, grid ** 2 + 1, width] + x = x + self.positional_embedding.to(x.dtype) + x = self.ln_pre(x) + + x = x.permute(1, 0, 2) # NLD -> LND + x = self.transformer(x) + x = x.permute(1, 0, 2) # LND -> NLD + + x = self.ln_post(x) + if self.proj is not None: + x = x @ self.proj + + return x + + +class CLIP(nn.Module): + def __init__(self, + embed_dim: int, + # vision + image_resolution: int, + vision_layers: Union[Tuple[int, int, int, int], int], + vision_width: int, + vision_patch_size: int, + # text + context_length: int, + vocab_size: int, + transformer_width: int, + transformer_heads: int, + transformer_layers: int + ): + super().__init__() + + self.context_length = context_length + + if isinstance(vision_layers, (tuple, list)): + vision_heads = vision_width * 32 // 64 + self.visual = ModifiedResNet( + layers=vision_layers, + output_dim=embed_dim, + heads=vision_heads, + input_resolution=image_resolution, + width=vision_width + ) + else: + vision_heads = vision_width // 64 + self.visual = VisionTransformer( + input_resolution=image_resolution, + patch_size=vision_patch_size, + width=vision_width, + layers=vision_layers, + heads=vision_heads, + output_dim=embed_dim + ) + + self.transformer = Transformer( + width=transformer_width, + layers=transformer_layers, + heads=transformer_heads, + attn_mask=self.build_attention_mask() + ) + + self.vocab_size = vocab_size + self.token_embedding = nn.Embedding(vocab_size, transformer_width) + self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width)) + self.ln_final = LayerNorm(transformer_width) + + self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim)) + self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) + + self.initialize_parameters() + + def initialize_parameters(self): + nn.init.normal_(self.token_embedding.weight, std=0.02) + nn.init.normal_(self.positional_embedding, std=0.01) + + if isinstance(self.visual, ModifiedResNet): + if self.visual.attnpool is not None: + std = self.visual.attnpool.c_proj.in_features ** -0.5 + nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std) + nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std) + nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std) + nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std) + + for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]: + for name, param in resnet_block.named_parameters(): + if name.endswith("bn3.weight"): + nn.init.zeros_(param) + + proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5) + attn_std = self.transformer.width ** -0.5 + fc_std = (2 * self.transformer.width) ** -0.5 + for block in self.transformer.resblocks: + nn.init.normal_(block.attn.in_proj_weight, std=attn_std) + nn.init.normal_(block.attn.out_proj.weight, std=proj_std) + nn.init.normal_(block.mlp.c_fc.weight, std=fc_std) + nn.init.normal_(block.mlp.c_proj.weight, std=proj_std) + + if self.text_projection is not None: + nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5) + + def build_attention_mask(self): + # lazily create causal attention mask, with full attention between the vision tokens + # pytorch uses additive attention mask; fill with -inf + mask = torch.empty(self.context_length, self.context_length) + mask.fill_(float("-inf")) + mask.triu_(1) # zero out the lower diagonal + return mask + + @property + def dtype(self): + return self.visual.conv1.weight.dtype + + def encode_image(self, image): + return self.visual(image.type(self.dtype)) + + def encode_text(self, text): + x = self.token_embedding(text).type(self.dtype) # [batch_size, n_ctx, d_model] + + x = x + self.positional_embedding.type(self.dtype) + x = x.permute(1, 0, 2) # NLD -> LND + x = self.transformer(x) + x = x.permute(1, 0, 2) # LND -> NLD + x = self.ln_final(x).type(self.dtype) + + # take features from the eot embedding (eot_token is the highest number in each sequence) + output = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection + output = torch.cat([output.unsqueeze(1), x], dim=1) + + return output + + def forward(self, image, text): + image_features = self.encode_image(image) + text_features = self.encode_text(text) + + # normalized features + image_features = image_features / image_features.norm(dim=1, keepdim=True) + text_features = text_features / text_features.norm(dim=1, keepdim=True) + + # cosine similarity as logits + logit_scale = self.logit_scale.exp() + logits_per_image = logit_scale * image_features @ text_features.t() + logits_per_text = logits_per_image.t() + + # shape = [global_batch_size, global_batch_size] + return logits_per_image, logits_per_text + + +class FMU(nn.Module): + """A Transformer decoder layer supporting two-stream attention (XLNet) + This implements a pre-LN decoder, as opposed to the post-LN default in PyTorch.""" + + def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='gelu', + layer_norm_eps=1e-5): + super().__init__() + self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm = nn.LayerNorm(d_model, eps=layer_norm_eps) + + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.dropout3 = nn.Dropout(dropout) + + self.activation = transformer._get_activation_fn(activation) + + def __setstate__(self, state): + if 'activation' not in state: + state['activation'] = F.gelu + super().__setstate__(state) + + def forward(self, query: Tensor, memory: Tensor): + """Forward pass for a single stream (i.e. content or query) + tgt_norm is just a LayerNorm'd tgt. Added as a separate parameter for efficiency. + Both tgt_kv and memory are expected to be LayerNorm'd too. + memory is LayerNorm'd by ViT. + """ + query1, ca_weights = self.cross_attn(query, memory, memory) + query = query + self.dropout1(query1) + + query2 = self.linear2(self.dropout2(self.activation(self.linear1(self.norm(query))))) + query = query + self.dropout3(query2) + + return query + + +class DecoderLayer(nn.Module): + """A Transformer decoder layer supporting two-stream attention (XLNet) This + implements a pre-LN decoder, as opposed to the post-LN default in + PyTorch.""" + + def __init__( + self, + d_model, + nhead, + dim_feedforward=2048, + dropout=0.1, + activation='gelu', + layer_norm_eps=1e-5, + ): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, + nhead, + dropout=dropout, + batch_first=True) + self.cross_attn = nn.MultiheadAttention(d_model, + nhead, + dropout=dropout, + batch_first=True) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps) + self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps) + self.norm_q = nn.LayerNorm(d_model, eps=layer_norm_eps) + self.norm_c = nn.LayerNorm(d_model, eps=layer_norm_eps) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.dropout3 = nn.Dropout(dropout) + + self.activation = transformer._get_activation_fn(activation) + + def __setstate__(self, state): + if 'activation' not in state: + state['activation'] = F.gelu + super().__setstate__(state) + + def forward_stream( + self, + tgt: Tensor, + tgt_norm: Tensor, + tgt_kv: Tensor, + memory: Tensor, + tgt_mask: Optional[Tensor], + tgt_key_padding_mask: Optional[Tensor], + ): + """Forward pass for a single stream (i.e. content or query) tgt_norm is + just a LayerNorm'd tgt. + + Added as a separate parameter for efficiency. Both tgt_kv and memory + are expected to be LayerNorm'd too. memory is LayerNorm'd by ViT. + """ + tgt2, sa_weights = self.self_attn( + tgt_norm, + tgt_kv, + tgt_kv, + attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask) + + tgt = tgt + self.dropout1(tgt2) + + tgt2, ca_weights = self.cross_attn(self.norm1(tgt), memory, memory) + self.attn_map = ca_weights + tgt = tgt + self.dropout2(tgt2) + + tgt2 = self.linear2( + self.dropout(self.activation(self.linear1(self.norm2(tgt))))) + tgt = tgt + self.dropout3(tgt2) + return tgt, sa_weights, ca_weights + + def forward( + self, + query, + content, + memory, + query_mask: Optional[Tensor] = None, + content_mask: Optional[Tensor] = None, + content_key_padding_mask: Optional[Tensor] = None, + update_content: bool = True, + ): + query_norm = self.norm_q(query) + content_norm = self.norm_c(content) + query = self.forward_stream(query, query_norm, content_norm, memory, + query_mask, content_key_padding_mask)[0] + if update_content: + content = self.forward_stream(content, content_norm, content_norm, + memory, content_mask, + content_key_padding_mask)[0] + return query, content + + +class Decoder(nn.Module): + __constants__ = ['norm'] + + def __init__(self, decoder_layer, num_layers, norm): + super().__init__() + self.layers = transformer._get_clones(decoder_layer, num_layers) + self.num_layers = num_layers + self.norm = norm + + def forward( + self, + query, + content, + memory, + query_mask: Optional[Tensor] = None, + content_mask: Optional[Tensor] = None, + content_key_padding_mask: Optional[Tensor] = None, + ): + for i, mod in enumerate(self.layers): + last = i == len(self.layers) - 1 + query, content = mod( + query, + content, + memory, + query_mask, + content_mask, + content_key_padding_mask, + update_content=not last, + ) + query = self.norm(query) + return query + + +class TokenEmbedding(nn.Module): + + def __init__(self, charset_size: int, embed_dim: int): + super().__init__() + self.embedding = nn.Embedding(charset_size, embed_dim) + self.embed_dim = embed_dim + + def forward(self, tokens: torch.Tensor): + return math.sqrt(self.embed_dim) * self.embedding(tokens) + + +def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None): + """Load a CLIP model + + Parameters + ---------- + name : str + A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict + + device : Union[str, torch.device] + The device to put the loaded model + + jit : bool + Whether to load the optimized JIT model or more hackable non-JIT model (default). + + download_root: str + path to download the model files; by default, it uses "~/.cache/clip" + + Returns + ------- + model : torch.nn.Module + The CLIP model + + preprocess : Callable[[PIL.Image], torch.Tensor] + A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input + """ + if name in _MODELS: + model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip")) + elif os.path.isfile(name): + model_path = name + else: + raise RuntimeError(f"Model {name} not found; available models = {available_models()}") + + with open(model_path, 'rb') as opened_file: + try: + # loading JIT archive + model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval() + state_dict = None + except RuntimeError: + # loading saved state dict + if jit: + warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead") + jit = False + state_dict = torch.load(opened_file, map_location="cpu") + + if not jit: + model = build_model(state_dict or model.state_dict()).to(device) + if str(device) == "cpu": + model.float() + return model, _transform(model.visual.input_resolution) + + # patch the device names + device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[]) + device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1] + + def patch_device(module): + try: + graphs = [module.graph] if hasattr(module, "graph") else [] + except RuntimeError: + graphs = [] + + if hasattr(module, "forward1"): + graphs.append(module.forward1.graph) + + for graph in graphs: + for node in graph.findAllNodes("prim::Constant"): + if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"): + node.copyAttributes(device_node) + + model.apply(patch_device) + patch_device(model.encode_image) + patch_device(model.encode_text) + + # patch dtype to float32 on CPU + if str(device) == "cpu": + float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[]) + float_input = list(float_holder.graph.findNode("aten::to").inputs())[1] + float_node = float_input.node() + + def patch_float(module): + try: + graphs = [module.graph] if hasattr(module, "graph") else [] + except RuntimeError: + graphs = [] + + if hasattr(module, "forward1"): + graphs.append(module.forward1.graph) + + for graph in graphs: + for node in graph.findAllNodes("aten::to"): + inputs = list(node.inputs()) + for i in [1, 2]: # dtype can be the second or third argument to aten::to() + if inputs[i].node()["value"] == 5: + inputs[i].node().copyAttributes(float_node) + + model.apply(patch_float) + patch_float(model.encode_image) + patch_float(model.encode_text) + + model.float() + + return model, _transform(model.input_resolution.item()) + +def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]: + """ + Returns the tokenized representation of given input string(s) + + Parameters + ---------- + texts : Union[str, List[str]] + An input string or a list of input strings to tokenize + + context_length : int + The context length to use; all CLIP models use 77 as the context length + + truncate: bool + Whether to truncate the text in case its encoding is longer than the context length + + Returns + ------- + A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]. + We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long. + """ + if isinstance(texts, str): + texts = [texts] + + sot_token = _tokenizer.encoder["<|startoftext|>"] + eot_token = _tokenizer.encoder["<|endoftext|>"] + all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts] + if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"): + result = torch.zeros(len(all_tokens), context_length, dtype=torch.long) + else: + result = torch.zeros(len(all_tokens), context_length, dtype=torch.int) + + for i, tokens in enumerate(all_tokens): + if len(tokens) > context_length: + if truncate: + tokens = tokens[:context_length] + tokens[-1] = eot_token + else: + raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}") + result[i, :len(tokens)] = torch.tensor(tokens) + + return result + +class DptrParseq(nn.Module): + + def __init__(self, + in_channels, + out_channels, + max_label_length=25, + embed_dim=512, + dec_num_heads=8, + dec_mlp_ratio=4, + dec_depth=6, + perm_num=6, + perm_forward=True, + perm_mirrored=True, + decode_ar=True, + refine_iters=1, + dropout=0.1, + is_pretrain=True, + ORP_path=None, + **kwargs: Any) -> None: + super().__init__() + self.pad_id = out_channels - 1 + self.eos_id = 0 + self.bos_id = out_channels - 2 + self.max_label_length = max_label_length + self.decode_ar = decode_ar + self.refine_iters = refine_iters + self.is_pretrain = is_pretrain + if not is_pretrain: + self.token_query = nn.Parameter(torch.Tensor(1, 26, embed_dim)) + self.fmu = FMU(embed_dim, dec_num_heads, embed_dim * dec_mlp_ratio, dropout) + + decoder_layer = DecoderLayer(embed_dim, dec_num_heads, embed_dim * dec_mlp_ratio, dropout) + self.decoder = Decoder(decoder_layer, + num_layers=dec_depth, + norm=nn.LayerNorm(embed_dim)) + + # Perm/attn mask stuff + self.rng = np.random.default_rng() + self.max_gen_perms = perm_num // 2 if perm_mirrored else perm_num + self.perm_forward = perm_forward + self.perm_mirrored = perm_mirrored + + # We don't predict
|---|