源代码
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image, ImageDraw, ImageFont, ImageFilter
import os
import random
import numpy as np# -------------------------- 1. 关键参数配置（可按需修改） --------------------------
# 目标识别图片路径
TARGET_IMG_PATH = r"C:\Users\黄楚玉\Desktop\杂七杂八\文字识别图.jpg"
# 自定义数据集路径（自动生成）
DATA_ROOT = r"C:\Users\黄楚玉\Desktop\文字识别数据集"
IMG_DIR = os.path.join(DATA_ROOT, "images")
LABEL_PATH = os.path.join(DATA_ROOT, "labels.txt")
# 文字库（识别范围：数字+大小写字母，共62类）
TEXT_CHARS = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
NUM_CLASSES = len(TEXT_CHARS) # 类别数：62
# 训练参数
BATCH_SIZE = 32
EPOCHS = 30
LEARNING_RATE = 0.001
# 图片尺寸（与CNN输入匹配）
IMG_WIDTH = 128
IMG_HEIGHT = 32# -------------------------- 2. 步骤1：自动生成自定义文字数据集（含标注） --------------------------
def generate_custom_dataset(total_samples=2000, test_ratio=0.2):"""生成训练/测试数据集，自动标注标签"""# 创建文件夹os.makedirs(IMG_DIR, exist_ok=True)# 分割训练/测试样本数train_samples = int(total_samples * (1 - test_ratio))test_samples = total_samples - train_samplesall_samples = []# 生成单字符图片（模拟真实文字）def draw_text_img(text):# 创建白色背景图img = Image.new("RGB", (IMG_WIDTH, IMG_HEIGHT), (255, 255, 255))draw = ImageDraw.Draw(img)# 加载字体（Windows默认黑体，支持中文；其他系统需修改路径）font = ImageFont.truetype(r"C:\Windows\Fonts\simhei.ttf", 22)# 使用 ImageDraw 的 textbbox 方法获取文本边界框，从而得到文本尺寸bbox = draw.textbbox((0, 0), text, font=font)text_w = bbox[2] - bbox[0]text_h = bbox[3] - bbox[1]# 文字位置（垂直居中，水平左移）x = 5y = (IMG_HEIGHT - text_h) // 2# 随机深色文字（与背景区分）text_color = (random.randint(0, 80), random.randint(0, 80), random.randint(0, 80))draw.text((x, y), text, fill=text_color, font=font)# 随机加噪声/模糊（增强鲁棒性）if random.random() < 0.3: # 30%概率加噪声img_np = np.array(img)img_np = img_np + np.random.normal(0, 15, img_np.shape)img = Image.fromarray(np.clip(img_np, 0, 255).astype(np.uint8))if random.random() < 0.2: # 20%概率加模糊，使用PIL的ImageFilterimg = img.filter(ImageFilter.GaussianBlur(radius=1))return img.convert("L") # 转为灰度图# 批量生成图片与标签for i in range(total_samples):# 随机选一个字符作为标签label = random.choice(TEXT_CHARS)# 生成图片img = draw_text_img(label)# 保存图片img_name = f"img_{i}.jpg"img.save(os.path.join(IMG_DIR, img_name))# 记录“图片名-标签”all_samples.append((img_name, label))# 分割训练/测试集并写入标签文件train_data = all_samples[:train_samples]test_data = all_samples[train_samples:]with open(LABEL_PATH, "w", encoding="utf-8") as f:f.write("train\n")for img_name, label in train_data:f.write(f"{img_name} {label}\n")f.write("test\n")for img_name, label in test_data:f.write(f"{img_name} {label}\n")print(f"数据集生成完成！训练集{train_samples}张，测试集{test_samples}张")return train_data, test_data# -------------------------- 3. 步骤2：数据加载与预处理 --------------------------
class TextDataset(Dataset):def __init__(self, img_dir, label_path, data_type="train", transform=None):self.img_dir = img_dirself.transform = transform# 建立标签到数字的映射（如"0"→0，"a"→10）self.label_map = {char: idx for idx, char in enumerate(TEXT_CHARS)}# 读取训练/测试数据self.data = []with open(label_path, "r", encoding="utf-8") as f:lines = f.readlines()# 找到训练/测试集的起始位置start_idx = lines.index("train\n") + 1 if data_type == "train" else lines.index("test\n") + 1# 读取对应数据（直到下一个分割符或文件结束）for line in lines[start_idx:]:line = line.strip()# 跳过空行和分割符行if not line or line in ["train\n", "test\n"]:continue# 分割行内容为图片名和标签parts = line.split()if len(parts) == 2:img_name, label = partsself.data.append((img_name, label))else:print(f"警告：跳过格式不正确的行：{line}")def __len__(self):return len(self.data)def __getitem__(self, idx):img_name, label = self.data[idx]# 读取图片img_path = os.path.join(self.img_dir, img_name)img = Image.open(img_path).convert("L") # 灰度图# 预处理if self.transform:img = self.transform(img)# 标签转为数字张量label = torch.tensor(self.label_map[label], dtype=torch.long)return img, label# 数据预处理流水线（与模型输入匹配）
transform = transforms.Compose([transforms.Resize((IMG_HEIGHT, IMG_WIDTH)), # 统一尺寸transforms.ToTensor(), # 转为[1,32,128]张量（通道数,高,宽）transforms.Normalize((0.5,), (0.5,)) # 灰度图标准化（均值0.5，标准差0.5）
])# 生成数据集
generate_custom_dataset()
# 加载训练/测试集
train_dataset = TextDataset(IMG_DIR, LABEL_PATH, data_type="train", transform=transform)
test_dataset = TextDataset(IMG_DIR, LABEL_PATH, data_type="test", transform=transform)
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)# -------------------------- 4. 步骤3：定义CNN模型 --------------------------
class TextCNN(nn.Module):def __init__(self, num_classes):super(TextCNN, self).__init__()# 卷积层：提取文字笔画/局部特征self.conv_layers = nn.Sequential(nn.Conv2d(1, 32, kernel_size=(3, 3), padding=1), # 输入1通道，输出32通道nn.ReLU(),nn.MaxPool2d(kernel_size=(2, 2)), # 下采样：(32,128)→(16,64)nn.Conv2d(32, 64, kernel_size=(3, 3), padding=1),nn.ReLU(),nn.MaxPool2d(kernel_size=(2, 2)), # 下采样：(16,64)→(8,32)nn.Conv2d(64, 128, kernel_size=(3, 3), padding=1),nn.ReLU(),nn.MaxPool2d(kernel_size=(2, 2)) # 下采样：(8,32)→(4,16))# 全连接层：映射到类别概率self.fc_layers = nn.Sequential(nn.Flatten(), # 展平：[128,4,16]→128*4*16=8192nn.Linear(128 * 4 * 16, 1024), # 8192→1024nn.ReLU(),nn.Dropout(0.5), # 防止过拟合nn.Linear(1024, num_classes) # 1024→62（输出62类概率）)def forward(self, x):x = self.conv_layers(x)x = self.fc_layers(x)return x# -------------------------- 5. 步骤4：初始化模型、损失函数、优化器 --------------------------
# 设备选择（优先GPU，无则用CPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备：{device}")# 初始化组件
model = TextCNN(num_classes=NUM_CLASSES).to(device)
criterion = nn.CrossEntropyLoss() # 多分类损失：交叉熵
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE) # 优化器：Adam# -------------------------- 6. 步骤5：模型训练 --------------------------
def train_model():print("\n开始训练模型...")model.train() # 训练模式（启用Dropout）for epoch in range(EPOCHS):total_loss = 0.0for images, labels in train_loader:# 数据移至设备images, labels = images.to(device), labels.to(device)# 前向传播：计算预测值outputs = model(images)# 计算损失loss = criterion(outputs, labels)# 反向传播+参数更新optimizer.zero_grad() # 清空梯度loss.backward() # 计算梯度optimizer.step() # 更新参数total_loss += loss.item() * images.size(0)# 计算本轮平均损失avg_loss = total_loss / len(train_dataset)print(f"第{epoch+1}/{EPOCHS}轮，训练损失：{avg_loss:.4f}")print("模型训练完成！")# -------------------------- 7. 步骤6：计算测试集正确率 --------------------------
def calculate_accuracy():model.eval() # 评估模式（关闭Dropout）correct = 0total = 0with torch.no_grad(): # 禁用梯度计算，节省内存for images, labels in test_loader:images, labels = images.to(device), labels.to(device)outputs = model(images)# 取概率最大的类别作为预测结果_, predicted = torch.max(outputs.data, 1)total += labels.size(0)correct += (predicted == labels).sum().item()# 计算正确率accuracy = 100 * correct / totalprint(f"\n测试集正确率：{accuracy:.2f}%")return accuracy# -------------------------- 8. 步骤7：识别目标图片（输出文字结果） --------------------------
def recognize_target_img(img_path):model.eval() # 评估模式# 预处理目标图片（与训练数据一致）img = Image.open(img_path).convert("L") # 转为灰度图img = transform(img).unsqueeze(0) # 增加batch维度：[1,1,32,128]img = img.to(device)# 预测with torch.no_grad():outputs = model(img)_, predicted_idx = torch.max(outputs.data, 1)# 从索引映射回文字predicted_text = TEXT_CHARS[predicted_idx.item()]print(f"\n目标图片识别结果：{predicted_text}")return predicted_text# -------------------------- 9. 主流程执行 --------------------------
if __name__ == "__main__":# 1. 训练模型train_model()# 2. 计算测试集正确率calculate_accuracy()# 3. 识别目标图片并输出文字recognize_target_img(TARGET_IMG_PATH)运行结果如下：
