Skip to content

Commit

Permalink
add new pipeline
Browse files Browse the repository at this point in the history
  • Loading branch information
horcham committed Apr 23, 2024
1 parent 7971ee0 commit ef52103
Show file tree
Hide file tree
Showing 47 changed files with 2,954 additions and 2 deletions.
45 changes: 45 additions & 0 deletions configs/cls/mobilenetv3/cls_mv3.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -142,3 +142,48 @@ eval:
drop_remainder: False
max_rowsize: 12
num_workers: 8

predict:
deive_target: Ascend
device_id: 0
max_device_memory: 8GB
amp_level: O2
mode: 0
ckpt_load_path: /root/.mindspore/models/cls_mobilenetv3-92db9c58.ckpt
dataset_sink_mode: False
dataset:
type: RecDataset
dataset_root: dir/to/dataset
data_dir: all_images
label_file: val_cls_gt.txt
sample_ratio: 1.0
shuffle: False
transform_pipeline:
- DecodeImage:
img_mode: BGR
to_float32: False
- Rotate90IfVertical:
threshold: 2.0
direction: counterclockwise
- ClsLabelEncode:
label_list: *label_list
- RecResizeImg:
image_shape: [48, 192] # H, W
padding: False # aspect ratio will be preserved if true.
- NormalizeImage:
bgr_to_rgb: True
is_hwc: True
mean : [127.0, 127.0, 127.0]
std : [127.0, 127.0, 127.0]
- ToCHWImage:
# the order of the dataloader list, matching the network input and the input labels for the loss function, and optional data for debug/visaulize
output_columns: ['image', 'label'] # TODO return text string padding w/ fixed length, and a scaler to indicate the length
net_input_column_index: [0] # input indices for network forward func in output_columns
label_column_index: [1] # input indices marked as label

loader:
shuffle: False
batch_size: 8
drop_remainder: False
max_rowsize: 12
num_workers: 8
7 changes: 6 additions & 1 deletion configs/det/dbnet/db_r50_icdar15.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,12 @@ eval:
num_workers: 2

predict:
ckpt_load_path: tmp_det/best.ckpt
deive_target: Ascend
device_id: 0
max_device_memory: 8GB
amp_level: O2
mode: 0
ckpt_load_path: /root/.mindspore/models/dbnet_resnet50-c3a4aa24.ckpt
output_save_dir: ./output
dataset_sink_mode: False
dataset:
Expand Down
7 changes: 6 additions & 1 deletion configs/rec/crnn/crnn_resnet34.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -150,7 +150,12 @@ eval:
num_workers: 8

predict:
ckpt_load_path: ./tmp_rec/best.ckpt
deive_target: Ascend
device_id: 0
max_device_memory: 8GB
amp_level: O2
mode: 0
ckpt_load_path: /root/.mindspore/models/crnn_resnet34-83f37f07.ckpt
vis_font_path: tools/utils/simfang.ttf
dataset_sink_mode: False
dataset:
Expand Down
Empty file added mindocr/infer/__init__.py
Empty file.
3 changes: 3 additions & 0 deletions mindocr/infer/classification/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
from .cls_infer_node import ClsInferNode
from .cls_post_node import ClsPostNode
from .cls_pre_node import ClsPreNode
63 changes: 63 additions & 0 deletions mindocr/infer/classification/classification.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
import logging
import os
import time
import sys
import numpy as np
import yaml
from addict import Dict
from typing import List

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.abspath(os.path.join(__dir__, "../../")))

from tools.infer.text.utils import get_ckpt_file
from mindocr.data.transforms import create_transforms, run_transforms
from mindocr.postprocess import build_postprocess
from mindocr.infer.utils.model import MSModel, LiteModel


algo_to_model_name = {
"MV3": "cls_mobilenet_v3_small_100_model",
}
logger = logging.getLogger("mindocr")

class ClsPreprocess(object):
def __init__(self, args):
self.args = args
with open(args.cls_model_name_or_config, "r") as f:
self.yaml_cfg = Dict(yaml.safe_load(f))
self.transforms = create_transforms(self.yaml_cfg.predict.dataset.transform_pipeline)

def __call__(self, img):
data = {"image": img}
# ZHQ TODO: [1:] ???
data = run_transforms(img, self.transforms[1:])
return data


class ClsModelMS(MSModel):
def __init__(self, args):
self.args = args
self.model_name = algo_to_model_name[args.cls_algorithm]
self.config_path = args.cls_config_path
self._init_model(self.model_name, self.config_path)


class ClsModelLite(LiteModel):
def __init__(self, args):
self.args = args
self.model_name = algo_to_model_name[args.cls_algorithm]
self.config_path = args.cls_config_path
self._init_model(self.model_name, self.config_path)

INFER_CLS_MAP = {"MindSporeLite": ClsModelLite, "MindSpore": ClsModelMS}

class ClsPostProcess(object):
def __init__(self, args):
self.args = args
with open(args.cls_model_name_or_config, "r") as f:
self.yaml_cfg = Dict(yaml.safe_load(f))
self.postprocessor = build_postprocess(self.yaml_cfg.postprocess)

