工业图像序列识别实战：基于PyTorch的OCR模型训练与优化

工业图像序列识别实战：基于PyTorch的OCR模型训练与优化

在工业视觉场景中，图像序列识别（如生产日期、批号、字符码等）是一项常见需求。本文将分享一个完整的工业OCR训练流程，包括数据预处理、模型设计、CTC损失训练和性能评估。代码已在实际项目中验证，适用于字符集有限的定长/不定长文本识别任务。

---

一、项目背景与目标

工业现场常需要识别钢印、喷码、激光刻印等字符。由于字符数量少（如0-9、字母组合），但存在倾斜、光照不均、背景干扰等问题，传统OCR引擎往往难以适应。本文构建了一个轻量级CNN+BiGRU的序列识别模型，使用CTC损失，支持变长文本识别，并通过数据增强提升鲁棒性。

字符集 ："013"（实际使用时可根据业务扩充，只需修改CHARSET即可）

---

二、数据集结构

假设数据集目录如下：

../captchas_gray_o/
├── labels.csv          # 标签文件，每行：文件名,标签
└── *.jpg               # 图像文件

labels.csv格式示例：

img001.jpg,013
img002.jpg,10

图像为彩色图（后续转为灰度），已预先旋转到合适方向，但代码中仍包含了一次cv2.rotate，可根据实际调整。

---

三、代码模块详解

1. 自定义数据增强类 RandomResizeCrop

class RandomResizeCrop:
    def __init__(self, h=128, w=512, scale=(0.8, 1.2)):
        self.h = h
        self.w = w
        self.scale = scale

    def __call__(self, img):
        s = random.uniform(*self.scale)
        new_h = int(self.h * s)
        new_w = int(self.w * s)
        img = F.resize(img, (new_h, new_w))
        if new_h < self.h or new_w < self.w:
            pad_h = max(self.h - new_h, 0)
            pad_w = max(self.w - new_w, 0)
            img = F.pad(img, (0,0,pad_w,pad_h))
            new_h, new_w = img.size[1], img.size[0]
        i = random.randint(0, new_h - self.h)
        j = random.randint(0, new_w - self.w)
        img = F.crop(img, i, j, self.h, self.w)
        return img

作用 ：先随机缩放图像尺寸（比例0.8~1.2），再随机裁剪到固定大小(128,512)。这种策略模拟了不同距离和视角的拍摄效果，增强模型对尺度变化的适应能力。

2. 数据集加载 CaptchaDataset

class CaptchaDataset(Dataset):
    def __init__(self, img_dir, csv_file, transform):
        df = pd.read_csv(csv_file, header=None)
        self.samples = [(img_dir/row[0], str(row[1])) for _, row in df.iterrows()]
        self.transform = transform

    def __getitem__(self, i):
        p, t = self.samples[i]
        img = cv2.imread(p)
        img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)  # 根据实际旋转
        img = resize_image_by_height_opencv(img, 128)           # 高度固定128，等比例缩放宽度
        height, width, _ = img.shape
        left_padding = np.random.randint(0, 512 - width)
        right_padding = 512 - width - left_padding
        img = cv2.copyMakeBorder(img, 0, 0, left_padding, right_padding, 
                                 cv2.BORDER_CONSTANT, value=(127,127,127))
        img = Image.fromarray(img).convert("L")                 # 转灰度
        img = self.transform(img)
        return img, t

关键点：

• 使用OpenCV读取图像，便于高效处理。
• resize_image_by_height_opencv：保持宽高比，固定高度为128。
• 随机左右填充至宽度512，使用中灰色(127)填充，模拟实际图像中可能出现的空白背景。
• 最后转为PIL图像并应用数据增强（随机裁剪、颜色抖动、仿射变换、高斯模糊等），然后转为Tensor。

3. 字符编码与CTC解码

CHARSET = "013"
BLANK = len(CHARSET)          # 空白符号索引
NUM_CLASSES = len(CHARSET)+1

def text_to_labels(text):
    m = {c:i for i,c in enumerate(CHARSET)}
    return torch.tensor([m[c] for c in text], dtype=torch.long)

