devbox
Gausst松鼠会
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

基于LeNet与YOLOV8在2024年mm杯甲骨文图像文字识别数据集研究_yolov8文本识别

作者：Gausst松鼠会 | 2024-06-15 22:11:44

赞

踩

yolov8文本识别

1、数据展示与处理

  • 数据集下载：https://pan.baidu.com/s/1cp-DUwgS3c_d8EFhBwXRHw?pwd=pf1t
  • 数据选择：下载后的数据集分为4个文件1_Pre_test 、2_Train、3_Test、4_Recognize
    由于本文只讲解如何搭建一个较为完整的图像识别分类模型，所以选取的数据集为2_Train与4_Recognize这二份数据集。
  • 其中2_Train文件夹下有对应的image以及json数据类型。但在json中所展示的数据集如下：
    {"img_name": "b02519Z", "ann": [[64.0, 194.0, 134.0, 277.0, 1.0], [150.0, 168.0, 204.0, 344.0, 1.0], [206.0, 170.0, 259.0, 349.0, 1.0]]}。ann每一行所对应的信息是对应img_name的一个图像框，也即为已经标注好的甲骨文信息。同时在4_Recongnize文件夹下有二个子文件夹，在训练集下有已经分好的图片类别，在测试集下有待识别文字的图片。接下来我们讨论如何处理这些数据集。
  • 在YOLOV8训练数据集中，数据类型应该为YOLO TXT类型，具体可参考这位博主：https://blog.csdn.net/wlh156423/article/details/118861987
    在处理的同时，可以看出在json文件下，并没有为我们识别文字信息为那一类，故而需要采取一类模型对文字进行分类处理，将训练好的模型对上述信息进行分类，最后转为YOLO TXT类型再进行YOLOV8的模型训练，即可得到想要模型。
  • 图像分类模型有多具体可以参考：https://zhuanlan.zhihu.com/p/562609779
    本文采用LeNet模型对所含信息进行分类。

2、具体实现过程

在开始实现以上操作时，我们需要准备python环境，具体流程可以参考：（本文所采用python=3.9的环境）https://blog.csdn.net/weixin_43366149/article/details/132206526

第一步训练LeNet模型

搭建LeNet模型

本文将此代码保存为LeNet.py。

  1. import torch
  2. import torch.nn as nn
  3.  
  4. class LeNet(nn.Module):
  5.     def __init__(self, num_classes=10):
  6.         super(LeNet, self).__init__()
  7.         self.conv1 = nn.Conv2d(1, 6, kernel_size=5)  # 输入通道数为1（灰度图像），输出通道数为6，卷积核大小为5x5
  8.         self.conv2 = nn.Conv2d(6, 16, kernel_size=5)  # 输入通道数为6，输出通道数为16，卷积核大小为5x5
  9.         self.fc1 = nn.Linear(16*5*5, 120)  # 全连接层，输入大小为16*5*5，输出大小为120
  10.         self.fc2 = nn.Linear(120, 84)  # 全连接层，输入大小为120，输出大小为84
  11.         self.fc3 = nn.Linear(84, num_classes)  # 全连接层，输入大小为84，输出大小为类别数
  12.  
  13.     def forward(self, x):
  14.         # 输入数据经过卷积层1，使用ReLU激活函数，再经过2x2最大池化层
  15.         x = torch.max_pool2d(torch.relu(self.conv1(x)), (2, 2))
  16.         # 输入数据经过卷积层2，使用ReLU激活函数，再经过2x2最大池化层
  17.         x = torch.max_pool2d(torch.relu(self.conv2(x)), (2, 2))
  18.         # 将特征图展平成一维向量
  19.         x = x.view(-1, 16*5*5)
  20.         # 经过全连接层1和ReLU激活函数
  21.         x = torch.relu(self.fc1(x))
  22.         # 经过全连接层2和ReLU激活函数
  23.         x = torch.relu(self.fc2(x))
  24.         # 经过全连接层3
  25.         x = self.fc3(x)
  26.         return x
  27.  
  28. # # 创建LeNet模型实例
  29. # model = LeNet(num_classes=10)
  30. # # 输出模型结构
  31. # print(model)

 处理数据集

