Development of the ocr part of AOI
Samo Penic
2018-11-17 cf921b251d8664900bc7c5b3068bcd7b0ce2b2b7
commit | author | age
e555c0 1 from pyzbar.pyzbar import decode
02e0f7 2 from sid_process import getSID
e555c0 3 import cv2
SP 4 import numpy as np
5 import math
6
7
511c2e 8 class Paper:
762a5e 9     def __init__(self, filename=None, sid_classifier=None, settings=None):
511c2e 10         self.filename = filename
SP 11         self.invalid = None
12         self.QRData = None
0436f6 13         self.settings={'answer_treshold':0.25,} if settings is None else settings
511c2e 14         self.errors = []
SP 15         self.warnings = []
0436f6 16         self.sid=None
762a5e 17         self.sid_classifier = sid_classifier
511c2e 18         if filename is not None:
SP 19             self.loadImage(filename)
20             self.runOcr()
e555c0 21
511c2e 22     def loadImage(self, filename, rgbchannel=0):
SP 23         self.img = cv2.imread(filename, rgbchannel)
24         if self.img is None:
25             self.errors.append("File could not be loaded!")
26             self.invalid = True
27             return
28         self.imgHeight, self.imgWidth = self.img.shape[0:2]
e555c0 29
511c2e 30     def saveImage(self, filename="debug_image.png"):
SP 31         cv2.imwrite(filename, self.img)
e555c0 32
511c2e 33     def runOcr(self):
SP 34         if self.invalid == True:
35             return
36         self.decodeQRandRotate()
37         self.imgTreshold()
38         skewAngle = 0
39         #         try:
40         #             skewAngle=self.getSkewAngle()
41         #         except:
42         #             self.errors.append("Could not determine skew angle!")
43         #         self.rotateAngle(skewAngle)
e555c0 44
511c2e 45         self.generateAnswerMatrix()
e555c0 46
511c2e 47         self.saveImage()
e555c0 48
511c2e 49     def decodeQRandRotate(self):
SP 50         if self.invalid == True:
51             return
52         blur = cv2.blur(self.img, (3, 3))
53         d = decode(blur)
54         self.img = blur
55         if len(d) == 0:
56             self.errors.append("QR code could not be found!")
57             self.data = None
58             self.invalid = True
59             return
60         self.QRDecode = d
61         self.QRData = d[0].data
62         xpos = d[0].rect.left
63         ypos = d[0].rect.top
64         # check if image is rotated wrongly
82ec6d 65         if xpos > self.imgHeight / 2.0 and ypos > self.imgWidth / 2.0:
511c2e 66             self.rotateAngle(180)
e555c0 67
511c2e 68     def rotateAngle(self, angle=0):
82ec6d 69         #rot_mat = cv2.getRotationMatrix2D(
SP 70         #    (self.imgHeight / 2, self.imgWidth / 2), angle, 1.0
71         #)
511c2e 72         rot_mat = cv2.getRotationMatrix2D(
82ec6d 73             (self.imgWidth/2, self.imgHeight/2), angle, 1.0
511c2e 74         )
SP 75         result = cv2.warpAffine(
76             self.img,
77             rot_mat,
82ec6d 78             (self.imgWidth, self.imgHeight),
511c2e 79             flags=cv2.INTER_CUBIC,
SP 80             borderMode=cv2.BORDER_CONSTANT,
81             borderValue=(255, 255, 255),
82         )
e555c0 83
511c2e 84         self.img = result
SP 85         self.imgHeight, self.imgWidth = self.img.shape[0:2]
e555c0 86
511c2e 87         # todo, make better tresholding
e555c0 88
511c2e 89     def imgTreshold(self):
SP 90         (self.thresh, self.bwimg) = cv2.threshold(
91             self.img, 128, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU
92         )
e555c0 93
511c2e 94     def getSkewAngle(self):
SP 95         neg = 255 - self.bwimg  # get negative image
96         cv2.imwrite("debug_1.png", neg)
e555c0 97
511c2e 98         angle_counter = 0  # number of angles
SP 99         angle = 0.0  # collects sum of angles
100         cimg = cv2.cvtColor(self.img, cv2.COLOR_GRAY2BGR)
e555c0 101
511c2e 102         # get all the Hough lines
SP 103         for line in cv2.HoughLinesP(neg, 1, np.pi / 180, 325):
104             x1, y1, x2, y2 = line[0]
105             cv2.line(cimg, (x1, y1), (x2, y2), (0, 0, 255), 2)
106             # calculate the angle (in radians)
107             this_angle = np.arctan2(y2 - y1, x2 - x1)
108             if this_angle and abs(this_angle) <= 10:
109                 # filtered zero degree and outliers
110                 angle += this_angle
111                 angle_counter += 1
e555c0 112
511c2e 113                 # the skew is calculated of the mean of the total angles, #try block helps with division by zero.
SP 114         try:
115             skew = np.rad2deg(
116                 angle / angle_counter
117             )  # the 1.2 factor is just experimental....
118         except:
119             skew = 0
e555c0 120
511c2e 121         cv2.imwrite("debug_2.png", cimg)
SP 122         return skew
e555c0 123
511c2e 124     def locateUpMarkers(self, threshold=0.85, height=200):
SP 125         template = cv2.imread("template.png", 0)
126         w, h = template.shape[::-1]
127         crop_img = self.img[0:height, :]
128         res = cv2.matchTemplate(crop_img, template, cv2.TM_CCOEFF_NORMED)
129         loc = np.where(res >= threshold)
130         cimg = cv2.cvtColor(crop_img, cv2.COLOR_GRAY2BGR)
131         # remove false matching of the squares in qr code
132         loc_filtered_x = []
133         loc_filtered_y = []
134         if len(loc[0]) == 0:
135             min_y = -1
136         else:
137             min_y = np.min(loc[0])
138             for pt in zip(*loc[::-1]):
139                 if pt[1] < min_y + 20:
140                     loc_filtered_y.append(pt[1])
141                     loc_filtered_x.append(pt[0])
142                     # order by x coordinate
143             loc_filtered_x, loc_filtered_y = zip(
144                 *sorted(zip(loc_filtered_x, loc_filtered_y))
145             )
02e0f7 146             # loc=[loc_filtered_y,loc_filtered_x]
SP 147             # remove duplicates
511c2e 148             a = np.diff(loc_filtered_x) > 40
SP 149             a = np.append(a, True)
150             loc_filtered_x = np.array(loc_filtered_x)
151             loc_filtered_y = np.array(loc_filtered_y)
152             loc = [loc_filtered_y[a], loc_filtered_x[a]]
153             for pt in zip(*loc[::-1]):
154                 cv2.rectangle(cimg, pt, (pt[0] + w, pt[1] + h), (0, 255, 255), 2)
e555c0 155
511c2e 156         cv2.imwrite("debug_3.png", cimg)
e555c0 157
511c2e 158         self.xMarkerLocations = loc
SP 159         return loc
e555c0 160
511c2e 161     def locateRightMarkers(self, threshold=0.85, width=200):
SP 162         template = cv2.imread("template.png", 0)
163         w, h = template.shape[::-1]
164         crop_img = self.img[:, -width:]
165         res = cv2.matchTemplate(crop_img, template, cv2.TM_CCOEFF_NORMED)
166         loc = np.where(res >= threshold)
167         cimg = cv2.cvtColor(crop_img, cv2.COLOR_GRAY2BGR)
168         # remove false matching of the squares in qr code
169         loc_filtered_x = []
170         loc_filtered_y = []
171         if len(loc[1]) == 0:
172             min_x = -1
173         else:
174             max_x = np.max(loc[1])
175             for pt in zip(*loc[::-1]):
176                 if pt[1] > max_x - 20:
177                     loc_filtered_y.append(pt[1])
178                     loc_filtered_x.append(pt[0])
179                     # order by y coordinate
180             loc_filtered_y, loc_filtered_x = zip(
181                 *sorted(zip(loc_filtered_y, loc_filtered_x))
182             )
183             # loc=[loc_filtered_y,loc_filtered_x]
184             # remove duplicates
185             a = np.diff(loc_filtered_y) > 40
186             a = np.append(a, True)
187             loc_filtered_x = np.array(loc_filtered_x)
188             loc_filtered_y = np.array(loc_filtered_y)
189             loc = [loc_filtered_y[a], loc_filtered_x[a]]
190             for pt in zip(*loc[::-1]):
191                 cv2.rectangle(cimg, pt, (pt[0] + w, pt[1] + h), (0, 255, 255), 2)
e555c0 192
511c2e 193         cv2.imwrite("debug_4.png", cimg)
e555c0 194
511c2e 195         self.yMarkerLocations = [loc[0], loc[1] + self.imgWidth - width]
SP 196         return self.yMarkerLocations
e555c0 197
511c2e 198     def generateAnswerMatrix(self):
SP 199         self.locateUpMarkers()
200         self.locateRightMarkers()
e555c0 201
511c2e 202         roixoff = 10
SP 203         roiyoff = 5
204         roiwidth = 50
205         roiheight = roiwidth
206         totpx = roiwidth * roiheight
e555c0 207
511c2e 208         self.answerMatrix = []
SP 209         for y in self.yMarkerLocations[0]:
210             oneline = []
211             for x in self.xMarkerLocations[1]:
212                 roi = self.bwimg[
213                     y - roiyoff : y + int(roiheight - roiyoff),
214                     x - roixoff : x + int(roiwidth - roixoff),
215                 ]
216                 # cv2.imwrite('ans_x'+str(x)+'_y_'+str(y)+'.png',roi)
217                 black = totpx - cv2.countNonZero(roi)
218                 oneline.append(black / totpx)
219             self.answerMatrix.append(oneline)
9efc18 220
SP 221     def get_enhanced_sid(self):
02e0f7 222         if self.sid_classifier is None:
SP 223             return "x"
762a5e 224         if self.settings is not None:
SP 225             sid_mask=self.settings.get("sid_mask", None)
5cb7c1 226         es,err,warn = getSID(
02e0f7 227             self.img[
d5c694 228                 int(0.04 * self.imgHeight) : int(0.095 * self.imgHeight),
02e0f7 229                 int(0.7 * self.imgWidth) : int(0.99 * self.imgWidth),
SP 230             ],
231             self.sid_classifier,
762a5e 232             sid_mask
02e0f7 233         )
5cb7c1 234         [self.errors.append(e) for e in err]
SP 235         [self.warnings.append(w) for w in warn]
02e0f7 236         return es
0436f6 237
SP 238
239     def get_code_data(self):
cf921b 240         if self.QRData is None:
SP 241             self.errors.append("Could not read QR or EAN code! Not an exam?")
242             retval = {'exam_id': None,
243                       'page_no': None,
244                       'paper_id': None,
245                       'faculty_id': None,
246                       'sid':None
247                       }
248             return retval
0436f6 249         qrdata = bytes.decode(self.QRData, 'utf8')
SP 250         if self.QRDecode[0].type=='EAN13':
251             return {'exam_id': int(qrdata[0:7]),
252                     'page_no': int(qrdata[7]),
253                     'paper_id': int(qrdata[-5:-1]),
254                     'faculty_id': None,
255                     'sid': None
256                     }
257         else:
258             data=qrdata.split(',')
259             retval={'exam_id': int(data[1]),
260                     'page_no': int(data[3]),
261                     'paper_id':int(data[2]),
262                     'faculty_id':int(data[0]),
263             }
264             if(len(data)>4):
265                 retval['sid']=data[4]
266
267             return retval
268
269     def get_paper_ocr_data(self):
270         data=self.get_code_data()
271         data['qr']=self.QRData
272         data['errors']=self.errors
273         data['warnings']=self.warnings
274         data['up_position']=(list(self.xMarkerLocations[1]/self.imgWidth), list(self.yMarkerLocations[1]/self.imgHeight))
275         data['right_position']=(list(self.xMarkerLocations[1]/self.imgWidth), list(self.yMarkerLocations[1]/self.imgHeight))
276         data['ans_matrix']=((np.array(self.answerMatrix)>self.settings['answer_treshold'])*1).tolist()
277         if data['sid'] is None:
278             data['sid']=self.get_enhanced_sid()
279         return data