Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+axmodel/backbone_encoder.axmodel filter=lfs diff=lfs merge=lfs -text
+axmodel/decoder.axmodel filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,88 @@
----
-license: bsd-3-clause-clear
----
+# satrn
+
+[original repo](https://github.com/open-mmlab/mmocr/blob/main/configs/textrecog/satrn/README.md)
+
+## Convert tools links:
+
+For those who are interested in model conversion, you can try to export onnx or axmodel through 
+
+[satrn.axera](https://github.com/AXERA-TECH/satrn.axera)
+
+
+## Support Platform
+
+- AX650
+  - [M4N-Dock(爱芯派Pro)](https://wiki.sipeed.com/hardware/zh/maixIV/m4ndock/m4ndock.html)
+  - [M.2 Accelerator card](https://axcl-docs.readthedocs.io/zh-cn/latest/doc_guide_hardware.html)
+
+
+ The speed measurements(under different NPU configurations ) of the two parts of SATRN:
+
+ (1) backbone+encoder
+
+ (2) decoder
+
+
+||backbone+encoder(ms)|decoder(ms)|
+|--|--|--|
+|NPU1|20.494|2.648|
+|NPU2|9.785|1.504|
+|NPU3|6.085|1.384|
+
+## How to use
+
+Download all files from this repository to the device
+
+```
+.
+├── axmodel
+│   ├── backbone_encoder.axmodel
+│   └── decoder.axmodel
+├── demo_text_recog.jpg
+├── onnx
+│   ├── satrn_backbone_encoder.onnx
+│   └── satrn_decoder_sim.onnx
+├── README.md
+├── run_axmodel.py
+├── run_model.py
+└── run_onnx.py
+```
+
+### python env requirement
+
+#### 1. pyaxengine
+
+https://github.com/AXERA-TECH/pyaxengine
+
+```
+wget https://github.com/AXERA-TECH/pyaxengine/releases/download/0.1.1rc0/axengine-0.1.1-py3-none-any.whl
+pip install axengine-0.1.1-py3-none-any.whl
+```
+
+#### 2. satrn
+
+[satrn installation](https://github.com/open-mmlab/mmocr/tree/main?tab=readme-ov-file#installation)
+
+#### Inference onnxmodel
+```
+python run_onnx.py
+```
+input:
+
+![](demo_text_recog.jpg)
+
+output:
+
+```
+pred_text: STAR
+score: [0.9384028315544128, 0.9574984908103943, 0.9993689656257629, 0.9994958639144897]
+```
+
+#### Inference with AX650 Host
+
+```
+python run_axmodel.py
+```
+
+
+

axmodel/backbone_encoder.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca1bc3964ad5b7d57a2c5b08b0ca53619127501aed402f02829a53c26b021756
+size 47589096

axmodel/decoder.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1914ff3a36e5e2d9d2e6174bdbb8e5c369374e3c8420e22c445771ab1406347e
+size 27697793

onnx/satrn_backbone_encoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66870bd86006213fcb0e5db1d5b0e376d6a4f30c0e20e9f34cda66d3c259f39c
+size 161383339

onnx/satrn_decoder_sim.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d417bc9618b68e4f691d3f9571c93f9d101b30da92e47924d4e25ad3e37f8198
+size 101341850

run_axmodel.py

@@ -0,0 +1,154 @@
+from mmocr.apis import MMOCRInferencer
+from mmocr.apis.inferencers.base_mmocr_inferencer import BaseMMOCRInferencer
+import torch
+from rich.progress import track
+import torch.nn as nn
+import axengine as axe
+import numpy as np
+
+onnx_bb_encoder = axe.InferenceSession("satrn_backbone_encoder.axmodel")
+onnx_decoder = axe.InferenceSession("satrn_decoder.axmodel")
+
+
+class BackboneEncoderOnly(nn.Module):
+    def __init__(self, original_model):
+        super().__init__()
+        # 保留 backbone 和 encoder
+        self.backbone = original_model.backbone
+        self.encoder = original_model.encoder
+        
+    def forward(self, x):
+        x = self.backbone(x)
+        return self.encoder(x)
+
+
+class DecoderOnly(nn.Module):
+    def __init__(self, original_model):
+        super().__init__()
+        # 保留 backbone 和 encoder
+        original_decoder = original_model.decoder
+        # self._attention = original_decoder._attention
+        self.classifier = original_decoder.classifier
+        self.trg_word_emb = original_decoder.trg_word_emb
+        self.position_enc = original_decoder.position_enc
+        self._get_target_mask = original_decoder._get_target_mask
+        self.dropout = original_decoder.dropout
+        self.layer_stack = original_decoder.layer_stack
+        self.layer_norm = original_decoder.layer_norm
+        self._get_source_mask = original_decoder._get_source_mask
+        self.postprocessor = original_decoder.postprocessor
+        self.start_idx = 90
+        self.padding_idx = 91
+        self.max_seq_len = 25
+        self.softmax = nn.Softmax(dim=-1)
+        
+        
+        
+    def forward(self, trg_seq,src,src_mask,step):
+        # decoder_output = self._attention(init_target_seq, out_enc, src_mask=src_mask)
+        trg_embedding = self.trg_word_emb(trg_seq)
+        trg_pos_encoded = self.position_enc(trg_embedding)
+        trg_mask = self._get_target_mask(trg_seq)
+        tgt_seq = self.dropout(trg_pos_encoded)
+
+        output = tgt_seq
+        for dec_layer in self.layer_stack:
+            output = dec_layer(
+                output,
+                src,
+                self_attn_mask=trg_mask,
+                dec_enc_attn_mask=src_mask)
+        output = self.layer_norm(output)
+        # bsz * seq_len * C
+        step_result = self.classifier(output[:, step, :])
+        return step_result
+
+
+
+def normalize_tensor(tensor: torch.Tensor) -> torch.Tensor:
+    """
+    对 uint8 张量进行标准化处理
+    参数:
+        tensor: 输入张量，形状为 [3, 32, 100]，数据类型为 uint8
+    返回:
+        标准化后的张量，形状不变，数据类型为 float32
+    """
+    # 检查输入张量的形状和数据类型
+    assert tensor.shape == (3, 32, 100), "输入张量形状必须为 [3, 32, 100]"
+    assert tensor.dtype == torch.uint8, "输入张量数据类型必须为 uint8"
+    
+    # 转换为 float32 类型
+    tensor = tensor.float()
+    
+    # 定义标准化参数（RGB 通道顺序）
+    mean = torch.tensor([123.675, 116.28, 103.53], dtype=torch.float32).view(3, 1, 1)
+    std = torch.tensor([58.395, 57.12, 57.375], dtype=torch.float32).view(3, 1, 1)
+    
+    # 执行标准化：(x - mean) / std
+    normalized_tensor = (tensor - mean) / std
+    
+    return normalized_tensor
+
+
+infer = MMOCRInferencer(rec='satrn')
+model = infer.textrec_inferencer.model
+model.eval()
+model.cpu()
+input_path = 'mmor_demo/demo/demo_text_recog.jpg'
+ori_inputs = infer._inputs_to_list([input_path])
+base = BaseMMOCRInferencer(model='satrn')
+chunked_inputs = base._get_chunk_data(ori_inputs, 1)
+for ori_inputs in track(chunked_inputs, description='Inference'):
+    input = ori_inputs[0][1]
+    input_img = input['inputs']
+    input_image = normalize_tensor(input_img).unsqueeze(0)
+    input_sample = input['data_samples']
+    
+    # backbone+encoder
+    model_backbone_encoder = BackboneEncoderOnly(model)
+    model_decoder = DecoderOnly(model)
+    
+    # out_enc = model_backbone_encoder(input_image)
+    out_enc = onnx_bb_encoder.run(None, {"input": np.array(input_image.cpu())})[0]
+    out_enc = torch.tensor(out_enc)
+    data_samples = None
+    
+    N = out_enc.size(0)
+    init_target_seq = torch.full((N, model_decoder.max_seq_len + 1),
+                                model_decoder.padding_idx,
+                                device=out_enc.device,
+                                dtype=torch.long)
+# bsz * seq_len
+    init_target_seq[:, 0] = model_decoder.start_idx
+
+    outputs = []
+    for step in range(0, model_decoder.max_seq_len):
+        valid_ratios = [1.0 for _ in range(out_enc.size(0))]
+        if data_samples is not None:
+            valid_ratios = []
+            for data_sample in data_samples:
+                valid_ratios.append(data_sample.get('valid_ratio'))
+        
+        src_mask = model_decoder._get_source_mask(out_enc, valid_ratios)
+        # step_result = model_decoder(init_target_seq,out_enc,src_mask,step)
+        step_result = onnx_decoder.run(None,{'init_target_seq':np.array(init_target_seq),
+                                             'out_enc':np.array(out_enc),
+                                             'src_mask':np.array(src_mask),
+                                             'step':np.array([step])})[0][0]
+        step_result = torch.tensor(step_result)
+        outputs.append(step_result)
+        _, step_max_index = torch.max(step_result, dim=-1)
+        init_target_seq[:, step + 1] = step_max_index
+    outputs = torch.stack(outputs, dim=1)
+    out_dec = model_decoder.softmax(outputs)
+    output = model_decoder.postprocessor(out_dec, [input_sample])
+    outstr = output[0].pred_text.item
+    outscore = output[0].pred_text.score
+    
+    print('pred_text:',outstr)
+    print('score:',outscore)
+    
+    
+    
+    
+    

run_model.py

@@ -0,0 +1,140 @@
+from mmocr.apis import MMOCRInferencer
+from mmocr.apis.inferencers.base_mmocr_inferencer import BaseMMOCRInferencer
+import torch
+from rich.progress import track
+import torch.nn as nn
+
+class BackboneEncoderOnly(nn.Module):
+    def __init__(self, original_model):
+        super().__init__()
+        # 保留 backbone 和 encoder
+        self.backbone = original_model.backbone
+        self.encoder = original_model.encoder
+        
+    def forward(self, x):
+        x = self.backbone(x)
+        return self.encoder(x)
+
+
+class DecoderOnly(nn.Module):
+    def __init__(self, original_model):
+        super().__init__()
+        # 保留 backbone 和 encoder
+        original_decoder = original_model.decoder
+        # self._attention = original_decoder._attention
+        self.classifier = original_decoder.classifier
+        self.trg_word_emb = original_decoder.trg_word_emb
+        self.position_enc = original_decoder.position_enc
+        self._get_target_mask = original_decoder._get_target_mask
+        self.dropout = original_decoder.dropout
+        self.layer_stack = original_decoder.layer_stack
+        self.layer_norm = original_decoder.layer_norm
+        self._get_source_mask = original_decoder._get_source_mask
+        self.postprocessor = original_decoder.postprocessor
+        self.start_idx = 90
+        self.padding_idx = 91
+        self.max_seq_len = 25
+        self.softmax = nn.Softmax(dim=-1)
+        
+        
+        
+    def forward(self, trg_seq,src,src_mask,step):
+        # decoder_output = self._attention(init_target_seq, out_enc, src_mask=src_mask)
+        trg_embedding = self.trg_word_emb(trg_seq)
+        trg_pos_encoded = self.position_enc(trg_embedding)
+        trg_mask = self._get_target_mask(trg_seq)
+        tgt_seq = self.dropout(trg_pos_encoded)
+
+        output = tgt_seq
+        for dec_layer in self.layer_stack:
+            output = dec_layer(
+                output,
+                src,
+                self_attn_mask=trg_mask,
+                dec_enc_attn_mask=src_mask)
+        output = self.layer_norm(output)
+        # bsz * seq_len * C
+        step_result = self.classifier(output[:, step, :])
+        return step_result
+
+
+def normalize_tensor(tensor: torch.Tensor) -> torch.Tensor:
+    """
+    对 uint8 张量进行标准化处理
+    参数:
+        tensor: 输入张量，形状为 [3, 32, 100]，数据类型为 uint8
+    返回:
+        标准化后的张量，形状不变，数据类型为 float32
+    """
+    # 检查输入张量的形状和数据类型
+    assert tensor.shape == (3, 32, 100), "输入张量形状必须为 [3, 32, 100]"
+    assert tensor.dtype == torch.uint8, "输入张量数据类型必须为 uint8"
+    
+    # 转换为 float32 类型
+    tensor = tensor.float()
+    
+    # 定义标准化参数（RGB 通道顺序）
+    mean = torch.tensor([123.675, 116.28, 103.53], dtype=torch.float32).view(3, 1, 1)
+    std = torch.tensor([58.395, 57.12, 57.375], dtype=torch.float32).view(3, 1, 1)
+    
+    # 执行标准化：(x - mean) / std
+    normalized_tensor = (tensor - mean) / std
+    
+    return normalized_tensor
+
+
+infer = MMOCRInferencer(rec='satrn')
+model = infer.textrec_inferencer.model
+model.eval()
+model.cpu()
+input_path = 'mmor_demo/demo/demo_text_recog.jpg'
+ori_inputs = infer._inputs_to_list([input_path])
+base = BaseMMOCRInferencer(model='satrn')
+chunked_inputs = base._get_chunk_data(ori_inputs, 1)
+for ori_inputs in track(chunked_inputs, description='Inference'):
+    input = ori_inputs[0][1]
+    input_img = input['inputs']
+    input_image = normalize_tensor(input_img).unsqueeze(0)
+    input_sample = input['data_samples']
+    
+    # backbone+encoder
+    model_backbone_encoder = BackboneEncoderOnly(model)
+    model_decoder = DecoderOnly(model)
+    
+    out_enc = model_backbone_encoder(input_image)
+    data_samples = None
+    
+    N = out_enc.size(0)
+    init_target_seq = torch.full((N, model_decoder.max_seq_len + 1),
+                                model_decoder.padding_idx,
+                                device=out_enc.device,
+                                dtype=torch.long)
+# bsz * seq_len
+    init_target_seq[:, 0] = model_decoder.start_idx
+
+    outputs = []
+    for step in range(0, model_decoder.max_seq_len):
+        valid_ratios = [1.0 for _ in range(out_enc.size(0))]
+        if data_samples is not None:
+            valid_ratios = []
+            for data_sample in data_samples:
+                valid_ratios.append(data_sample.get('valid_ratio'))
+        
+        src_mask = model_decoder._get_source_mask(out_enc, valid_ratios)
+        step_result = model_decoder(init_target_seq,out_enc,src_mask,step)
+        outputs.append(step_result)
+        _, step_max_index = torch.max(step_result, dim=-1)
+        init_target_seq[:, step + 1] = step_max_index
+    outputs = torch.stack(outputs, dim=1)
+    out_dec = model_decoder.softmax(outputs)
+    output = model_decoder.postprocessor(out_dec, [input_sample])
+    outstr = output[0].pred_text.item
+    outscore = output[0].pred_text.score
+    
+    print('pred_text:',outstr)
+    print('score:',outscore)
+    
+    
+    
+    
+    

run_onnx.py

@@ -0,0 +1,154 @@
+from mmocr.apis import MMOCRInferencer
+from mmocr.apis.inferencers.base_mmocr_inferencer import BaseMMOCRInferencer
+import torch
+from rich.progress import track
+import torch.nn as nn
+import onnxruntime as ort
+import numpy as np
+
+onnx_bb_encoder = ort.InferenceSession("onnx/satrn_backbone_encoder.onnx")
+onnx_decoder = ort.InferenceSession("onnx/satrn_decoder_sim.onnx")
+
+
+class BackboneEncoderOnly(nn.Module):
+    def __init__(self, original_model):
+        super().__init__()
+        # 保留 backbone 和 encoder
+        self.backbone = original_model.backbone
+        self.encoder = original_model.encoder
+        
+    def forward(self, x):
+        x = self.backbone(x)
+        return self.encoder(x)
+
+
+class DecoderOnly(nn.Module):
+    def __init__(self, original_model):
+        super().__init__()
+        # 保留 backbone 和 encoder
+        original_decoder = original_model.decoder
+        # self._attention = original_decoder._attention
+        self.classifier = original_decoder.classifier
+        self.trg_word_emb = original_decoder.trg_word_emb
+        self.position_enc = original_decoder.position_enc
+        self._get_target_mask = original_decoder._get_target_mask
+        self.dropout = original_decoder.dropout
+        self.layer_stack = original_decoder.layer_stack
+        self.layer_norm = original_decoder.layer_norm
+        self._get_source_mask = original_decoder._get_source_mask
+        self.postprocessor = original_decoder.postprocessor
+        self.start_idx = 90
+        self.padding_idx = 91
+        self.max_seq_len = 25
+        self.softmax = nn.Softmax(dim=-1)
+        
+        
+        
+    def forward(self, trg_seq,src,src_mask,step):
+        # decoder_output = self._attention(init_target_seq, out_enc, src_mask=src_mask)
+        trg_embedding = self.trg_word_emb(trg_seq)
+        trg_pos_encoded = self.position_enc(trg_embedding)
+        trg_mask = self._get_target_mask(trg_seq)
+        tgt_seq = self.dropout(trg_pos_encoded)
+
+        output = tgt_seq
+        for dec_layer in self.layer_stack:
+            output = dec_layer(
+                output,
+                src,
+                self_attn_mask=trg_mask,
+                dec_enc_attn_mask=src_mask)
+        output = self.layer_norm(output)
+        # bsz * seq_len * C
+        step_result = self.classifier(output[:, step, :])
+        return step_result
+
+
+
+def normalize_tensor(tensor: torch.Tensor) -> torch.Tensor:
+    """
+    对 uint8 张量进行标准化处理
+    参数:
+        tensor: 输入张量，形状为 [3, 32, 100]，数据类型为 uint8
+    返回:
+        标准化后的张量，形状不变，数据类型为 float32
+    """
+    # 检查输入张量的形状和数据类型
+    assert tensor.shape == (3, 32, 100), "输入张量形状必须为 [3, 32, 100]"
+    assert tensor.dtype == torch.uint8, "输入张量数据类型必须为 uint8"
+    
+    # 转换为 float32 类型
+    tensor = tensor.float()
+    
+    # 定义标准化参数（RGB 通道顺序）
+    mean = torch.tensor([123.675, 116.28, 103.53], dtype=torch.float32).view(3, 1, 1)
+    std = torch.tensor([58.395, 57.12, 57.375], dtype=torch.float32).view(3, 1, 1)
+    
+    # 执行标准化：(x - mean) / std
+    normalized_tensor = (tensor - mean) / std
+    
+    return normalized_tensor
+
+
+infer = MMOCRInferencer(rec='satrn')
+model = infer.textrec_inferencer.model
+model.eval()
+model.cpu()
+input_path = 'demo_text_recog.jpg'
+ori_inputs = infer._inputs_to_list([input_path])
+base = BaseMMOCRInferencer(model='satrn')
+chunked_inputs = base._get_chunk_data(ori_inputs, 1)
+for ori_inputs in track(chunked_inputs, description='Inference'):
+    input = ori_inputs[0][1]
+    input_img = input['inputs']
+    input_image = normalize_tensor(input_img).unsqueeze(0)
+    input_sample = input['data_samples']
+    
+    # backbone+encoder
+    model_backbone_encoder = BackboneEncoderOnly(model)
+    model_decoder = DecoderOnly(model)
+    
+    # out_enc = model_backbone_encoder(input_image)
+    out_enc = onnx_bb_encoder.run(None, {"input": np.array(input_image.cpu())})[0]
+    out_enc = torch.tensor(out_enc)
+    data_samples = None
+    
+    N = out_enc.size(0)
+    init_target_seq = torch.full((N, model_decoder.max_seq_len + 1),
+                                model_decoder.padding_idx,
+                                device=out_enc.device,
+                                dtype=torch.long)
+# bsz * seq_len
+    init_target_seq[:, 0] = model_decoder.start_idx
+
+    outputs = []
+    for step in range(0, model_decoder.max_seq_len):
+        valid_ratios = [1.0 for _ in range(out_enc.size(0))]
+        if data_samples is not None:
+            valid_ratios = []
+            for data_sample in data_samples:
+                valid_ratios.append(data_sample.get('valid_ratio'))
+        
+        src_mask = model_decoder._get_source_mask(out_enc, valid_ratios)
+        # step_result = model_decoder(init_target_seq,out_enc,src_mask,step)
+        step_result = onnx_decoder.run(None,{'init_target_seq':np.array(init_target_seq),
+                                             'out_enc':np.array(out_enc),
+                                             'src_mask':np.array(src_mask),
+                                             'step':np.array([step])})[0][0]
+        step_result = torch.tensor(step_result)
+        outputs.append(step_result)
+        _, step_max_index = torch.max(step_result, dim=-1)
+        init_target_seq[:, step + 1] = step_max_index
+    outputs = torch.stack(outputs, dim=1)
+    out_dec = model_decoder.softmax(outputs)
+    output = model_decoder.postprocessor(out_dec, [input_sample])
+    outstr = output[0].pred_text.item
+    outscore = output[0].pred_text.score
+    
+    print('pred_text:',outstr)
+    print('score:',outscore)
+    
+    
+    
+    
+