本文将其命名为Mydata.py。

  1. import os
  2. from PIL import Image
  3. from torch.utils.data import Dataset
  4.  
  5.  
  6. class MyDataset(Dataset):
  7.     def __init__(self, data_dir, transform=None):
  8.         """
  9.         初始化 MyDataset 类。
  11.         Args:
  12.             data_dir (str): 数据目录的路径。
  13.             transform (callable, optional): 对图像进行预处理的转换函数。
  14.         """
  15.         self.data_dir = data_dir
  16.         self.transform = transform
  17.  
  18.         # 获取类别文件夹和类别名称
  19.         self.classes = sorted([d.name for d in os.scandir(data_dir) if d.is_dir()])
  20.         self.class_to_idx = {cls_name: idx for idx, cls_name in enumerate(self.classes)}
  21.  
  22.         # 获取图像文件路径和对应的类别标签
  23.         self.image_paths = []
  24.         self.labels = []
  25.  
  26.         for class_name in self.classes:
  27.             class_dir = os.path.join(data_dir, class_name)
  28.             for img_name in os.listdir(class_dir):
  29.                 img_path = os.path.join(class_dir, img_name)
  30.                 self.image_paths.append(img_path)
  31.                 self.labels.append(self.class_to_idx[class_name])
  32.  
  33.     def __len__(self):
  34.         """
  35.         返回数据集中样本的数量。
  36.         """
  37.         return len(self.image_paths)
  38.  
  39.     def __getitem__(self, index):
  40.         """
  41.         根据索引获取样本（图像和标签）。
  43.         Args:
  44.             index (int): 样本的索引。
  46.         Returns:
  47.             tuple: (image, label) 图像和标签。
  48.         """
  49.         # 获取图像路径和标签
  50.         img_path = self.image_paths[index]
  51.         label = self.labels[index]
  52.  
  53.         # 打开图像并进行预处理
  54.         image = Image.open(img_path)
  55.         if self.transform:
  56.             image = self.transform(image)
  57.  
  58.         return image, label
  59.

搭建训练主程序

本文将其命名为train_lenet.py。

  1. import torch
  2. import torch.nn as nn
  3. import torch.optim as optim
  4. from torch.utils.data import DataLoader
  5. from LeNet import LeNet
  6. from Mydata import MyDataset
  7. from torchvision import transforms
  8. from torch.utils.tensorboard import SummaryWriter
  9.  
  10. # 定义训练参数
  11. batch_size = 64
  12. learning_rate = 0.001
  13. num_epochs = 5
  14.  
  15. # 检查是否可以使用CUDA加速
  16. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  17.  
  18. # 创建数据集实例
  19. data_dir = "4_Recognize/训练集"
  20. transform = transforms.Compose([
  21.     transforms.Resize((32, 32)),
  22.     transforms.Grayscale(num_output_channels=1),  # Convert to grayscale
  23.     transforms.ToTensor(),
  24.     transforms.Normalize((0.5,), (0.5,))
  25. ])
  26.  
  27. dataset = MyDataset(data_dir, transform=transform)
  28. data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
  29.  
  30. # 创建LeNet模型实例
  31. model = LeNet(num_classes=76)
  32. model.to(device)
  33.  
  34. # 定义损失函数和优化器
  35. criterion = nn.CrossEntropyLoss()
  36. optimizer = optim.Adam(model.parameters(), lr=learning_rate)
  37.  
  38. # 创建SummaryWriter实例
  39. writer = SummaryWriter()
  40.  
  41. # 训练模型
  42. total_steps = len(data_loader)
  43. for epoch in range(num_epochs):
  44.     for i, (images, labels) in enumerate(data_loader):
  45.         images = images.to(device)
  46.         labels = labels.to(device)
  47.  
  48.         # 前向传播
  49.         outputs = model(images)
  50.         loss = criterion(outputs, labels)
  51.  
  52.         # 反向传播和优化
  53.         optimizer.zero_grad()
  54.         loss.backward()
  55.         optimizer.step()
  56.  
  57.         if (i + 1) % 100 == 0:
  58.             print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{total_steps}], Loss: {loss.item():.4f}')
  59.  
  60.             # 将损失写入日志
  61.             writer.add_scalar('Loss/train', loss.item(), epoch * total_steps + i)
  62.  
  63. print('Finished Training')
  64.  
  65. # 关闭SummaryWriter
  66. writer.close()
  67.  
  68. # 保存模型
  69. torch.save(model.state_dict(), 'lenet_model.pth')

最后保存所需模型文件即可，本文主要讨论的是方法以及具体实现思路，本文未进行模型评估检验，感兴趣的小伙伴可以外加代码对模型精度进行检验。

第二步处理2_Train文件

JSON转YOLO TXT 

