提交 144b022f 编写于 作者: W WenmuZhou

添加分类模型

上级 567c74c5
Global: Global:
algorithm: CLS algorithm: CLS
use_gpu: false use_gpu: False
epoch_num: 30 epoch_num: 100
log_smooth_window: 20 log_smooth_window: 20
print_batch_step: 10 print_batch_step: 100
save_model_dir: output/cls_mb3 save_model_dir: output/cls_mv3
save_epoch_step: 3 save_epoch_step: 3
eval_batch_step: 100 eval_batch_step: 500
train_batch_size_per_card: 256 train_batch_size_per_card: 512
test_batch_size_per_card: 256 test_batch_size_per_card: 512
image_shape: [3, 32, 100] image_shape: [3, 48, 192]
label_list: [0,180] label_list: ['0','180']
distort: True
reader_yml: ./configs/cls/cls_reader.yml reader_yml: ./configs/cls/cls_reader.yml
pretrain_weights: pretrain_weights:
checkpoints: /Users/zhoujun20/Desktop/code/class_model/cls_mb3_ultra_small_0.35/best_accuracy checkpoints:
save_inference_dir: save_inference_dir:
infer_img: /Users/zhoujun20/Desktop/code/PaddleOCR/doc/imgs_words/ch/word_1.jpg infer_img:
Architecture: Architecture:
function: ppocr.modeling.architectures.cls_model,ClsModel function: ppocr.modeling.architectures.cls_model,ClsModel
...@@ -23,7 +24,7 @@ Architecture: ...@@ -23,7 +24,7 @@ Architecture:
Backbone: Backbone:
function: ppocr.modeling.backbones.rec_mobilenet_v3,MobileNetV3 function: ppocr.modeling.backbones.rec_mobilenet_v3,MobileNetV3
scale: 0.35 scale: 0.35
model_name: Ultra_small model_name: small
Head: Head:
function: ppocr.modeling.heads.cls_head,ClsHead function: ppocr.modeling.heads.cls_head,ClsHead
...@@ -38,6 +39,6 @@ Optimizer: ...@@ -38,6 +39,6 @@ Optimizer:
beta1: 0.9 beta1: 0.9
beta2: 0.999 beta2: 0.999
decay: decay:
function: piecewise_decay function: cosine_decay
boundaries: [20,30] step_each_epoch: 1169
decay_rate: 0.1 total_epoch: 100
\ No newline at end of file
TrainReader: TrainReader:
reader_function: ppocr.data.cls.dataset_traversal,SimpleReader reader_function: ppocr.data.cls.dataset_traversal,SimpleReader
num_workers: 1 num_workers: 8
img_set_dir: / img_set_dir: ./train_data/cls
label_file_path: /Users/zhoujun20/Downloads/direction/rotate_ver/train.txt label_file_path: ./train_data/cls/train.txt
EvalReader: EvalReader:
reader_function: ppocr.data.cls.dataset_traversal,SimpleReader reader_function: ppocr.data.cls.dataset_traversal,SimpleReader
img_set_dir: / img_set_dir: ./train_data/cls
label_file_path: /Users/zhoujun20/Downloads/direction/rotate_ver/train.txt label_file_path: ./train_data/cls/test.txt
TestReader: TestReader:
reader_function: ppocr.data.cls.dataset_traversal,SimpleReader reader_function: ppocr.data.cls.dataset_traversal,SimpleReader
...@@ -57,6 +57,8 @@ public: ...@@ -57,6 +57,8 @@ public:
this->char_list_file.assign(config_map_["char_list_file"]); this->char_list_file.assign(config_map_["char_list_file"]);
this->use_angle_cls = bool(stoi(config_map_["use_angle_cls"]));
this->cls_model_dir.assign(config_map_["cls_model_dir"]); this->cls_model_dir.assign(config_map_["cls_model_dir"]);
this->cls_thresh = stod(config_map_["cls_thresh"]); this->cls_thresh = stod(config_map_["cls_thresh"]);
...@@ -88,6 +90,8 @@ public: ...@@ -88,6 +90,8 @@ public:
std::string rec_model_dir; std::string rec_model_dir;
bool use_angle_cls;
std::string char_list_file; std::string char_list_file;
std::string cls_model_dir; std::string cls_model_dir;
......
...@@ -58,7 +58,7 @@ public: ...@@ -58,7 +58,7 @@ public:
void LoadModel(const std::string &model_dir); void LoadModel(const std::string &model_dir);
void Run(std::vector<std::vector<std::vector<int>>> boxes, cv::Mat &img, void Run(std::vector<std::vector<std::vector<int>>> boxes, cv::Mat &img,
Classifier &cls); Classifier *cls);
private: private:
std::shared_ptr<PaddlePredictor> predictor_; std::shared_ptr<PaddlePredictor> predictor_;
......
...@@ -53,10 +53,15 @@ int main(int argc, char **argv) { ...@@ -53,10 +53,15 @@ int main(int argc, char **argv) {
config.cpu_math_library_num_threads, config.use_mkldnn, config.cpu_math_library_num_threads, config.use_mkldnn,
config.use_zero_copy_run, config.max_side_len, config.det_db_thresh, config.use_zero_copy_run, config.max_side_len, config.det_db_thresh,
config.det_db_box_thresh, config.det_db_unclip_ratio, config.visualize); config.det_db_box_thresh, config.det_db_unclip_ratio, config.visualize);
Classifier cls(config.cls_model_dir, config.use_gpu, config.gpu_id,
Classifier *cls = nullptr;
if (config.use_angle_cls == true) {
cls = new Classifier(config.cls_model_dir, config.use_gpu, config.gpu_id,
config.gpu_mem, config.cpu_math_library_num_threads, config.gpu_mem, config.cpu_math_library_num_threads,
config.use_mkldnn, config.use_zero_copy_run, config.use_mkldnn, config.use_zero_copy_run,
config.cls_thresh); config.cls_thresh);
}
CRNNRecognizer rec(config.rec_model_dir, config.use_gpu, config.gpu_id, CRNNRecognizer rec(config.rec_model_dir, config.use_gpu, config.gpu_id,
config.gpu_mem, config.cpu_math_library_num_threads, config.gpu_mem, config.cpu_math_library_num_threads,
config.use_mkldnn, config.use_zero_copy_run, config.use_mkldnn, config.use_zero_copy_run,
......
...@@ -17,7 +17,7 @@ ...@@ -17,7 +17,7 @@
namespace PaddleOCR { namespace PaddleOCR {
void CRNNRecognizer::Run(std::vector<std::vector<std::vector<int>>> boxes, void CRNNRecognizer::Run(std::vector<std::vector<std::vector<int>>> boxes,
cv::Mat &img, Classifier &cls) { cv::Mat &img, Classifier *cls) {
cv::Mat srcimg; cv::Mat srcimg;
img.copyTo(srcimg); img.copyTo(srcimg);
cv::Mat crop_img; cv::Mat crop_img;
...@@ -27,8 +27,9 @@ void CRNNRecognizer::Run(std::vector<std::vector<std::vector<int>>> boxes, ...@@ -27,8 +27,9 @@ void CRNNRecognizer::Run(std::vector<std::vector<std::vector<int>>> boxes,
int index = 0; int index = 0;
for (int i = boxes.size() - 1; i >= 0; i--) { for (int i = boxes.size() - 1; i >= 0; i--) {
crop_img = GetRotateCropImage(srcimg, boxes[i]); crop_img = GetRotateCropImage(srcimg, boxes[i]);
if (cls != nullptr) {
crop_img = cls.Run(crop_img); crop_img = cls->Run(crop_img);
}
float wh_ratio = float(crop_img.cols) / float(crop_img.rows); float wh_ratio = float(crop_img.cols) / float(crop_img.rows);
......
...@@ -4,23 +4,23 @@ gpu_id 0 ...@@ -4,23 +4,23 @@ gpu_id 0
gpu_mem 4000 gpu_mem 4000
cpu_math_library_num_threads 10 cpu_math_library_num_threads 10
use_mkldnn 0 use_mkldnn 0
use_zero_copy_run 1 use_zero_copy_run 0
# det config # det config
max_side_len 960 max_side_len 960
det_db_thresh 0.3 det_db_thresh 0.3
det_db_box_thresh 0.5 det_db_box_thresh 0.5
det_db_unclip_ratio 2.0 det_db_unclip_ratio 2.0
det_model_dir ./inference/det_db det_model_dir ../model/det
# cls config # cls config
cls_model_dir ./inference/cls use_angle_cls 1
cls_model_dir ../model/cls
cls_thresh 0.9 cls_thresh 0.9
# rec config # rec config
rec_model_dir ./inference/rec_crnn rec_model_dir ../model/rec
char_list_file ../../ppocr/utils/ppocr_keys_v1.txt char_list_file ../model/ppocr_keys_v1.txt
# show the detection results # show the detection results
visualize 1 visualize 1
\ No newline at end of file
## 文字角度分类
### 数据准备
请按如下步骤设置数据集:
训练数据的默认存储路径是 `PaddleOCR/train_data/cls`,如果您的磁盘上已有数据集,只需创建软链接至数据集目录:
```
ln -sf <path/to/dataset> <path/to/paddle_ocr>/train_data/cls/dataset
```
请参考下文组织您的数据。
- 训练集
首先请将训练图片放入同一个文件夹(train_images),并用一个txt文件(cls_gt_train.txt)记录图片路径和标签。
**注意:** 默认请将图片路径和图片标签用 `\t` 分割,如用其他方式分割将造成训练报错
0和180分别表示图片的角度为0度和180度
```
" 图像文件名 图像标注信息 "
train_data/cls/word_001.jpg 0
train_data/cls/word_002.jpg 180
```
最终训练集应有如下文件结构:
```
|-train_data
|-cls
|- cls_gt_train.txt
|- train
|- word_001.png
|- word_002.jpg
|- word_003.jpg
| ...
```
- 测试集
同训练集类似,测试集也需要提供一个包含所有图片的文件夹(test)和一个cls_gt_test.txt,测试集的结构如下所示:
```
|-train_data
|-cls
|- 和一个cls_gt_test.txt
|- test
|- word_001.jpg
|- word_002.jpg
|- word_003.jpg
| ...
```
### 启动训练
PaddleOCR提供了训练脚本、评估脚本和预测脚本。
开始训练:
*如果您安装的是cpu版本,请将配置文件中的 `use_gpu` 字段修改为false*
```
# 设置PYTHONPATH路径
export PYTHONPATH=$PYTHONPATH:.
# GPU训练 支持单卡,多卡训练,通过CUDA_VISIBLE_DEVICES指定卡号
export CUDA_VISIBLE_DEVICES=0,1,2,3
# 启动训练
python3 tools/train.py -c configs/cls/cls_mv3.yml
```
- 数据增强
PaddleOCR提供了多种数据增强方式,如果您希望在训练时加入扰动,请在配置文件中设置 `distort: true`
默认的扰动方式有:颜色空间转换(cvtColor)、模糊(blur)、抖动(jitter)、噪声(Gasuss noise)、随机切割(random crop)、透视(perspective)、颜色反转(reverse),随机数据增强(RandAugment)。
训练过程中除随机数据增强外每种扰动方式以50%的概率被选择,具体代码实现请参考:
[randaugment.py.py](https://github.com/PaddlePaddle/PaddleOCR/blob/develop/ppocr/data/cls/randaugment.py)
[img_tools.py](https://github.com/PaddlePaddle/PaddleOCR/blob/develop/ppocr/data/rec/img_tools.py)
*由于OpenCV的兼容性问题,扰动操作暂时只支持linux*
### 训练
PaddleOCR支持训练和评估交替进行, 可以在 `configs/cls/cls_mv3.yml` 中修改 `eval_batch_step` 设置评估频率,默认每500个iter评估一次。评估过程中默认将最佳acc模型,保存为 `output/cls_mv3/best_accuracy`
如果验证集很大,测试将会比较耗时,建议减少评估次数,或训练完再进行评估。
**注意,预测/评估时的配置文件请务必与训练一致。**
### 评估
评估数据集可以通过`configs/cls/cls_reader.yml` 修改EvalReader中的 `label_file_path` 设置。
*注意* 评估时必须确保配置文件中 infer_img 字段为空
```
export CUDA_VISIBLE_DEVICES=0
# GPU 评估, Global.checkpoints 为待测权重
python3 tools/eval.py -c configs/cls/cls_mv3.yml -o Global.checkpoints={path/to/weights}/best_accuracy
```
### 预测
* 训练引擎的预测
使用 PaddleOCR 训练好的模型,可以通过以下脚本进行快速预测。
默认预测图片存储在 `infer_img` 里,通过 `-o Global.checkpoints` 指定权重:
```
# 预测分类结果
python3 tools/infer_cls.py -c configs/cls/cls_mv3.yml -o Global.checkpoints={path/to/weights}/best_accuracy Global.infer_img=doc/imgs_words/en/word_1.png
```
预测图片:
![](../imgs_words/en/word_1.png)
得到输入图像的预测结果:
```
infer_img: doc/imgs_words/en/word_1.png
scores: [[0.93161047 0.06838956]]
label: [0]
```
## TEXT ANGLE CLASSIFICATION
### DATA PREPARATION
Please organize the dataset as follows:
The default storage path for training data is `PaddleOCR/train_data/cls`, if you already have a dataset on your disk, just create a soft link to the dataset directory:
```
ln -sf <path/to/dataset> <path/to/paddle_ocr>/train_data/cls/dataset
```
please refer to the following to organize your data.
- Training set
First put the training images in the same folder (train_images), and use a txt file (cls_gt_train.txt) to store the image path and label.
* Note: by default, the image path and image label are split with `\t`, if you use other methods to split, it will cause training error
0 and 180 indicate that the angle of the image is 0 degrees and 180 degrees, respectively.
```
" Image file name Image annotation "
train_data/word_001.jpg 0
train_data/word_002.jpg 180
```
The final training set should have the following file structure:
```
|-train_data
|-cls
|- cls_gt_train.txt
|- train
|- word_001.png
|- word_002.jpg
|- word_003.jpg
| ...
```
- Test set
Similar to the training set, the test set also needs to be provided a folder
containing all images (test) and a cls_gt_test.txt. The structure of the test set is as follows:
```
|-train_data
|-cls
|- cls_gt_test.txt
|- test
|- word_001.jpg
|- word_002.jpg
|- word_003.jpg
| ...
```
### TRAINING
PaddleOCR provides training scripts, evaluation scripts, and prediction scripts.
Start training:
```
# Set PYTHONPATH path
export PYTHONPATH=$PYTHONPATH:.
# GPU training Support single card and multi-card training, specify the card number through CUDA_VISIBLE_DEVICES
export CUDA_VISIBLE_DEVICES=0,1,2,3
# Training icdar15 English data
python3 tools/train.py -c configs/cls/cls_mv3.yml
```
- Data Augmentation
PaddleOCR provides a variety of data augmentation methods. If you want to add disturbance during training, please set `distort: true` in the configuration file.
The default perturbation methods are: cvtColor, blur, jitter, Gasuss noise, random crop, perspective, color reverse, RandAugment.
Except for RandAugment, each disturbance method is selected with a 50% probability during the training process. For specific code implementation, please refer to:
[randaugment.py.py](https://github.com/PaddlePaddle/PaddleOCR/blob/develop/ppocr/data/cls/randaugment.py)
[img_tools.py](https://github.com/PaddlePaddle/PaddleOCR/blob/develop/ppocr/data/rec/img_tools.py)
- Training
PaddleOCR supports alternating training and evaluation. You can modify `eval_batch_step` in `configs/cls/cls_mv3.yml` to set the evaluation frequency. By default, it is evaluated every 500 iter and the best acc model is saved under `output/cls_mv3/best_accuracy` during the evaluation process.
If the evaluation set is large, the test will be time-consuming. It is recommended to reduce the number of evaluations, or evaluate after training.
**Note that the configuration file for prediction/evaluation must be consistent with the training.**
### EVALUATION
The evaluation data set can be modified via `configs/cls/cls_reader.yml` setting of `label_file_path` in EvalReader.
```
export CUDA_VISIBLE_DEVICES=0
# GPU evaluation, Global.checkpoints is the weight to be tested
python3 tools/eval.py -c configs/cls/cls_mv3.yml -o Global.checkpoints={path/to/weights}/best_accuracy
```
### PREDICTION
* Training engine prediction
Using the model trained by paddleocr, you can quickly get prediction through the following script.
The default prediction picture is stored in `infer_img`, and the weight is specified via `-o Global.checkpoints`:
```
# Predict English results
python3 tools/infer_rec.py -c configs/cls/cls_mv3.yml -o Global.checkpoints={path/to/weights}/best_accuracy TestReader.infer_img=doc/imgs_words/en/word_1.jpg
```
Input image:
![](../imgs_words/en/word_1.png)
Get the prediction result of the input image:
```
infer_img: doc/imgs_words/en/word_1.png
scores: [[0.93161047 0.06838956]]
label: [0]
```
...@@ -14,6 +14,7 @@ ...@@ -14,6 +14,7 @@
import os import os
import sys import sys
import math
import random import random
import numpy as np import numpy as np
import cv2 import cv2
...@@ -23,7 +24,18 @@ from ppocr.utils.utility import get_image_file_list ...@@ -23,7 +24,18 @@ from ppocr.utils.utility import get_image_file_list
logger = initial_logger() logger = initial_logger()
from ppocr.data.rec.img_tools import warp, resize_norm_img from ppocr.data.rec.img_tools import resize_norm_img, warp
from ppocr.data.cls.randaugment import RandAugment
def random_crop(img):
img_h, img_w = img.shape[:2]
if img_w > img_h * 4:
w = random.randint(img_h * 2, img_w)
i = random.randint(0, img_w - w)
img = img[:, i:i + w, :]
return img
class SimpleReader(object): class SimpleReader(object):
...@@ -39,7 +51,8 @@ class SimpleReader(object): ...@@ -39,7 +51,8 @@ class SimpleReader(object):
self.image_shape = params['image_shape'] self.image_shape = params['image_shape']
self.mode = params['mode'] self.mode = params['mode']
self.infer_img = params['infer_img'] self.infer_img = params['infer_img']
self.use_distort = False self.use_distort = params['mode'] == 'train' and params['distort']
self.randaug = RandAugment()
self.label_list = params['label_list'] self.label_list = params['label_list']
if "distort" in params: if "distort" in params:
self.use_distort = params['distort'] and params['use_gpu'] self.use_distort = params['distort'] and params['use_gpu']
...@@ -76,6 +89,7 @@ class SimpleReader(object): ...@@ -76,6 +89,7 @@ class SimpleReader(object):
if img.shape[-1] == 1 or len(list(img.shape)) == 2: if img.shape[-1] == 1 or len(list(img.shape)) == 2:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
norm_img = resize_norm_img(img, self.image_shape) norm_img = resize_norm_img(img, self.image_shape)
norm_img = norm_img[np.newaxis, :] norm_img = norm_img[np.newaxis, :]
yield norm_img yield norm_img
else: else:
...@@ -97,6 +111,8 @@ class SimpleReader(object): ...@@ -97,6 +111,8 @@ class SimpleReader(object):
for img_id in range(process_id, img_num, self.num_workers): for img_id in range(process_id, img_num, self.num_workers):
label_infor = label_infor_list[img_id_list[img_id]] label_infor = label_infor_list[img_id_list[img_id]]
substr = label_infor.decode('utf-8').strip("\n").split("\t") substr = label_infor.decode('utf-8').strip("\n").split("\t")
label = self.label_list.index(substr[1])
img_path = self.img_set_dir + "/" + substr[0] img_path = self.img_set_dir + "/" + substr[0]
img = cv2.imread(img_path) img = cv2.imread(img_path)
if img is None: if img is None:
...@@ -105,12 +121,14 @@ class SimpleReader(object): ...@@ -105,12 +121,14 @@ class SimpleReader(object):
if img.shape[-1] == 1 or len(list(img.shape)) == 2: if img.shape[-1] == 1 or len(list(img.shape)) == 2:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
label = substr[1]
if self.use_distort: if self.use_distort:
# if random.randint(1, 100)>= 50:
# img = random_crop(img)
img = warp(img, 10) img = warp(img, 10)
img = self.randaug(img)
norm_img = resize_norm_img(img, self.image_shape) norm_img = resize_norm_img(img, self.image_shape)
norm_img = norm_img[np.newaxis, :] norm_img = norm_img[np.newaxis, :]
yield (norm_img, self.label_list.index(int(label))) yield (norm_img, label)
def batch_iter_reader(): def batch_iter_reader():
batch_outs = [] batch_outs = []
......
# 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.
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
from PIL import Image, ImageEnhance, ImageOps
import numpy as np
import random
import six
class RawRandAugment(object):
def __init__(self, num_layers=2, magnitude=5, fillcolor=(128, 128, 128)):
self.num_layers = num_layers
self.magnitude = magnitude
self.max_level = 10
abso_level = self.magnitude / self.max_level
self.level_map = {
"shearX": 0.3 * abso_level,
"shearY": 0.3 * abso_level,
"translateX": 150.0 / 331 * abso_level,
"translateY": 150.0 / 331 * abso_level,
"rotate": 30 * abso_level,
"color": 0.9 * abso_level,
"posterize": int(4.0 * abso_level),
"solarize": 256.0 * abso_level,
"contrast": 0.9 * abso_level,
"sharpness": 0.9 * abso_level,
"brightness": 0.9 * abso_level,
"autocontrast": 0,
"equalize": 0,
"invert": 0
}
# from https://stackoverflow.com/questions/5252170/
# specify-image-filling-color-when-rotating-in-python-with-pil-and-setting-expand
def rotate_with_fill(img, magnitude):
rot = img.convert("RGBA").rotate(magnitude)
return Image.composite(rot,
Image.new("RGBA", rot.size, (128, ) * 4),
rot).convert(img.mode)
rnd_ch_op = random.choice
self.func = {
"shearX": lambda img, magnitude: img.transform(
img.size,
Image.AFFINE,
(1, magnitude * rnd_ch_op([-1, 1]), 0, 0, 1, 0),
Image.BICUBIC,
fillcolor=fillcolor),
"shearY": lambda img, magnitude: img.transform(
img.size,
Image.AFFINE,
(1, 0, 0, magnitude * rnd_ch_op([-1, 1]), 1, 0),
Image.BICUBIC,
fillcolor=fillcolor),
"translateX": lambda img, magnitude: img.transform(
img.size,
Image.AFFINE,
(1, 0, magnitude * img.size[0] * rnd_ch_op([-1, 1]), 0, 1, 0),
fillcolor=fillcolor),
"translateY": lambda img, magnitude: img.transform(
img.size,
Image.AFFINE,
(1, 0, 0, 0, 1, magnitude * img.size[1] * rnd_ch_op([-1, 1])),
fillcolor=fillcolor),
"rotate": lambda img, magnitude: rotate_with_fill(img, magnitude),
"color": lambda img, magnitude: ImageEnhance.Color(img).enhance(
1 + magnitude * rnd_ch_op([-1, 1])),
"posterize": lambda img, magnitude:
ImageOps.posterize(img, magnitude),
"solarize": lambda img, magnitude:
ImageOps.solarize(img, magnitude),
"contrast": lambda img, magnitude:
ImageEnhance.Contrast(img).enhance(
1 + magnitude * rnd_ch_op([-1, 1])),
"sharpness": lambda img, magnitude:
ImageEnhance.Sharpness(img).enhance(
1 + magnitude * rnd_ch_op([-1, 1])),
"brightness": lambda img, magnitude:
ImageEnhance.Brightness(img).enhance(
1 + magnitude * rnd_ch_op([-1, 1])),
"autocontrast": lambda img, magnitude:
ImageOps.autocontrast(img),
"equalize": lambda img, magnitude: ImageOps.equalize(img),
"invert": lambda img, magnitude: ImageOps.invert(img)
}
def __call__(self, img):
avaiable_op_names = list(self.level_map.keys())
for layer_num in range(self.num_layers):
op_name = np.random.choice(avaiable_op_names)
img = self.func[op_name](img, self.level_map[op_name])
return img
class RandAugment(RawRandAugment):
""" RandAugment wrapper to auto fit different img types """
def __init__(self, *args, **kwargs):
if six.PY2:
super(RandAugment, self).__init__(*args, **kwargs)
else:
super().__init__(*args, **kwargs)
def __call__(self, img):
if not isinstance(img, Image.Image):
img = np.ascontiguousarray(img)
img = Image.fromarray(img)
if six.PY2:
img = super(RandAugment, self).__call__(img)
else:
img = super().__call__(img)
if isinstance(img, Image.Image):
img = np.asarray(img)
return img
...@@ -16,12 +16,9 @@ from __future__ import absolute_import ...@@ -16,12 +16,9 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
import logging
import numpy as np import numpy as np
import paddle.fluid as fluid __all__ = ['eval_cls_run']
__all__ = ['eval_class_run']
import logging import logging
...@@ -52,7 +49,8 @@ def eval_cls_run(exe, eval_info_dict): ...@@ -52,7 +49,8 @@ def eval_cls_run(exe, eval_info_dict):
fetch_list=eval_info_dict['fetch_varname_list'], \ fetch_list=eval_info_dict['fetch_varname_list'], \
return_numpy=False) return_numpy=False)
softmax_outs = np.array(outs[1]) softmax_outs = np.array(outs[1])
if len(softmax_outs.shape) != 1:
softmax_outs = np.array(outs[0])
acc, acc_num = cal_cls_acc(softmax_outs, label_list) acc, acc_num = cal_cls_acc(softmax_outs, label_list)
total_acc_num += acc_num total_acc_num += acc_num
total_sample_num += len(label_list) total_sample_num += len(label_list)
......
...@@ -108,7 +108,7 @@ class TextClassifier(object): ...@@ -108,7 +108,7 @@ class TextClassifier(object):
score = prob_out[rno][label_idx] score = prob_out[rno][label_idx]
label = self.label_list[label_idx] label = self.label_list[label_idx]
cls_res[indices[beg_img_no + rno]] = [label, score] cls_res[indices[beg_img_no + rno]] = [label, score]
if label == 180: if '180' in label and score > 0.9999:
img_list[indices[beg_img_no + rno]] = cv2.rotate( img_list[indices[beg_img_no + rno]] = cv2.rotate(
img_list[indices[beg_img_no + rno]], 1) img_list[indices[beg_img_no + rno]], 1)
return img_list, cls_res, predict_time return img_list, cls_res, predict_time
...@@ -130,12 +130,6 @@ def main(args): ...@@ -130,12 +130,6 @@ def main(args):
img_list.append(img) img_list.append(img)
try: try:
img_list, cls_res, predict_time = text_classifier(img_list) img_list, cls_res, predict_time = text_classifier(img_list)
print(cls_res)
from matplotlib import pyplot as plt
for img, angle in zip(img_list, cls_res):
plt.title(str(angle))
plt.imshow(img)
plt.show()
except Exception as e: except Exception as e:
print(e) print(e)
exit() exit()
......
...@@ -40,6 +40,8 @@ class TextSystem(object): ...@@ -40,6 +40,8 @@ class TextSystem(object):
def __init__(self, args): def __init__(self, args):
self.text_detector = predict_det.TextDetector(args) self.text_detector = predict_det.TextDetector(args)
self.text_recognizer = predict_rec.TextRecognizer(args) self.text_recognizer = predict_rec.TextRecognizer(args)
self.use_angle_cls = args.use_angle_cls
if self.use_angle_cls:
self.text_classifier = predict_cls.TextClassifier(args) self.text_classifier = predict_cls.TextClassifier(args)
def get_rotate_crop_image(self, img, points): def get_rotate_crop_image(self, img, points):
...@@ -95,10 +97,12 @@ class TextSystem(object): ...@@ -95,10 +97,12 @@ class TextSystem(object):
tmp_box = copy.deepcopy(dt_boxes[bno]) tmp_box = copy.deepcopy(dt_boxes[bno])
img_crop = self.get_rotate_crop_image(ori_im, tmp_box) img_crop = self.get_rotate_crop_image(ori_im, tmp_box)
img_crop_list.append(img_crop) img_crop_list.append(img_crop)
img_rotate_list, angle_list, elapse = self.text_classifier( if self.use_angle_cls:
img_crop_list, angle_list, elapse = self.text_classifier(
img_crop_list) img_crop_list)
print("cls num : {}, elapse : {}".format(len(img_rotate_list), elapse)) print("cls num : {}, elapse : {}".format(
rec_res, elapse = self.text_recognizer(img_rotate_list) len(img_crop_list), elapse))
rec_res, elapse = self.text_recognizer(img_crop_list)
print("rec_res num : {}, elapse : {}".format(len(rec_res), elapse)) print("rec_res num : {}, elapse : {}".format(len(rec_res), elapse))
# self.print_draw_crop_rec_res(img_crop_list, rec_res) # self.print_draw_crop_rec_res(img_crop_list, rec_res)
return dt_boxes, rec_res return dt_boxes, rec_res
......
...@@ -15,6 +15,7 @@ ...@@ -15,6 +15,7 @@
import argparse import argparse
import os, sys import os, sys
from ppocr.utils.utility import initial_logger from ppocr.utils.utility import initial_logger
logger = initial_logger() logger = initial_logger()
from paddle.fluid.core import PaddleTensor from paddle.fluid.core import PaddleTensor
from paddle.fluid.core import AnalysisConfig from paddle.fluid.core import AnalysisConfig
...@@ -31,34 +32,34 @@ def parse_args(): ...@@ -31,34 +32,34 @@ def parse_args():
return v.lower() in ("true", "t", "1") return v.lower() in ("true", "t", "1")
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
#params for prediction engine # params for prediction engine
parser.add_argument("--use_gpu", type=str2bool, default=True) parser.add_argument("--use_gpu", type=str2bool, default=True)
parser.add_argument("--ir_optim", type=str2bool, default=True) parser.add_argument("--ir_optim", type=str2bool, default=True)
parser.add_argument("--use_tensorrt", type=str2bool, default=False) parser.add_argument("--use_tensorrt", type=str2bool, default=False)
parser.add_argument("--gpu_mem", type=int, default=8000) parser.add_argument("--gpu_mem", type=int, default=8000)
#params for text detector # params for text detector
parser.add_argument("--image_dir", type=str) parser.add_argument("--image_dir", type=str)
parser.add_argument("--det_algorithm", type=str, default='DB') parser.add_argument("--det_algorithm", type=str, default='DB')
parser.add_argument("--det_model_dir", type=str) parser.add_argument("--det_model_dir", type=str)
parser.add_argument("--det_max_side_len", type=float, default=960) parser.add_argument("--det_max_side_len", type=float, default=960)
#DB parmas # DB parmas
parser.add_argument("--det_db_thresh", type=float, default=0.3) parser.add_argument("--det_db_thresh", type=float, default=0.3)
parser.add_argument("--det_db_box_thresh", type=float, default=0.5) parser.add_argument("--det_db_box_thresh", type=float, default=0.5)
parser.add_argument("--det_db_unclip_ratio", type=float, default=2.0) parser.add_argument("--det_db_unclip_ratio", type=float, default=2.0)
#EAST parmas # EAST parmas
parser.add_argument("--det_east_score_thresh", type=float, default=0.8) parser.add_argument("--det_east_score_thresh", type=float, default=0.8)
parser.add_argument("--det_east_cover_thresh", type=float, default=0.1) parser.add_argument("--det_east_cover_thresh", type=float, default=0.1)
parser.add_argument("--det_east_nms_thresh", type=float, default=0.2) parser.add_argument("--det_east_nms_thresh", type=float, default=0.2)
#SAST parmas # SAST parmas
parser.add_argument("--det_sast_score_thresh", type=float, default=0.5) parser.add_argument("--det_sast_score_thresh", type=float, default=0.5)
parser.add_argument("--det_sast_nms_thresh", type=float, default=0.2) parser.add_argument("--det_sast_nms_thresh", type=float, default=0.2)
parser.add_argument("--det_sast_polygon", type=bool, default=False) parser.add_argument("--det_sast_polygon", type=bool, default=False)
#params for text recognizer # params for text recognizer
parser.add_argument("--rec_algorithm", type=str, default='CRNN') parser.add_argument("--rec_algorithm", type=str, default='CRNN')
parser.add_argument("--rec_model_dir", type=str) parser.add_argument("--rec_model_dir", type=str)
parser.add_argument("--rec_image_shape", type=str, default="3, 32, 320") parser.add_argument("--rec_image_shape", type=str, default="3, 32, 320")
...@@ -72,13 +73,14 @@ def parse_args(): ...@@ -72,13 +73,14 @@ def parse_args():
parser.add_argument("--use_space_char", type=bool, default=True) parser.add_argument("--use_space_char", type=bool, default=True)
# params for text classifier # params for text classifier
parser.add_argument("--use_angle_cls", type=str2bool, default=True)
parser.add_argument("--cls_model_dir", type=str) parser.add_argument("--cls_model_dir", type=str)
parser.add_argument("--cls_image_shape", type=str, default="3, 32, 100") parser.add_argument("--cls_image_shape", type=str, default="3, 48, 192")
parser.add_argument("--label_list", type=list, default=[0, 180]) parser.add_argument("--label_list", type=list, default=['0', '180'])
parser.add_argument("--cls_batch_num", type=int, default=30) parser.add_argument("--cls_batch_num", type=int, default=30)
parser.add_argument("--enable_mkldnn", type=bool, default=False) parser.add_argument("--enable_mkldnn", type=str2bool, default=False)
parser.add_argument("--use_zero_copy_run", type=bool, default=False) parser.add_argument("--use_zero_copy_run", type=str2bool, default=False)
return parser.parse_args() return parser.parse_args()
...@@ -112,7 +114,7 @@ def create_predictor(args, mode): ...@@ -112,7 +114,7 @@ def create_predictor(args, mode):
if args.enable_mkldnn: if args.enable_mkldnn:
config.enable_mkldnn() config.enable_mkldnn()
#config.enable_memory_optim() # config.enable_memory_optim()
config.disable_glog_info() config.disable_glog_info()
if args.use_zero_copy_run: if args.use_zero_copy_run:
......
...@@ -85,9 +85,10 @@ def main(): ...@@ -85,9 +85,10 @@ def main():
feed={"image": img}, feed={"image": img},
fetch_list=fetch_varname_list, fetch_list=fetch_varname_list,
return_numpy=False) return_numpy=False)
for k in predict: scores = np.array(predict[0])
k = np.array(k) label = np.array(predict[1])
print(k) logger.info('\t scores: {}'.format(scores))
logger.info('\t label: {}'.format(label))
# save for inference model # save for inference model
target_var = [] target_var = []
for key, values in outputs.items(): for key, values in outputs.items():
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册