# 二值化图像
|
import os
|
import random
|
|
import cv2
|
import matplotlib.pyplot as plt
|
import numpy
|
|
SHOW_PLT = False
|
|
|
def gray_img(img):
|
if img.ndim == 2:
|
return img
|
result = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
return result
|
|
|
def is_col_empty(img, col, start_row, end_row, thresh_hold=0):
|
for r in range(start_row, end_row + 1):
|
if img[r][col] > thresh_hold:
|
return False
|
return True
|
|
|
def is_col_full(img, col, start_row, end_row, thresh_hold_start=64, thresh_hold_end=255, thresh_empty_count=1):
|
rows, cols = img.shape
|
empty_count = 0
|
for r in range(start_row, end_row + 1):
|
if img[r][col] < thresh_hold_start or img[r][col] > thresh_hold_end:
|
empty_count += 1
|
if empty_count > thresh_empty_count:
|
return False
|
return True
|
|
|
# 行是否是满的
|
def is_row_full(img, row, start_col, end_col, thresh_hold_start=64, thresh_hold_end=255):
|
rows, cols = img.shape
|
for c in range(start_col, end_col + 1):
|
if img[row][c] < thresh_hold_start or img[row][c] > thresh_hold_end:
|
return False
|
return True
|
|
|
def format_num_img(img):
|
# 先调整为方形
|
rows, cols = img.shape
|
img_new = numpy.zeros((max(rows, cols), max(rows, cols)), numpy.uint8)
|
img_new.fill(0)
|
start_col = (max(rows, cols) - cols) // 2
|
if rows > cols:
|
for r in range(rows):
|
for c in range(cols):
|
img_new[r][c + start_col] = img[r][c]
|
else:
|
pass
|
# 设置为28X28像素
|
return cv2.resize(img_new, (28, 28), interpolation=cv2.INTER_AREA)
|
|
|
# 分隔数字
|
def split_nums_img(img):
|
# 获取代码的最大色值
|
rows, cols = img.shape
|
max_color = 0
|
for r in range(0, rows):
|
for c in range(0, cols):
|
if img[r][c] > max_color:
|
max_color = img[r][c]
|
|
start_index = -1
|
end_index = -1
|
codes_pos = []
|
for col in range(0, cols):
|
if not is_col_empty(img, col, 0, rows - 1, max_color * 2 // 3):
|
if start_index < 0:
|
start_index = col
|
end_index = col
|
else:
|
if end_index >= 0 and start_index >= 0:
|
codes_pos.append((start_index, end_index + 1))
|
end_index = -1
|
start_index = -1
|
img_detail = []
|
for i in range(len(codes_pos)):
|
temp_img = img[0:rows - 1, codes_pos[i][0]:codes_pos[i][1]]
|
temp_img = format_num_img(temp_img)
|
img_detail.append(temp_img)
|
return img_detail
|
|
|
# 分离同花顺的代码
|
def clip_ths_code_area(img):
|
img = gray_img(img)
|
rows, cols = img.shape
|
|
# 行分隔
|
full_row = -1
|
start_row = -1
|
end_row = -1
|
for r in range(0, rows, 1):
|
if is_row_full(img, r, 0, cols // 2, 38, 38):
|
# print("找到分隔：", col)
|
full_row = r
|
if start_row >= 0 and r - start_row > 10:
|
end_row = r - 1
|
break
|
|
else:
|
if full_row > -1:
|
full_row = -1
|
if start_row < 0:
|
start_row = r
|
if end_row < 0:
|
end_row = rows - 1
|
|
if start_row < 0:
|
raise Exception("没找到上分割线")
|
|
start_col = cols - 1
|
for c in range(cols - 1, -1, -1):
|
if is_col_full(img, c, start_row, end_row, 38, 38, 2):
|
# print("找到分隔：", col)
|
start_col = c + 1
|
break
|
# 往前找数字分隔
|
content_start = -1
|
empty_start = -1
|
empty_end = -1
|
start_index = -1
|
end_index = -1
|
codes_pos = []
|
# 获取代码的最大色值
|
max_color = 0
|
for r in range(start_row, end_row + 1):
|
for c in range(cols - 1, start_col + (cols - start_col) // 3, -1):
|
if img[r][c] > max_color:
|
max_color = img[r][c]
|
# 从后往前找
|
for col in range(cols - 1, start_col, -1):
|
if not is_col_empty(img, col, start_row, end_row, max_color * 2 // 3):
|
if start_index < 0:
|
start_index = col
|
end_index = col
|
else:
|
if end_index >= 0 and start_index >= 0:
|
codes_pos.append((end_index - 1, start_index + 1))
|
end_index = -1
|
start_index = -1
|
codes_pos = codes_pos[:6]
|
codes_pos.sort(key=lambda x: x[0])
|
|
if SHOW_PLT:
|
plt.figure(figsize=(10, 4))
|
img_detail = []
|
for i in range(len(codes_pos)):
|
if SHOW_PLT:
|
plt.subplot(2, 5, i + 1)
|
plt.title(f'pred {i}')
|
plt.axis('off')
|
temp_img = img[start_row:end_row, codes_pos[i][0]:codes_pos[i][1]]
|
temp_img = format_num_img(temp_img)
|
img_detail.append(temp_img)
|
if SHOW_PLT:
|
plt.imshow(temp_img, cmap='gray')
|
if SHOW_PLT:
|
plt.show()
|
|
clip_img = img[start_row:end_row, codes_pos[0][0]:codes_pos[-1][1]]
|
# cv2.imwrite("test1.png", clip_img)
|
return clip_img, img_detail
|
|
# print(clip_img.shape)
|
# ret1, p1 = cv2.threshold(src=clip_img, thresh=100, maxval=255, type=cv2.THRESH_BINARY)
|
# cv2.imwrite("D:/workspace/GP/trade_desk/test3.png", p1)
|
|
|
def __test4():
|
files = os.listdir("datas/test4/")
|
for file in files:
|
code = file[:-4]
|
img = cv2.imread(f"datas/test4/{file}", cv2.IMREAD_GRAYSCALE)
|
img_details = split_nums_img(img)
|
for d in range(0, len(img_details)):
|
cv2.imwrite(f"C:/Users/Administrator/Desktop/ocr/codes/{code}_{random.randint(0, 100000)}.png",
|
img_details[d])
|
plt.figure(figsize=(10, 4))
|
for i in range(0, len(img_details)):
|
plt.subplot(2, 5, i + 1)
|
plt.title(f'pred {i}')
|
plt.axis('off')
|
plt.imshow(img_details[i], cmap='gray')
|
plt.show()
|
|
|
def __test3():
|
files = os.listdir("datas/test3/")
|
for file in files:
|
code = file[:-4]
|
img = cv2.imread(f"datas/test3/{file}", cv2.IMREAD_GRAYSCALE)
|
rows, cols = img.shape
|
for r in range(rows):
|
for c in range(cols):
|
img[r][c] = 255 - img[r][c]
|
|
img_details = split_nums_img(img)
|
for d in range(0, len(img_details)):
|
cv2.imwrite(f"C:/Users/Administrator/Desktop/ocr/codes/{code}_{random.randint(0, 100000)}.png",
|
img_details[d])
|
plt.figure(figsize=(10, 4))
|
for i in range(0, len(img_details)):
|
plt.subplot(2, 5, i + 1)
|
plt.title(f'pred {i}')
|
plt.axis('off')
|
plt.imshow(img_details[i], cmap='gray')
|
plt.show()
|
|
|
if __name__ == '__main__':
|
gray_img = cv2.imread('D:/test2.png', cv2.IMREAD_GRAYSCALE)
|
clip_ths_code_area(gray_img)
|
|
if __name__ == "__main__1":
|
#
|
files = os.listdir("datas/test/")
|
for file in files:
|
code = file[:6]
|
img = cv2.imread(f"datas/test/{file}", cv2.IMREAD_GRAYSCALE)
|
img, img_details = clip_ths_code_area(img)
|
for d in range(0, len(img_details)):
|
cv2.imwrite(f"C:/Users/Administrator/Desktop/ocr/codes/{code}_{random.randint(0, 100000)}.png",
|
img_details[d])
|
if SHOW_PLT:
|
plt.figure(figsize=(1, 1))
|
plt.subplot(1, 1, 1)
|
plt.title(f"test")
|
plt.axis('off')
|
plt.imshow(img)
|
plt.show()
|
pass
|
# img = gray_img(cv2.imread("C:\\Users\\Administrator\\Desktop\\ocr\\code_test.png"))
|
# h = img.shape[0]
|
# w = img.shape[1]
|
# img = cv2.resize(img, (int(w * 2), int(h * 2)), interpolation=cv2.INTER_AREA)
|
# cv2.imwrite("C:\\Users\\Administrator\\Desktop\\ocr\\code_test_gray.png", img)
|
