Development of the ocr part of AOI
Samo Penic
2018-11-16 9efc18e0eed88475c205abf9fc564ee0f5c73cb6
commit | author | age
e555c0 1 from pyzbar.pyzbar import decode
9efc18 2 from sid_process import enhanceSID
e555c0 3 import cv2
SP 4 import numpy as np
5 import math
6
7
511c2e 8 class Paper:
SP 9     def __init__(self, filename=None):
10         self.filename = filename
11         self.invalid = None
12         self.QRData = None
13         self.errors = []
14         self.warnings = []
15         if filename is not None:
16             self.loadImage(filename)
17             self.runOcr()
e555c0 18
511c2e 19     def loadImage(self, filename, rgbchannel=0):
SP 20         self.img = cv2.imread(filename, rgbchannel)
21         if self.img is None:
22             self.errors.append("File could not be loaded!")
23             self.invalid = True
24             return
25         self.imgHeight, self.imgWidth = self.img.shape[0:2]
e555c0 26
511c2e 27     def saveImage(self, filename="debug_image.png"):
SP 28         cv2.imwrite(filename, self.img)
e555c0 29
511c2e 30     def runOcr(self):
SP 31         if self.invalid == True:
32             return
33         self.decodeQRandRotate()
34         self.imgTreshold()
35         skewAngle = 0
36         #         try:
37         #             skewAngle=self.getSkewAngle()
38         #         except:
39         #             self.errors.append("Could not determine skew angle!")
40         #         self.rotateAngle(skewAngle)
e555c0 41
511c2e 42         self.generateAnswerMatrix()
e555c0 43
511c2e 44         self.saveImage()
e555c0 45
511c2e 46     def decodeQRandRotate(self):
SP 47         if self.invalid == True:
48             return
49         blur = cv2.blur(self.img, (3, 3))
50         d = decode(blur)
51         self.img = blur
52         if len(d) == 0:
53             self.errors.append("QR code could not be found!")
54             self.data = None
55             self.invalid = True
56             return
57         self.QRDecode = d
58         self.QRData = d[0].data
59         xpos = d[0].rect.left
60         ypos = d[0].rect.top
61         # check if image is rotated wrongly
62         if xpos > self.imgHeight / 2.0 and ypost > self.imgWidth / 2.0:
63             self.rotateAngle(180)
e555c0 64
511c2e 65     def rotateAngle(self, angle=0):
SP 66         rot_mat = cv2.getRotationMatrix2D(
67             (self.imgHeight / 2, self.imgWidth / 2), angle, 1.0
68         )
69         result = cv2.warpAffine(
70             self.img,
71             rot_mat,
72             (self.imgHeight, self.imgWidth),
73             flags=cv2.INTER_CUBIC,
74             borderMode=cv2.BORDER_CONSTANT,
75             borderValue=(255, 255, 255),
76         )
e555c0 77
511c2e 78         self.img = result
SP 79         self.imgHeight, self.imgWidth = self.img.shape[0:2]
e555c0 80
511c2e 81         # todo, make better tresholding
e555c0 82
511c2e 83     def imgTreshold(self):
SP 84         (self.thresh, self.bwimg) = cv2.threshold(
85             self.img, 128, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU
86         )
e555c0 87
511c2e 88     def getSkewAngle(self):
SP 89         neg = 255 - self.bwimg  # get negative image
90         cv2.imwrite("debug_1.png", neg)
e555c0 91
511c2e 92         angle_counter = 0  # number of angles
SP 93         angle = 0.0  # collects sum of angles
94         cimg = cv2.cvtColor(self.img, cv2.COLOR_GRAY2BGR)
e555c0 95
511c2e 96         # get all the Hough lines
SP 97         for line in cv2.HoughLinesP(neg, 1, np.pi / 180, 325):
98             x1, y1, x2, y2 = line[0]
99             cv2.line(cimg, (x1, y1), (x2, y2), (0, 0, 255), 2)
100             # calculate the angle (in radians)
101             this_angle = np.arctan2(y2 - y1, x2 - x1)
102             if this_angle and abs(this_angle) <= 10:
103                 # filtered zero degree and outliers
104                 angle += this_angle
105                 angle_counter += 1
e555c0 106
511c2e 107                 # the skew is calculated of the mean of the total angles, #try block helps with division by zero.
SP 108         try:
109             skew = np.rad2deg(
110                 angle / angle_counter
111             )  # the 1.2 factor is just experimental....
112         except:
113             skew = 0
e555c0 114
511c2e 115         cv2.imwrite("debug_2.png", cimg)
SP 116         return skew
e555c0 117
511c2e 118     def locateUpMarkers(self, threshold=0.85, height=200):
SP 119         template = cv2.imread("template.png", 0)
120         w, h = template.shape[::-1]
121         crop_img = self.img[0:height, :]
122         res = cv2.matchTemplate(crop_img, template, cv2.TM_CCOEFF_NORMED)
123         loc = np.where(res >= threshold)
124         cimg = cv2.cvtColor(crop_img, cv2.COLOR_GRAY2BGR)
125         # remove false matching of the squares in qr code
126         loc_filtered_x = []
127         loc_filtered_y = []
128         if len(loc[0]) == 0:
129             min_y = -1
130         else:
131             min_y = np.min(loc[0])
132             for pt in zip(*loc[::-1]):
133                 if pt[1] < min_y + 20:
134                     loc_filtered_y.append(pt[1])
135                     loc_filtered_x.append(pt[0])
136                     # order by x coordinate
137             loc_filtered_x, loc_filtered_y = zip(
138                 *sorted(zip(loc_filtered_x, loc_filtered_y))
139             )
140         # loc=[loc_filtered_y,loc_filtered_x]
141         # remove duplicates
142             a = np.diff(loc_filtered_x) > 40
143             a = np.append(a, True)
144             loc_filtered_x = np.array(loc_filtered_x)
145             loc_filtered_y = np.array(loc_filtered_y)
146             loc = [loc_filtered_y[a], loc_filtered_x[a]]
147             for pt in zip(*loc[::-1]):
148                 cv2.rectangle(cimg, pt, (pt[0] + w, pt[1] + h), (0, 255, 255), 2)
e555c0 149
511c2e 150         cv2.imwrite("debug_3.png", cimg)
e555c0 151
511c2e 152         self.xMarkerLocations = loc
SP 153         return loc
e555c0 154
511c2e 155     def locateRightMarkers(self, threshold=0.85, width=200):
SP 156         template = cv2.imread("template.png", 0)
157         w, h = template.shape[::-1]
158         crop_img = self.img[:, -width:]
159         res = cv2.matchTemplate(crop_img, template, cv2.TM_CCOEFF_NORMED)
160         loc = np.where(res >= threshold)
161         cimg = cv2.cvtColor(crop_img, cv2.COLOR_GRAY2BGR)
162         # remove false matching of the squares in qr code
163         loc_filtered_x = []
164         loc_filtered_y = []
165         if len(loc[1]) == 0:
166             min_x = -1
167         else:
168             max_x = np.max(loc[1])
169             for pt in zip(*loc[::-1]):
170                 if pt[1] > max_x - 20:
171                     loc_filtered_y.append(pt[1])
172                     loc_filtered_x.append(pt[0])
173                     # order by y coordinate
174             loc_filtered_y, loc_filtered_x = zip(
175                 *sorted(zip(loc_filtered_y, loc_filtered_x))
176             )
177             # loc=[loc_filtered_y,loc_filtered_x]
178             # remove duplicates
179             a = np.diff(loc_filtered_y) > 40
180             a = np.append(a, True)
181             loc_filtered_x = np.array(loc_filtered_x)
182             loc_filtered_y = np.array(loc_filtered_y)
183             loc = [loc_filtered_y[a], loc_filtered_x[a]]
184             for pt in zip(*loc[::-1]):
185                 cv2.rectangle(cimg, pt, (pt[0] + w, pt[1] + h), (0, 255, 255), 2)
e555c0 186
511c2e 187         cv2.imwrite("debug_4.png", cimg)
e555c0 188
511c2e 189         self.yMarkerLocations = [loc[0], loc[1] + self.imgWidth - width]
SP 190         return self.yMarkerLocations
e555c0 191
511c2e 192     def generateAnswerMatrix(self):
SP 193         self.locateUpMarkers()
194         self.locateRightMarkers()
e555c0 195
511c2e 196         roixoff = 10
SP 197         roiyoff = 5
198         roiwidth = 50
199         roiheight = roiwidth
200         totpx = roiwidth * roiheight
e555c0 201
511c2e 202         self.answerMatrix = []
SP 203         for y in self.yMarkerLocations[0]:
204             oneline = []
205             for x in self.xMarkerLocations[1]:
206                 roi = self.bwimg[
207                     y - roiyoff : y + int(roiheight - roiyoff),
208                     x - roixoff : x + int(roiwidth - roixoff),
209                 ]
210                 # cv2.imwrite('ans_x'+str(x)+'_y_'+str(y)+'.png',roi)
211                 black = totpx - cv2.countNonZero(roi)
212                 oneline.append(black / totpx)
213             self.answerMatrix.append(oneline)
9efc18 214
SP 215     def get_enhanced_sid(self):
216         es= enhanceSID(self.img[int(0.04*self.imgHeight):int(0.08*self.imgHeight), int(0.7*self.imgWidth):int(0.99*self.imgWidth)])
217         cv2.imwrite("enhancedSID.png",es)