def ctc_decode(log_probs):
    best = log_probs.argmax(2)            # (T, B)
    T, B = best.shape
    out = []
    for b in range(B):
        prev = None
        s = []
        for i in best[:, b]:
            i = i.item()
            if i == prev:
                continue
            if i != BLANK:
                s.append(CHARSET[i])
            prev = i
        out.append("".join(s))
    return out

• 字符集直接硬编码，空白索引为len(CHARSET)。
• ctc_decode实现贪心解码，合并重复字符并去除空白。

4. 模型架构 UltraOCR

class UltraOCR(nn.Module):
    def __init__(self):
        super().__init__()
        self.cnn = nn.Sequential(
            ConvBlock(1,32), nn.MaxPool2d(2,2),
            ConvBlock(32,64), nn.MaxPool2d(2,2),
            ConvBlock(64,128), nn.MaxPool2d((2,1),(2,1)),
            ConvBlock(128,256), nn.MaxPool2d((2,1),(2,1)),
            ConvBlock(256,256),
        )
        self.rnn = nn.GRU(input_size=256, hidden_size=128, 
                          num_layers=2, bidirectional=True, dropout=0.2)
        self.fc = nn.Linear(256, NUM_CLASSES)
        self.log_softmax = nn.LogSoftmax(2)

    def forward(self, x):
        f = self.cnn(x)                     # (B, C, H, W)
        f = torch.mean(f, dim=2, keepdim=True)  # 沿高度方向平均，消除高度维度 -> (B, C, 1, W)
        b,c,h,w = f.shape
        f = f.squeeze(2)                    # (B, C, W)
        f = f.permute(2,0,1)                # (W, B, C) 符合RNN输入格式
        y, _ = self.rnn(f)                  # (W, B, H*2)
        y = self.fc(y)                      # (W, B, NUM_CLASSES)
        return self.log_softmax(y)

设计思路：

• CNN部分 ：提取空间特征，逐步下采样。特别地，最后两个池化层使用(2,1)核，保持宽度方向特征不丢失过多。
• 高度压缩 ：通过torch.mean(f, dim=2)将高度特征聚合，得到特征向量序列，序列长度等于原图宽度方向的特征点数。
• RNN部分：双向GRU捕捉上下文依赖，输出每个时间步的分类概率。
• 输出：LogSoftmax，配合CTC损失。

5. 训练与验证流程

• 损失函数 ：nn.CTCLoss，空白索引为BLANK，zero_infinity=True避免梯度爆炸。
• 优化器：Adam，学习率1e-4。
• 数据加载 ：使用多进程num_workers=16，pin_memory=True加速数据传输。
• 动态计算输入长度 ：由于网络有4次下采样（2,2,2,2）但其中两次为(2,1)，实际宽度缩减因子为22 1*1=4？需核对：两个(2,2)池化->宽/4，两个(2,1)池化->宽不变，所以最终特征图宽度为原图宽度的1/4。因此T = imgs.shape[-1] // 4。

6. 数据增强配置

transform = transforms.Compose([
    RandomResizeCrop(128,512,(0.85,1.15)),
    transforms.RandomApply([transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2)], p=0.5),
    transforms.RandomApply([transforms.RandomAffine(degrees=8, translate=(0.03, 0.08), shear=8)], p=0.5),
    transforms.RandomApply([transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 0.5))], p=0.3),
    transforms.ToTensor(),
])

• RandomResizeCrop：随机缩放裁剪。
• ColorJitter：模拟光照变化。
• RandomAffine：模拟拍摄角度偏移。
• GaussianBlur：模拟图像模糊。

---

四、训练技巧与注意事项

1. 数据集划分：随机取10%作为验证集。
2. 模型保存：保存验证准确率最高时的模型，同时每epoch保存一次。
3. 设备 ：代码中指定CUDA_VISIBLE_DEVICES=1，可根据需要修改。
4. 预处理流水线 ：在__getitem__中完成固定高度缩放+随机填充，而后应用数据增强，确保每张图独立随机。
5. CTC输入长度 ：务必确保input_lengths与模型输出的时间步数一致，且target_lengths为每个标签的真实长度。

---

