main.py
3.76 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import cv2
import os
import numpy as np
filenames = ['13.jpg', '14.jpg', '15.jpg']
def generate_dataset(filename):
src = cv2.imread(filename, cv2.IMREAD_GRAYSCALE) # 그레이 스케일로 변환
src = cv2.resize(src, (int(src.shape[1] / 5), int(src.shape[0] / 5))) # 이미지의 크기가 워낙 크기 때문에 5분의1로 줄여서 확인
cv2.imshow('gray', src)
k = cv2.waitKey(0)
cv2.destroyAllWindows()
src = src[0:src.shape[0] - 10, 15:src.shape[1] - 25] # 양 옆의 노이즈를 제거
ret, binary = cv2.threshold(src, 170, 255, cv2.THRESH_BINARY_INV) # 영상 이진화
cv2.imshow('binary', binary)
k = cv2.waitKey(0)
cv2.destroyAllWindows()
# binary = cv2.morphologyEx(binary , cv2.MORPH_OPEN , cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2,2)), iterations = 2)
# cv2.imshow('binary',binary)
# k = cv2.waitKey(0)
# cv2.destroyAllWindows()
contours, hierarchy = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # 외곽선 검출
color = cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR) # 이진화 이미지를 color이미지로 복사(확인용)
cv2.drawContours(color, contours, -1, (0, 255, 0), 3) # 초록색으로 외곽선을 그려준다.
# 리스트연산을 위해 초기변수 선언
bR_arr = []
digit_arr = []
digit_arr2 = []
count = 0
# 검출한 외곽선에 사각형을 그려서 배열에 추가
for i in range(len(contours)):
bin_tmp = binary.copy()
x, y, w, h = cv2.boundingRect(contours[i])
bR_arr.append([x, y, w, h])
bR_arr = sorted(bR_arr, key=lambda num: num[0], reverse=False)
print(bR_arr[:5])
print(len(bR_arr))
# 작은 노이즈데이터 버림,사각형그리기, 10개씩 리스트로 다시 묶어서 저장
for x, y, w, h in bR_arr:
tmp_y = bin_tmp[y - 2:y + h + 2, x - 2:x + w + 2].shape[0]
tmp_x = bin_tmp[y - 2:y + h + 2, x - 2:x + w + 2].shape[1]
if tmp_x and tmp_y > 10:
count += 1
cv2.rectangle(color, (x - 2, y - 2), (x + w + 2, y + h + 2), (0, 0, 255), 1)
digit_arr.append(bin_tmp[y - 2:y + h + 2, x - 2:x + w + 2])
if count == 10:
digit_arr2.append(digit_arr)
digit_arr = []
count = 0
cv2.imshow('contours', color)
k = cv2.waitKey(0)
cv2.destroyAllWindows()
for i in range(0, len(digit_arr2)):
for j in range(len(digit_arr2[i])):
count += 1
if i == 0: # 1일 경우 비율 유지를 위해 마스크를 만들어 그위에 얹어줌
width = digit_arr2[i][j].shape[1]
height = digit_arr2[i][j].shape[0]
tmp = (height - width) / 2
mask = np.zeros((height, height))
mask[0:height, int(tmp):int(tmp) + width] = digit_arr2[i][j]
digit_arr2[i][j] = cv2.resize(mask, (28, 28))
else:
digit_arr2[i][j] = cv2.resize(digit_arr2[i][j], (28, 28))
# if i == 9 : i = -1
cv2.imwrite(f"./juno/data/{i}/" + str(i) + '_' + str(j) + f"_{filename.split('.')[0]}.png", digit_arr2[i][j])
for filename in filenames:
generate_dataset(filename)
# def createdataset(directory): #sklearn사용을 위해 데이터세트를 생성
# files = os.listdir(directory)
# x = []
# y = []
# for file in files:
# attr_x = cv2.imread(directory+file, cv2.IMREAD_GRAYSCALE)
# attr_x = attr_x.flatten()
# attr_y = file[0]
# x.append(attr_x)
# y.append(attr_y)
#
# x = np.array(x)
# y = np.array(y)
# return x , y
#
# train_dir = './juno/data/'
# train_x ,train_y = createdataset(train_dir)
# test_dir = './juno/test/'
# test_x , test_y = createdataset(test_dir)