def __call__(self, pred):
return self.postprocessor(pred)
57 changes: 57 additions & 0 deletions mindocr/infer/classification/cls_infer_node.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
import os
import time
import sys
import numpy as np
import yaml
from addict import Dict

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.abspath(os.path.join(__dir__, "../../")))

from pipeline.framework.module_base import ModuleBase
from pipeline.tasks import TaskType
from infer.classification.classification import INFER_CLS_MAP


class ClsInferNode(ModuleBase):
def __init__(self, args, msg_queue, tqdm_info):
super(ClsInferNode, self).__init__(args, msg_queue, tqdm_info)
self.args = args
self.cls_model = None
self.task_type = self.args.task_type

def init_self_args(self):
with open(self.args.cls_model_name_or_config, "r") as f:
self.yaml_cfg = Dict(yaml.safe_load(f))
self.batch_size = self.yaml_cfg.predict.loader.batch_size
ClsModel = INFER_CLS_MAP[self.yaml_cfg.predict.backend]
self.cls_model = ClsModel(self.args)
super().init_self_args()

def process(self, input_data):
"""
Input:
- input_data.data: [np.ndarray], shape:[3,w,h], e.g. [3,48,192]
Output:
- input_data.data: [np.ndarray], shape:[?,2]
"""
if input_data.skip:
self.send_to_next_module(input_data)
return

data = input_data.data
data = [np.expand_dims(d, 0) for d in data if len(d.shape) == 3]
data = np.concatenate(data, axis=0)

preds = []
for batch_i in range(data.shape[0] // self.batch_size + 1):
start_i = batch_i * self.batch_size
end_i = (batch_i + 1) * self.batch_size
d = data[start_i:end_i]
if d.shape[0] == 0:
continue
pred = self.cls_model([d])
preds.append(pred[0])
preds = np.concatenate(preds, axis=0)
input_data.data = {"pred": preds}
self.send_to_next_module(input_data)
63 changes: 63 additions & 0 deletions mindocr/infer/classification/cls_post_node.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
import os
import time
import sys
import numpy as np
import yaml
from addict import Dict

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.abspath(os.path.join(__dir__, "../../")))

import cv2
import numpy as np

from pipeline.framework.module_base import ModuleBase
from pipeline.tasks import TaskType
from infer.classification.classification import ClsPostProcess
from tools.infer.text.utils import crop_text_region
from pipeline.data_process.utils.cv_utils import crop_box_from_image


class ClsPostNode(ModuleBase):
def __init__(self, args, msg_queue, tqdm_info):
super(ClsPostNode, self).__init__(args, msg_queue, tqdm_info)
self.cls_postprocess = ClsPostProcess(args)
self.task_type = self.args.task_type
self.cls_thresh = 0.9

def init_self_args(self):
super().init_self_args()

def process(self, input_data):
"""
Input:
- input_data.data: [np.ndarray], shape:[?,2]
Output:
- input_data.sub_image_list: [np.ndarray], shape:[1,3,-1,-1], e.g. [1,3,48,192]
- input_data.data = None
or
- input_data.infer_result = [(str, float)]
"""
if input_data.skip:
self.send_to_next_module(input_data)
return

data = input_data.data
pred = data["pred"]
output = self.cls_postprocess(pred)
angles = output["angles"]
scores = np.array(output["scores"]).tolist()

batch = input_data.sub_image_size
if self.task_type.value == TaskType.DET_CLS_REC.value:
sub_images = input_data.sub_image_list
for i in range(batch):
angle, score = angles[i], scores[i]
if "180" == angle and score > self.cls_thresh:
sub_images[i] = cv2.rotate(sub_images[i], cv2.ROTATE_180)
input_data.sub_image_list = sub_images
else:
input_data.infer_result = [(angle, score) for angle, score in zip(angles, scores)]

input_data.data = None
self.send_to_next_module(input_data)
41 changes: 41 additions & 0 deletions mindocr/infer/classification/cls_pre_node.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
import argparse
import os
import time
import sys
import numpy as np

__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.abspath(os.path.join(__dir__, "../../")))

from pipeline.framework.module_base import ModuleBase
from pipeline.tasks import TaskType
from infer.classification.classification import ClsPreProcess


class ClsPreNode(ModuleBase):
def __init__(self, args, msg_queue, tqdm_info):
super(ClsPreNode, self).__init__(args, msg_queue, tqdm_info)
self.cls_preprocesser = ClsPreProcess(args)
self.task_type = self.args.task_type

def init_self_args(self):
super().init_self_args()
return {"batch_size": 1}

def process(self, input_data):
if input_data.skip:
self.send_to_next_module(input_data)
return

if self.task_type.value == TaskType.REC.value:
image = input_data.frame[0]
data = [self.cls_preprocesser(image)["image"]]
input_data.sub_image_size = 1
input_data.data = data
self.send_to_next_module(input_data)
else:
sub_image_list = input_data.sub_image_list
data = [self.cls_preprocesser(image)["image"] for split_image in sub_image_list]
input_data.sub_image_size = len(sub_image_list)
input_data.data = data
self.send_to_next_module(input_data)
3 changes: 3 additions & 0 deletions mindocr/infer/common/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
from .collect_node import CollectNode
from .decode_node import DecodeNode
from .handout_node import HandoutNode
Loading

0 comments on commit ef52103

Please sign in to comment.