五、实验结果与调优建议

通过训练200个epoch，模型在验证集上达到较高准确率（具体数值根据实际数据而定）。以下是几点调优经验：

• 字符集扩充：当字符集增多时，可适当增加RNN隐层维度或添加更多卷积层。
• 图像归一化 ：当前只做了ToTensor，未进行均值和标准差归一化，若数据分布变化大可考虑添加transforms.Normalize。
• 平衡正负样本：如果某些字符出现频率极低，可考虑加权CTC损失或增加该类样本。
• 模型压缩：工业部署时，可对模型进行量化或剪枝。

---

六、总结

本文完整展示了一个工业级OCR模型的训练流程，涵盖了数据处理、数据增强、模型构建、CTC训练等核心环节。代码结构清晰，易于扩展，可直接用于生产环境中的字符识别任务。希望此分享能为从事工业视觉和OCR识别的开发者提供参考。

完整代码：已在文中展示，可根据实际需求调整字符集、图像尺寸和超参数。

训练代码：

import os
os.environ["CUDA_VISIBLE_DEVICES"]="1"

import math
import torch
import pandas as pd
from pathlib import Path
from PIL import Image
from torch import nn
from torch.utils.data import Dataset,DataLoader
from torchvision import transforms
from tqdm import tqdm
import random
import torchvision.transforms.functional as F
import cv2
import numpy as np


class RandomResizeCrop:

    def __init__(self, h=128, w=512, scale=(0.8, 1.2)):

        self.h = h
        self.w = w
        self.scale = scale

    def __call__(self, img):

        s = random.uniform(*self.scale)

        new_h = int(self.h * s)
        new_w = int(self.w * s)

        img = F.resize(img, (new_h, new_w))

        if new_h < self.h or new_w < self.w:

            pad_h = max(self.h - new_h, 0)
            pad_w = max(self.w - new_w, 0)

            img = F.pad(img, (0,0,pad_w,pad_h))

            new_h, new_w = img.size[1], img.size[0]

        i = random.randint(0, new_h - self.h)
        j = random.randint(0, new_w - self.w)

        img = F.crop(img, i, j, self.h, self.w)

        return img

# =========================================================
# charset
# =========================================================

CHARSET="013"
BLANK=len(CHARSET)
NUM_CLASSES=len(CHARSET)+1


def text_to_labels(text):
    m={c:i for i,c in enumerate(CHARSET)}
    return torch.tensor([m[c] for c in text],dtype=torch.long)


def ctc_decode(log_probs):

    best=log_probs.argmax(2)

    T,B=best.shape

    out=[]

    for b in range(B):

        prev=None
        s=[]

        for i in best[:,b]:

            i=i.item()

            if i==prev:
                continue

            if i!=BLANK:
                s.append(CHARSET[i])

            prev=i

        out.append("".join(s))

    return out


# =========================================================
# Dataset
# =========================================================
def resize_image_by_height_opencv(image, target_height, output_path=None):
    """
    使用 OpenCV 将图片按等比例缩放到指定高度。

    参数:
        image (str or numpy.ndarray): 图片文件路径或 OpenCV 图像（BGR 格式的 numpy 数组）。
        target_height (int): 目标高度（像素），必须为正整数。
        output_path (str, optional): 如果提供，将缩放后的图片保存到此路径。

    返回:
        numpy.ndarray: 缩放后的图像（BGR 格式）。

    异常:
        ValueError: 如果 target_height 不是正整数。
        FileNotFoundError: 如果 image 是路径但文件不存在。
        ValueError: 如果 OpenCV 无法读取图片。
    """
    # 参数检查
    if not isinstance(target_height, int) or target_height <= 0:
        raise ValueError("target_height 必须是正整数")

    # 加载图片
    if isinstance(image, str):
        if not os.path.exists(image):
            raise FileNotFoundError(f"图片文件不存在: {image}")
        img = cv2.imread(image)
        if img is None:
            raise ValueError(f"无法读取图片文件: {image}")
    else:
        img = image

    # 获取原始尺寸
    original_height, original_width = img.shape[:2]

    # 计算缩放比例和新宽度
    ratio = target_height / original_height
    new_width = int(original_width * ratio)

    # 使用 Lanczos 插值进行高质量缩放（对应 PIL 的 LANCZOS）
    resized_img = cv2.resize(img, (new_width, target_height), interpolation=cv2.INTER_LANCZOS4)

    # 可选保存
    if output_path:
        cv2.imwrite(output_path, resized_img)

    return resized_img



