经验首页 前端设计 程序设计 Java相关 移动开发 数据库/运维 软件/图像 大数据/云计算 其他经验
当前位置:技术经验 » 大数据/云/AI » 人工智能基础 » 查看文章
OpenCV + sklearnSVM 实现手写数字分割和识别
来源:cnblogs  作者:凪风sama  时间:2024/6/17 15:17:50  对本文有异议

这学期机器学习考核方式以大作业的形式进行考核,而且只能使用一些传统的机器学习算法。
综合再三,选择了自己比较熟悉的MNIST数据集以及OpenCV来完成手写数字的分割和识别作为大作业。

1. 数据集准备

MNIST数据集是一个手写数字的数据库,包含60000张训练图片和10000张测试图片,每张图片大小为28x28像素,每张图片都是一个
灰度图,像素取值范围在0-255之间。

这里使用pytorch的torchvision.datasets模块来读取MNIST数据集。

  1. from torchvision import datasets
  2. mnist_set = datasets.MNIST(root="./MNIST", train=True, download=True)

具体参数说明请自行搜索。注意若donwload=True,则torchvision会通过内置链接自动下载数据集,
但是有时会失效。因此可以自己去网络上下载并解压后排列成指定文件树,如下

  1. MNIST
  2. ├── MNSIT
  3. ├── raw
  4. ├── t10k-images-idx3-ubyte.gz
  5. ├── t10k-labels-idx1-ubyte.gz
  6. ├── train-images-idx3-ubyte.gz
  7. └── train-labels-idx1-ubyte.gz

然后使用如下语句去读取数据集

  1. img, target = minst_set[0]

其中每个img类型为PILimage,target类型为int,代表该图片对应的数字。

但是在喂给SVM训练时需要的是[batch_size, data]大小的numpy数组,因此需要做一些预处理

  1. x_, y_ = list(zip(*([(np.array(img).reshape(28*28), target) for img, target in mnist_set])))

上面的语句实现了将MNIST数据集转换成numpy数组的形式,其中x_是每个成员为[1, 784]的numpy数组,y_为对应的数字所组成的列表。

2. SVM训练

支持向量机(support vector machine,SVM)是经典的机器学习算法,其通过选取两个n维支持向量(support vector)之间的n维超平面来对两类对象进行二分类。而专注于分类的SVM又称作Support Vector Classification,SVC。

求解SVM是一个很复杂的问题,但是万幸的是sklearn中有封装的很好的模块,可以很简单的直接使用

  1. from sklearn.svm import SVC
  2. svc = SVC(kernel='rbf', C=1)
  3. svc.fit(x_, y_)

其中fit接口接受两个参数,第一个参数为训练数据[batch_size, data],第二个参数为训练标签[batch_size,1]。
SVC的构造函数如下

  1. SVC(C=1.0, kernel='rbf', degree=3, gamma='scale', coef0=0.0, shrinking=True, probability=False, tol=0.001, cache_size=200, class_weight=None, verbose=False, max_iter=-1, decision_function_shape='ovr', random_state=None)

比较重要的参数有kernel,C,decison_function_shape等。

  • kernel参数指定了核函数,常用的有linear,poly,rbf,sigmoid等。
  • C为惩罚系数,C越大,对误分类的惩罚越大,模型越保守,C越小,对误分类的惩罚越小,模型越宽松,也就是较大的C在训练集上会有更高的正确率,较小的C会容许噪声的存在,泛化能力较强。
  • decision_function_shape参数指定了决策函数的形状,ovr表示one-vs-rest,ovo表示one-vs-one,具体的意思可以网络查阅

4. 数字分割

数字分割是指将图像中的数字部分分割出来,然后一个一个喂给SVM进行分类

这里就是使用opencv对拍摄的图像进行轮廓提取后拟合外接矩形,借此来提取数字部分的ROI。

这里选择进行Canny边缘检测后去进行轮廓提取,然后拟合外接矩形,因为相较于直接二值化后去提取数字部分的ROI,
边缘检测对数字与纸张的边界更加敏感,即便在光照不均匀的情况下,也能较好的提取出数字的边缘。鲁棒性强。

5. 杂项与代码

这里还有一些杂项,比如保存模型,加载模型

使用pickle模块对训练好的模型对象进行序列化保存与加载,可以将训练好的模型保存到本地,以便后续使用。

最后贴出代码