通过参考https://blog.csdn.net/wlh156423/article/details/118861987 ​​​​​​本文对其进行改进，即可得到如下代码实现：

  1. import os
  2. import json
  3. import shutil
  4. from PIL import Image
  5.  
  6. def convert_bbox(img_size, box):
  7.     dw = 1. / img_size[0]
  8.     dh = 1. / img_size[1]
  9.     x = (box[0] + box[2]) / 2.0 - 1
  10.     y = (box[1] + box[3]) / 2.0 - 1
  11.     w = box[2] - box[0]
  12.     h = box[3] - box[1]
  13.     x = x * dw
  14.     w = w * dw
  15.     y = y * dh
  16.     h = h * dh
  17.     return (x, y, w, h)
  18.  
  19. def get_image_size(image_path):
  20.     with Image.open(image_path) as img:
  21.         width, height = img.size
  22.     return width, height
  23.  
  24. def extract_data_from_json(json_folder, train_folder):
  25.     images_folder = os.path.join(train_folder, 'images')
  26.     labels_folder = os.path.join(train_folder, 'labels')
  27.     os.makedirs(images_folder, exist_ok=True)
  28.     os.makedirs(labels_folder, exist_ok=True)
  29.  
  30.     json_files = [f for f in os.listdir(json_folder) if f.endswith('.json')]
  31.  
  32.     for json_file in json_files:
  33.         with open(os.path.join(json_folder, json_file)) as f:
  34.             data = json.load(f)
  35.             img_name = data.get('img_name')
  36.             ann = data.get('ann')
  37.  
  38.             img_path = os.path.join('2_Train', img_name + '.jpg')
  39.             img_w, img_h = get_image_size(img_path)
  40.             shutil.copy(img_path, os.path.join(images_folder, img_name + '.jpg'))
  41.  
  42.             txt_path = os.path.join(labels_folder, img_name + '.txt')
  43.             with open(txt_path, 'w') as txt_file:
  44.                 for bbox in ann:
  45.                     bbox = convert_bbox((img_w, img_h), bbox[:-1])  # Ignore the class label for now
  46.                     txt_file.write(f"0 {' '.join(map(str, bbox))}\n")  # Assuming class index is 0
  47.  
  48. if __name__ == "__main__":
  49.     json_folder_path = '2_Train'
  50.     train_folder_path = 'train'
  51.     extract_data_from_json(json_folder_path, train_folder_path)

将JSON文件转化为所需要的YOLO TXT文件，且保存再同项目train文件夹下。其中txt数据类如下

0 0.3091482649842271 0.5269662921348315 0.22082018927444794 0.1865168539325843
0 0.555205047318612 0.5730337078651686 0.17034700315457413 0.39550561797752815
0 0.7302839116719243 0.5808988764044944 0.167192429022082 0.402247191011236

将对应的信息框进行分类

  1. import torch
  2. import torchvision.transforms as transforms
  3. from PIL import Image
  4. from LeNet import LeNet
  5. import os
  6.  
  7. # 定义预处理操作
  8. transform = transforms.Compose([
  9.     transforms.Grayscale(),  # 转为灰度图像
  10.     transforms.Resize((32, 32)),  # 改变尺寸以适配LeNet模型输入
  11.     transforms.ToTensor(),  # 转为张量
  12. ])
  13.  
  14. # 初始化模型
  15. model = LeNet(num_classes=76)  
  16. model.load_state_dict(torch.load("lenet_model.pth"))
  17. model.eval()  # 设置为评估模式
  18.  
  19. # 文件和文件夹路径
  20. label_dir = 'B题/labels/train'
  21. image_dir = 'B题/images/train'
  22. output_dir = 'new_data'
  23.  
  24. if not os.path.exists(output_dir):
  25.     os.makedirs(output_dir)
  26.  
  27. # 遍历标签文件夹
  28. for label_file in os.listdir(label_dir):
  29.     label_path = os.path.join(label_dir, label_file)
  30.     image_path = os.path.join(image_dir, label_file.replace('.txt', '.jpg'))  # 假设图片文件名与标签相同，仅扩展名不同
  31.  
  32.     if os.path.exists(image_path):
  33.         img = Image.open(image_path)  # 加载图像
  34.         with open(label_path, 'r') as file:
  35.             lines = file.readlines()
  36.  
  37.         new_lines = []
  38.         for line in lines:
  39.             parts = line.strip().split()
  40.             # 假设标签格式为：class x_center y_center width height
  41.             _, x_center, y_center, width, height = map(float, parts)
  42.             bbox = [
  43.                 (x_center - width / 2) * img.width,
  44.                 (y_center - height / 2) * img.height,
  45.                 width * img.width,
  46.                 height * img.height
  47.             ]  # 计算实际像素坐标和尺寸
  48.  
  49.             # 裁剪并转换图像
  50.             cropped_img = img.crop((bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]))
  51.             input_tensor = transform(cropped_img).unsqueeze(0)
  52.  
  53.             # 模型预测
  54.             with torch.no_grad():
  55.                 output = model(input_tensor)
  56.                 predicted_class = output.argmax(1).item()
  57.  
  58.             # 更新行
  59.             new_line = f"{predicted_class} {x_center} {y_center} {width} {height}\n"
  60.             new_lines.append(new_line)
  61.  
  62.         # 保存新的标签文件
  63.         new_label_path = os.path.join(output_dir, label_file)
  64.         with open(new_label_path, 'w') as new_file:
  65.             new_file.writelines(new_lines)