class CaptchaDataset(Dataset):

    def __init__(self,img_dir,csv_file,transform):

        df=pd.read_csv(csv_file,header=None)

        self.samples=[
            (img_dir/row[0],str(row[1]))
            for _,row in df.iterrows()
        ]

        self.transform=transform

    def __len__(self):
        return len(self.samples)

    def __getitem__(self,i):

        p,t=self.samples[i] 
        img = cv2.imread(p)
        img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
        img = resize_image_by_height_opencv(img, 128)
        height, width, _ = img.shape
        left_padding = np.random.randint(0, 512 - width)
        right_padding = 512 - width - left_padding
        img = cv2.copyMakeBorder(img, 0, 0, left_padding, right_padding, cv2.BORDER_CONSTANT, value=(127, 127, 127))
        #img=Image.open(p).convert("L")
        img = Image.fromarray(img).convert("L")
        # img.save("123.jpg")
        # exit()
        img=self.transform(img)

        return img,t


def collate(batch):

    imgs,txts=zip(*batch)

    imgs=torch.stack(imgs)

    return imgs,txts


# =========================================================
# UltraOCR-V6 Model
# =========================================================

class ConvBlock(nn.Module):

    def __init__(self,in_c,out_c):

        super().__init__()

        self.net=nn.Sequential(
            nn.Conv2d(in_c,out_c,3,1,1,bias=False),
            nn.BatchNorm2d(out_c),
            nn.ReLU(inplace=True)
        )

    def forward(self,x):
        return self.net(x)


class UltraOCR(nn.Module):

    def __init__(self):

        super().__init__()

        self.cnn=nn.Sequential(

            ConvBlock(1,32),
            nn.MaxPool2d(2,2),

            ConvBlock(32,64),
            nn.MaxPool2d(2,2),

            ConvBlock(64,128),
            nn.MaxPool2d((2,1),(2,1)),

            ConvBlock(128,256),
            nn.MaxPool2d((2,1),(2,1)),

            ConvBlock(256,256),

            # nn.AdaptiveAvgPool2d((1,None))
        )

        self.rnn=nn.GRU(
            input_size=256,
            hidden_size=128,
            num_layers=2,
            bidirectional=True,
            dropout=0.2
        )

        self.fc=nn.Linear(256,NUM_CLASSES)

        self.log_softmax=nn.LogSoftmax(2)

    def forward(self,x):

        f=self.cnn(x)

        f = torch.mean(f, dim=2, keepdim=True)

        b,c,h,w=f.shape

        f=f.squeeze(2)

        f=f.permute(2,0,1)

        y,_=self.rnn(f)

        y=self.fc(y)

        return self.log_softmax(y)


# =========================================================
# prepare targets
# =========================================================

def prepare_targets(texts,device):

    labels=[text_to_labels(t) for t in texts]

    targets=torch.cat(labels).to(device)

    lengths=torch.tensor([len(x) for x in labels],dtype=torch.long).to(device)

    return targets,lengths


# =========================================================
# train
# =========================================================

def train_epoch(model,loader,optimizer,criterion,device):

    model.train()

    total_loss=0
    correct=0
    total=0

    for imgs,txts in tqdm(loader):

        imgs=imgs.to(device)

        targets,lengths=prepare_targets(txts,device)

        T=imgs.shape[-1]//4

        input_lengths=torch.full(
            (imgs.size(0),),
            T,
            dtype=torch.long,
            device=device
        )

        optimizer.zero_grad()

        log_probs=model(imgs)

        loss=criterion(log_probs,targets,input_lengths,lengths)

        loss.backward()

        optimizer.step()

        total_loss+=loss.item()*imgs.size(0)

        preds=ctc_decode(log_probs.detach())

        for p,t in zip(preds,txts):
            if p==t:
                correct+=1

        total+=imgs.size(0)

    return total_loss/total,correct/total