代码
  1. import os.path
  2. import cv2
  3. import numpy as np
  4. from matplotlib import pyplot as plt
  5. from torchvision import datasets
  6. from torchvision import transforms
  7. from sklearn import svm
  8. from sklearn import preprocessing
  9. from sklearnex import patch_sklearn
  10. import pickle
  11. from torch.utils.tensorboard import SummaryWriter
  12. from sklearn.model_selection import learning_curve
  13. '''
  14. @brief 加载MNIST数据集并转换格式成二值图
  15. @param train: 是否为训练集
  16. @param data_enhance: 是否进行数据增强
  17. @return 二值图集和标签集
  18. '''
  19. def LoadMnistDataset(train=True, data_enhance=False):
  20. mnist_set = datasets.MNIST(root="./MNIST", train=train, download=True)
  21. x_, y_ = list(zip(*([(np.array(img), target) for img, target in mnist_set])))
  22. sets_raw = []
  23. sets_r20 = []
  24. sets_invr20 = []
  25. y = []
  26. y_r20 = []
  27. y_invr20 = []
  28. sets = []
  29. matrix_r20 = cv2.getRotationMatrix2D((14, 14), 25, 1.0)
  30. matrix_invr20 = cv2.getRotationMatrix2D((14, 14), -25, 1.0)
  31. select = 0
  32. for idx in range(len(x_)):
  33. # 对图像进行二值化以及数据增强
  34. _, img = cv2.threshold(x_[idx], 255, 255, cv2.THRESH_OTSU)
  35. sets_raw.append(np.array(img.data).reshape(784))
  36. y.append(y_[idx])
  37. if data_enhance:
  38. if select % 2 == 0:
  39. img_r20 = ~cv2.warpAffine(~img, matrix_r20, (28, 28), borderValue=(255, 255, 255))
  40. sets_r20.append(np.array(img_r20.data).reshape(784))
  41. y_r20.append(y_[idx])
  42. else:
  43. img_invr20 = ~cv2.warpAffine(~img, matrix_invr20, (28, 28), borderValue=(255, 255, 255))
  44. sets_invr20.append(np.array(img_invr20.data).reshape(784))
  45. y_invr20.append(y_[idx])
  46. select += 1
  47. # 数据增强
  48. sets = sets_raw + sets_r20 + sets_invr20
  49. sets = np.array(sets)
  50. print(sets.shape)
  51. if data_enhance:
  52. y = y + y_r20 + y_invr20
  53. return sets, y
  54. '''
  55. @brief 保存SVM模型
  56. @param svc_model: SVM模型
  57. @param file_path: 模型保存路径,默认为./SVC
  58. @return None
  59. '''
  60. def SaveSvcModel(svc_model, file_path="./SVC"):
  61. with open(file_path, 'wb') as fs:
  62. pickle.dump(svc_model, fs)
  63. '''
  64. @brief 加载SVM模型
  65. @param file_path: 模型保存路径,默认为./SVC
  66. @return SVM模型
  67. '''
  68. def LoadSvcModel(file_path="./SVC"):
  69. if not os.path.exists(file_path):
  70. assert "Model Do Not Exist"
  71. with open(file_path, 'rb') as fs:
  72. svc_model = pickle.load(fs)
  73. return svc_model
  74. '''
  75. @brief 训练SVM模型
  76. @param c: SVM参数C
  77. @param enhance: 是否进行数据增强
  78. @return acc: 在测试集上的准确率
  79. svc_model: SVM模型
  80. '''
  81. def TrainSvc(c, enhance):
  82. # 读取数据集,训练集及测试集
  83. images_train, targets_train = LoadMnistDataset(train=True, data_enhance=enhance)
  84. images_test, targets_test = LoadMnistDataset(train=False, data_enhance=enhance)
  85. # 训练
  86. svc_model = svm.SVC(C=c,kernel='rbf', decision_function_shape='ovr')
  87. svc_model.fit(images_train, targets_train)
  88. # 在测试集上测试准确度
  89. res = svc_model.predict(images_test)
  90. correct = (res == targets_test).sum()
  91. accuracy = correct / len(images_test)
  92. print(f"测试集上的准确率为{accuracy * 100}%")
  93. return svc_model
  94. '''
  95. @brief 预处理比较粗的字体
  96. @param image: 输入图像
  97. @:param show: 是否显示预处理后的图像
  98. @:param thresh: 二值化阈值
  99. @return 预处理后的图像数据
  100. '''
  101. def PreProcessFatFont(image, show=False):
  102. # 白底黑字转黑底白字
  103. pre_ = ~image
  104. # 转单通道灰度
  105. pre_ = cv2.cvtColor(pre_, cv2.COLOR_BGR2GRAY)
  106. # 二值化
  107. _, pre_ = cv2.threshold(pre_, thresh=0, maxval=255, type=cv2.THRESH_OTSU)
  108. # resize后添加黑色边框,亲测可提高识别率
  109. pre_ = cv2.resize(pre_, (112, 112))
  110. _, pre_ = cv2.threshold(pre_, thresh=0, maxval=255, type=cv2.THRESH_OTSU)
  111. back = np.zeros((300, 300), np.uint8)
  112. back[29:141, 29:141] = pre_
  113. pre_ = back
  114. if show:
  115. cv2.imshow("show", pre_)
  116. cv2.waitKey(0)
  117. # 做一次开运算(腐蚀 + 膨胀)
  118. kernel = np.ones((2, 2), np.uint8)
  119. pre_ = cv2.erode(pre_, kernel, iterations=1)
  120. kernel = np.ones((3, 3), np.uint8)
  121. pre_ = cv2.dilate(pre_, kernel, iterations=1)
  122. # 第二次resize
  123. pre_ = cv2.resize(pre_, (56, 56))
  124. _, pre_ = cv2.threshold(pre_, thresh=0, maxval=255, type=cv2.THRESH_OTSU)
  125. # 做一次开运算(腐蚀 + 膨胀)
  126. kernel = np.ones((2, 2), np.uint8)
  127. pre_ = cv2.erode(pre_, kernel, iterations=1)
  128. kernel = np.ones((3, 3), np.uint8)
  129. pre_ = cv2.dilate(pre_, kernel, iterations=1)
  130. # resize成输入规格
  131. pre_ = cv2.resize(pre_, (28, 28))
  132. _, pre_ = cv2.threshold(pre_, thresh=0, maxval=255, type=cv2.THRESH_OTSU)
  133. # 转换为SVM的输入格式
  134. pre_ = np.array(pre_).flatten().reshape(1, -1)
  135. return pre_
  136. '''
  137. @brief 预处理细的字体
  138. @param image: 输入图像
  139. @param show: 是否显示预处理后的图像
  140. @param thresh: 二值化阈值
  141. @return 预处理后的图像数据
  142. '''
  143. def PreProcessThinFont(image, show=False):
  144. # 白底黑字转黑底白字
  145. pre_ = ~image
  146. # 转灰度图
  147. pre_ = cv2.cvtColor(pre_, cv2.COLOR_BGR2GRAY)
  148. # 增加黑色边框
  149. pre_ = cv2.resize(pre_, (112, 112))
  150. _, pre_ = cv2.threshold(pre_,thresh=0, maxval=255, type=cv2.THRESH_OTSU)
  151. back = np.zeros((170, 170), dtype=np.uint8) # 这里不指明类型会导致后续矩阵强转为float64,无法使用大津法阈值
  152. back[29:141, 29:141] = pre_
  153. pre_ = back
  154. if show:
  155. cv2.imshow("show", pre_)
  156. cv2.waitKey(0)
  157. # 对细字体先膨胀一下
  158. kernel = np.ones((3, 3), np.uint8)
  159. pre_ = cv2.dilate(pre_, kernel, iterations=2)
  160. # 第二次resize
  161. pre_ = cv2.resize(pre_, (56, 56))
  162. _, pre_ = cv2.threshold(pre_, thresh=0, maxval=255, type=cv2.THRESH_OTSU)
  163. # 做一次开运算(腐蚀 + 膨胀)
  164. kernel = np.ones((2, 2), np.uint8)
  165. pre_ = cv2.erode(pre_, kernel, iterations=1)
  166. kernel = np.ones((3, 3), np.uint8)
  167. pre_ = cv2.dilate(pre_, kernel, iterations=1)
  168. # resize成输入规格
  169. pre_ = cv2.resize(pre_, (28, 28))
  170. _, pre_ = cv2.threshold(pre_, thresh=0, maxval=255, type=cv2.THRESH_OTSU)
  171. # 转换为SVM输入格式
  172. pre_ = np.array(pre_).flatten().reshape(1, -1)
  173. return pre_
  174. '''
  175. @brief 在空白背景上显示提取出的roi
  176. @param res_list: roi列表
  177. @return None
  178. '''
  179. def ShowRoi(res_list):
  180. back = 255 * np.ones((1000, 1500, 3), dtype=np.uint8)
  181. # 图片x轴偏移量
  182. tlx = 0
  183. for roi in res_list:
  184. if tlx + roi.shape[1] > back.shape[1]:
  185. break
  186. # 每次在原图上加上一个roi
  187. back[0:roi.shape[0], tlx:tlx + roi.shape[1], :] = roi
  188. tlx += roi.shape[1]
  189. cv2.imshow("show", back)
  190. cv2.waitKey(0)
  191. '''
  192. @brief 寻找数字轮廓并提取roi
  193. @param src: 输入图像
  194. @param thin: 是否为细字体
  195. @param thresh: 二值化阈值
  196. @return roi列表
  197. '''
  198. def FindNumbers(src, thin=True):
  199. # 拷贝
  200. dst = src.copy()
  201. paint = src.copy()
  202. roi = src.copy()
  203. dst = ~dst
  204. # 预处理
  205. paint = cv2.resize(paint, (448, 448))
  206. dst = cv2.resize(dst, (448, 448))
  207. # 记录缩放比例,后来看这一步好像没啥意义
  208. fx = src.shape[1] / 448
  209. fy = src.shape[0] / 448
  210. # 转单通道
  211. dst = cv2.cvtColor(dst, cv2.COLOR_BGR2GRAY)
  212. # 边缘检测后二值化,直接二值化的话由于采光不同的原因灰度直方图峰与峰之间可能会差距过大,导致二值图的分割不准确
  213. # 而边缘检测对像素突变更加敏感,因此采用Canny边缘检测后二值化
  214. cv2.Canny(dst, 200, 200, dst)
  215. # 对于平常笔写的字太细,膨胀一下
  216. if thin:
  217. kernel = np.ones((5, 5), np.uint8)
  218. dst = cv2.dilate(dst, kernel, iterations=1)
  219. # 寻找外围轮廓
  220. contours, _ = cv2.findContours(dst, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  221. # 提取roi
  222. roi_list = []
  223. rect_list = []
  224. for contour in contours:
  225. rect = cv2.boundingRect(contour)
  226. if not ((rect[2] * rect[3] < 400 or rect[2] * rect[3] > 448 * 448 / 2.5) or (rect[3] < rect[2])):
  227. cv2.rectangle(paint, rect, (255, 0, 0), 1)
  228. x_min = rect[0] * fx
  229. x_max = (rect[0] + rect[2]) * fx
  230. y_min = rect[1] * fy
  231. y_max = (rect[1] + rect[3]) * fy
  232. roi_list.append(roi[int(y_min):int(y_max), int(x_min):int(x_max)].copy())
  233. rect_list.append(rect)
  234. return paint, roi_list, rect_list
  235. '''
  236. @brief 以txt形式显示数据
  237. @param data: 数据集
  238. @return None
  239. '''
  240. def ShowDataTxt(data):
  241. print("----------------------------------------------------------")
  242. for i in range(28):
  243. for j in range(28):
  244. print(0 if data[0][i * 28 + j] == 255 else 1, end='')
  245. print('\n')
  246. print("----------------------------------------------------------")
  247. if __name__ == "__main__":
  248. # 加载
  249. patch_sklearn()
  250. model_path = "./SVC_C1_enhance.pkl"
  251. if os.path.exists(model_path):
  252. print("Model Exist, Load Form Serialization")
  253. model = LoadSvcModel(model_path)
  254. else:
  255. print("Model Do Not Exist, Train")
  256. # 训练
  257. model = TrainSvc(1, False)
  258. # 保存
  259. SaveSvcModel(model, model_path)
  260. # 测试
  261. paint, nums, rects = FindNumbers(cv2.imread("test_final.jpg"))
  262. predict_nums = []
  263. for img in nums:
  264. data = PreProcessThinFont(img, show=False)
  265. # ShowDataTxt(data)
  266. predict_nums.append(model.predict(data)[0])
  267. for i in range(len(predict_nums)):
  268. cv2.putText(paint,str(predict_nums[i]), (rects[i][0], rects[i][1]), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
  269. cv2.imshow("show", paint)
  270. cv2.waitKey(0)

给出几个识别后的效果:
image

原文链接:https://www.cnblogs.com/CrescentWind/p/18252449

 友情链接:直通硅谷  点职佳  北美留学生论坛

本站QQ群:前端 618073944 | Java 606181507 | Python 626812652 | C/C++ 612253063 | 微信 634508462 | 苹果 692586424 | C#/.net 182808419 | PHP 305140648 | 运维 608723728

W3xue 的所有内容仅供测试,对任何法律问题及风险不承担任何责任。通过使用本站内容随之而来的风险与本站无关。
关于我们  |  意见建议  |  捐助我们  |  报错有奖  |  广告合作、友情链接(目前9元/月)请联系QQ:27243702 沸活量
皖ICP备17017327号-2 皖公网安备34020702000426号