这段代码会将原本的标签进行处理并且以txt形式保存在new_data文件夹下，小伙伴们运行时记得更改路径为自己的路径。转化后的图片txt信息如下

59 0.3091482649842271 0.5269662921348315 0.22082018927444794 0.1865168539325843
69 0.555205047318612 0.5730337078651686 0.17034700315457413 0.39550561797752815
6 0.7302839116719243 0.5808988764044944 0.167192429022082 0.402247191011236

提取类别标签

基于yolov8训练数据集时需要类别标签，本文运用如下代码进行相关处理

  1. import os
  2.  
  3. # 指定文件夹路径
  4. folder_path = '4_Recognize/训练集'
  5.  
  6. # 获取子文件夹名称并按照指定格式保存
  7. with open('name.txt', 'w') as file:
  8.     for idx, subfolder in enumerate(sorted(os.listdir(folder_path))):
  9.         file.write(f"{idx}: {subfolder}\n")

第三步进行YOLOV8模型训练

训练操作可以参考https://blog.csdn.net/qq_42452134/article/details/135149531

以下是实现过程

建立一个python文件且输入如下代码

  1. from ultralytics import YOLO
  2.  
  3. # 加载模型
  4. model = YOLO('yolov8n.yaml').load('yolov8n.pt')  # 从YAML构建并转移权重
  5.  
  6. if __name__ == '__main__':
  7.     # 训练模型
  8.     results = model.train(data='fnal.yaml', epochs=10, imgsz=256)
  9.  
  10.     metrics = model.val()

建立对应的.yaml文件，这里就需要用到我们提取的类别标签。（本文命名为fnal.yaml）

  1. # 这里的path需要指向你项目中数据集的目录
  2. path: E:/dell/比赛/new/train/
  3. # 这里分别指向你训练、验证、测试的文件地址，只需要指向图片的文件夹即可。但是要注意图片和labels名称要对应
  4. train: images/train  # train images (relative to 'path') 128 images
  5. val: images/val  # val images (relative to 'path') 128 images
  6. test: images/test # test images (optional)
  7.  
  8. # Classes
  9. names:
  10.   0: 万
  11.   1: 丘
  12.   2: 丙
  13.   3: 丧
  14.   4: 乘
  15.   5: 亦
  16.   6: 人
  17.   7: 今
  18.   8: 介
  19.   9: 从
  20.   10: 令
  21.   11: 以
  22.   12: 伊
  23.   13: 何
  24.   14: 余
  25.   15: 允
  26.   16: 元
  27.   17: 兄
  28.   18: 光
  29.   19: 兔
  30.   20: 入
  31.   21: 凤
  32.   22: 化
  33.   23: 北
  34.   24: 印
  35.   25: 及
  36.   26: 取
  37.   27: 口
  38.   28: 吉
  39.   29: 囚
  40.   30: 夫
  41.   31: 央
  42.   32: 宗
  43.   33: 宾
  44.   34: 尞
  45.   35: 巳
  46.   36: 帽
  47.   37: 并
  48.   38: 彘
  49.   39: 往
  50.   40: 御
  51.   41: 微
  52.   42: 旨
  53.   43: 昃
  54.   44: 木
  55.   45: 朿
  56.   46: 涎
  57.   47: 灾
  58.   48: 焦
  59.   49: 爽
  60.   50: 牝
  61.   51: 牡
  62.   52: 牧
  63.   53: 生
  64.   54: 田
  65.   55: 疑
  66.   56: 祝
  67.   57: 福
  68.   58: 立
  69.   59: 羊
  70.   60: 羌
  71.   61: 翌
  72.   62: 翼
  73.   63: 老
  74.   64: 艰
  75.   65: 艺
  76.   66: 若
  77.   67: 莫
  78.   68: 获
  79.   69: 衣
  80.   70: 逆
  81.   71: 门
  82.   72: 降
  83.   73: 陟
  84.   74: 雍
  85.   75: 鹿