# =========================================================
# eval
# =========================================================

def evaluate(model,loader,criterion,device):

    model.eval()

    total_loss=0
    correct=0
    total=0

    with torch.no_grad():

        for imgs,txts in loader:

            # imgs=imgs.to(device)
            imgs = imgs.to(device, non_blocking=True)
            targets,lengths=prepare_targets(txts,device)

            T=imgs.shape[-1]//4

            input_lengths=torch.full(
                (imgs.size(0),),
                T,
                dtype=torch.long,
                device=device
            )

            log_probs=model(imgs)

            loss=criterion(log_probs,targets,input_lengths,lengths)

            total_loss+=loss.item()*imgs.size(0)

            preds=ctc_decode(log_probs)

            for p,t in zip(preds,txts):
                if p==t:
                    correct+=1

            total+=imgs.size(0)

    return total_loss/total,correct/total


# =========================================================
# main
# =========================================================

def main():

    device="cuda" if torch.cuda.is_available() else "cpu"

    img_dir=Path("../captchas_gray_o")
    csv_file=img_dir/"labels.csv"

    transform=transforms.Compose([
        RandomResizeCrop(128,512,(0.85,1.15)),
        transforms.RandomApply([
                transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2)
            ], p=0.5),
        transforms.RandomApply([
                transforms.RandomAffine(degrees=8, translate=(0.03, 0.08), shear=8)
            ], p=0.5),
        transforms.RandomApply([
                transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 0.5))
            ], p=0.3), 

        transforms.ToTensor(),
    ])
    

    

    ds=CaptchaDataset(img_dir,csv_file,transform)

    n=len(ds)

    val=int(n*0.1)

    train_ds,val_ds=torch.utils.data.random_split(ds,[n-val,val])

    train_loader=DataLoader(train_ds,8*32,shuffle=True,collate_fn=collate,num_workers=16, pin_memory=True, persistent_workers=True)
    val_loader=DataLoader(val_ds,8*32,collate_fn=collate, num_workers=16, pin_memory=True, persistent_workers=True)

    model=UltraOCR().to(device)
    model.load_state_dict(torch.load("n_ultraocr_v6_best.pt", map_location=torch.device('cuda')))

    criterion=nn.CTCLoss(blank=BLANK,zero_infinity=True)

    optimizer=torch.optim.Adam(model.parameters(),1e-4)

    best=0

    for epoch in range(200):

        print("epoch",epoch)

        tr_loss,tr_acc=train_epoch(model,train_loader,optimizer,criterion,device)

        val_loss,val_acc=evaluate(model,val_loader,criterion,device)

        print(tr_loss,tr_acc,val_loss,val_acc)

        if val_acc>best:

            best=val_acc

            torch.save(model.state_dict(),"n_ultraocr_v6_best.pt")

        torch.save(model.state_dict(),f"{epoch}_ultraocr_v6_last.pt")


if __name__=="__main__":
    main()

转onnx

import torch
import torch.nn as nn
import torch.nn.functional as F
from train_gray_10_1_1 import UltraOCR

CHARSET = "013"
BLANK_INDEX = len(CHARSET)


# -------------------------
# Preprocess
# -------------------------
class Preprocess(nn.Module):

    def __init__(self):
        super().__init__()

    def rot90_ccw(self, x):
        # rot90替代实现（ONNX支持）
        x = x.transpose(2, 3)
        x = torch.flip(x, [2])
        return x

    def forward(self, x):
        

        # x (B,H,W,3) uint8
        x = x.float() / 255.0
        # x = x.float()

        # B,H,W,C -> B,C,H,W
        x = x.permute(0, 3, 1, 2)

        # RGB->GRAY
        # r, g, b = x[:,0:1], x[:,1:2], x[:,2:3]
        # x = 0.299*r + 0.587*g + 0.114*b

        # rotate 90 CCW
        x = self.rot90_ccw(x)

        B,C,H,W = x.shape

        # resize height ->128
        scale = 128.0 / H
        new_w = (W * scale).long()

        x = F.interpolate(x, size=(128, new_w), mode="bicubic", align_corners=False)

        # pad / crop width ->512
        target_w = 512

        if new_w < target_w:

            pad = target_w - new_w
            left = pad // 2
            right = pad - left
            x = F.pad(x, (left,right,0,0), value=127.0 / 255.0)
            # x = F.pad(x, (0,pad,0,0), value=127.0 / 255.0)

        else:

            x = x[:,:,:,:target_w]

        return x


# -------------------------
# CTC Greedy Decode
# -------------------------
class GreedyCTC(nn.Module):

    def __init__(self):
        super().__init__()
        self.blank_index = BLANK_INDEX

    def forward(self, log_probs):

        # (T,B,C) -> (T,B)
        best = torch.argmax(log_probs, dim=2)

        # (T,B) -> (B,T)
        best = best.permute(1,0)

        # ----- collapse repeat -----
        prev = F.pad(best[:, :-1], (1,0), value=-1)

        keep = best != prev

        best = torch.where(keep, best, torch.full_like(best, -1))

        # ----- remove blank -----
        keep_blank = best != self.blank_index

        best = torch.where(keep_blank, best, torch.full_like(best, -1))
        
        best = best[best != -1]

        return best


# -------------------------
# End2End
# -------------------------
class OCR_End2End(nn.Module):

    def __init__(self, crnn):
        super().__init__()

        self.pre = Preprocess()
        self.crnn = crnn
        self.ctc = GreedyCTC()

    def forward(self, x):

        x = self.pre(x)

        log_probs = self.crnn(x)

        pred = self.ctc(log_probs)

        return pred
    

device = "cpu"

model=UltraOCR()

model.load_state_dict(torch.load("ultraocr_v6_best.pt", map_location=torch.device('cpu')))

model.eval()

end2end = OCR_End2End(model).eval()

dummy = torch.randint(
    0,255,(1,128,512,1),dtype=torch.uint8
)

torch.onnx.export(
    end2end,
    dummy,
    "bo_pian_det.onnx",
    input_names=["image"],
    output_names=["pred"],
    dynamic_axes={
        "image":{
            0:"batch",
            1:"height",
            2:"width"
        },
        "pred":{
            0:"batch"
        }
    },
    opset_version=18
)

print("export success")

onnx推理

import onnxruntime as ort
import numpy as np
import cv2

# 分类字典
INDEX_CLASS = {0: "NORMAL", 1: "ABNORMAL", 2: "NO"}
INDEX_CLASS_CN = {0: "0", 1: "1", 2: "3"}

class BopianDetector:
    def __init__(self, model_path: str, providers=None):
        """
        初始化 ONNX 推理会话
        """
        if providers is None:
            providers = ["CPUExecutionProvider"]
        self.session = ort.InferenceSession(model_path, providers=providers)

    def predict(self, img: np.ndarray):
        """
        对输入灰度图像进行 OCR 推理
        """
        # 1. 灰度图
        if len(img.shape) == 3 and img.shape[2] == 3:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        # 2. 添加 channel 和 batch 维度
        input_img = img[None, ..., None].astype(np.uint8)

        # 3. ONNX 推理
        pred = self.session.run(None, {"image": input_img})[0]  # shape (seq_len,)

        # 4. 构造结果
        result = {}
        text = ""
        for idx, p in enumerate(pred):
            if p >= 0:
                result[idx] = INDEX_CLASS[p]
                text += INDEX_CLASS_CN[p]

        return result, text


if __name__ == "__main__":
    detector = BopianDetector("n_det.onnx")

    
    lines = open(r"labels.txt", "r").readlines()

    corrext_num = 0

    for line in lines:
        img_path, label = line.strip().split(",")
        img = cv2.imread("D:/imgs/bopian/imgs/"+img_path)
        result_dict, text = detector.predict(img)
    
        if text != label:
            # print(res, len(res), label)
            print(img_path)
            # print(img_path)
            print("预测：",text, len(text))
            print("标注：", label, len(label))
        else:
            corrext_num += 1

    print(f"正确率：{corrext_num/len(lines):.4f}")
        # break