同时要注意文件夹类型的指向，文件夹的命名也很关键，具体操作可以看https://blog.csdn.net/qq_42452134/article/details/135181244这位博主所讲内容。

这样我们就可以开始训练了

3、所得结果

训练的图片展示

当我们进行训练后，可以在项目路径下的runs文件中观看所得参数图像，以及python运行框中查看模型是否训练完全。

预测

再训练好模型后如何进行预测呢？可以参考官方文档https://docs.ultralytics.com/modes/predict/#key-features-of-predict-mode

本文给出相应的二份预测处理的代码

  1. from PIL import Image
  2. from ultralytics import YOLO
  3.  
  4. # Load a pretrained YOLOv8n model
  5. model = YOLO('yolov8n.pt')
  6.  
  7. # Run inference on 'bus.jpg'
  8. results = model(['4_Recognize/测试集/w01790.jpg', '4_Recognize/测试集/w01791.jpg'])  # results list
  9.  
  10. # Visualize the results
  11. for i, r in enumerate(results):
  12.     # Plot results image
  13.     im_bgr = r.plot()  # BGR-order numpy array
  14.     im_rgb = Image.fromarray(im_bgr[..., ::-1])  # RGB-order PIL image
  15.  
  16.     # Show results to screen (in supported environments)
  17.     r.show()
  18.  
  19.     # Save results to disk
  20.     r.save(filename=f'results{i}.jpg')
  1. # -*- coding: gbk -*-
  2. import os
  3. from PIL import Image
  4. from tqdm import tqdm
  5. from glob import glob
  6. from ultralytics import YOLO
  7.  
  8. # 定义测试集图片文件夹和保存预测图片的文件夹
  9. test_images_folder = "4_Recognize/测试集"
  10. output_folder = "new_images"
  11.  
  12. # 加载训练好的YOLOv8模型
  13. model = YOLO("yolov8n.pt")  # 替换为你训练好的模型的路径
  14.  
  15. # 确保保存预测图片的文件夹存在
  16. os.makedirs(output_folder, exist_ok=True)
  17.  
  18. # 获取测试集中所有图片的路径
  19. test_image_paths = glob(os.path.join(test_images_folder, "*.jpg")) + glob(os.path.join(test_images_folder, "*.png"))
  20.  
  21. # 对每张图片进行预测并保存预测后的图片
  22. # 对每张图片进行预测并保存预测后的图片
  23. for image_path in tqdm(test_image_paths, desc="Predicting"):
  24.     # 读取图片
  25.     image = Image.open(image_path)
  26.  
  27.     # 对图片进行预测
  28.     predicted_images = model.predict(image)
  29.  
  30.     # 处理每个预测结果并保存到文件夹中
  31.     for idx, predicted_image in enumerate(predicted_images):
  32.         # 构建保存预测图片的路径，这里假设以预测结果的索引命名文件
  33.         output_path = os.path.join(output_folder, f"{os.path.basename(image_path)}_{idx}.jpg")
  34.  
  35.         # 保存预测后的图片
  36.         predicted_image.save(output_path)
  37.  
  38. print("预测完成，并且预测图片已保存到'new_images'文件夹中。")

4、结语

本文只是提供一种参考思路且对对应数据集进行了处理。在进行相关的图像识别中，最让人头疼的应该是数据集，例如本次比赛的数据集处理起来较为复杂，如果不熟悉YOLO数据类型，以及对应关系转化，难以将题目进行。

同时作者参与此次比赛收获颇多，虽然是赛后补救，但也希望提供给大家一种思路，以及运用YOLOV8的细节。

  1. YOLOV8运行需要用所需的YOLO TXT类型文件，（即txt文件名为图像文件名，且txt文件中装有标签类与信息框数据）
  2. YOLOV8进行数据训练需要注意.yaml的路径设置要注意，以及对应路径下所包含的文件名不能改变，且文件摆放不能改变具体可以参考上文我推荐的文章。

第一次发表文章，望对大家有所帮助。

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/Gausst松鼠会/article/detail/724046
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号