Browse Source

Add bounding box debug image

main
Ashton Charbonneau 3 years ago
parent
commit
99c74903b7
1 changed files with 26 additions and 25 deletions
  1. 26
    25
      rok-reader.py

+ 26
- 25
rok-reader.py View File

4
 import os
4
 import os
5
 import multiprocessing
5
 import multiprocessing
6
 import numpy
6
 import numpy
7
-from PIL import Image, ImageFilter
8
-from PIL.ImageOps import autocontrast, invert, grayscale, contain, expand
7
+from PIL import Image, ImageDraw, ImageFilter
8
+from PIL.ImageOps import autocontrast, contain, expand, grayscale, invert
9
 import pytesseract
9
 import pytesseract
10
 import signal
10
 import signal
11
 import sys
11
 import sys
49
 # Coordinates
49
 # Coordinates
50
 # ----
50
 # ----
51
 
51
 
52
-# Name, Top Left, Bottom Right, Number, Invert, BonusRightTrim
52
+# Name, Top Left, Bottom Right, IsNumber, Invert, BonusRightTrim
53
 
53
 
54
 PROFILE_TARGETS = [
54
 PROFILE_TARGETS = [
55
-    ("ID", (1246, 375), (1445, 430), True, True, -10),
55
+    ("ID", (1242, 375), (1445, 430), True, True, -10),
56
     ("Power", (1435, 585), (1733, 634), True, True, 0),
56
     ("Power", (1435, 585), (1733, 634), True, True, 0),
57
     ("Kill Points", (1806, 585), (2112, 633), True, True, 0),
57
     ("Kill Points", (1806, 585), (2112, 633), True, True, 0),
58
     ("Alliance", (1025, 584), (1427, 637), False, True, 0),
58
     ("Alliance", (1025, 584), (1427, 637), False, True, 0),
61
 
61
 
62
 MOREINFO_TARGETS = [
62
 MOREINFO_TARGETS = [
63
     ("Power", (1305, 223), (1540, 274), True, True, 0),
63
     ("Power", (1305, 223), (1540, 274), True, True, 0),
64
-    ("Kill Points", (1931, 222), (2188, 276), True, True, 0),
64
+    ("Kill Points", (1933, 222), (2188, 276), True, True, 0),
65
     ("Highest Power", (1815, 416), (2105, 483), True, True, 0),
65
     ("Highest Power", (1815, 416), (2105, 483), True, True, 0),
66
     ("Victories", (1815, 515), (2105, 580), True, True, 0),
66
     ("Victories", (1815, 515), (2105, 580), True, True, 0),
67
     ("Defeats", (1815, 613), (2105, 675), True, True, 0),
67
     ("Defeats", (1815, 613), (2105, 675), True, True, 0),
74
 
74
 
75
 KILLS_TARGETS = [
75
 KILLS_TARGETS = [
76
     ("Kill Points", (1418, 312), (1694, 352), True, False, 0),
76
     ("Kill Points", (1418, 312), (1694, 352), True, False, 0),
77
-    ("T1 Kills", (1325, 637), (1538, 684), True, False, 0),
77
+    ("T1 Kills", (1323, 637), (1538, 684), True, False, 0),
78
     ("T1 Kill Points", (1986, 637), (2212, 684), True, False, 0),
78
     ("T1 Kill Points", (1986, 637), (2212, 684), True, False, 0),
79
-    ("T2 Kills", (1325, 702), (1538, 755), True, False, 0),
79
+    ("T2 Kills", (1323, 702), (1538, 755), True, False, 0),
80
     ("T2 Kill Points", (1986, 702), (2212, 755), True, False, 0),
80
     ("T2 Kill Points", (1986, 702), (2212, 755), True, False, 0),
81
-    ("T3 Kills", (1325, 770), (1538, 824), True, False, 0),
81
+    ("T3 Kills", (1323, 770), (1538, 824), True, False, 0),
82
     ("T3 Kill Points", (1986, 770), (2212, 824), True, False, 0),
82
     ("T3 Kill Points", (1986, 770), (2212, 824), True, False, 0),
83
-    ("T4 Kills", (1325, 847), (1538, 897), True, False, 0),
83
+    ("T4 Kills", (1323, 847), (1538, 897), True, False, 0),
84
     ("T4 Kill Points", (1986, 847), (2212, 897), True, False, 0),
84
     ("T4 Kill Points", (1986, 847), (2212, 897), True, False, 0),
85
-    ("T5 Kills", (1325, 918), (1538, 968), True, False, 0),
85
+    ("T5 Kills", (1323, 918), (1538, 968), True, False, 0),
86
     ("T5 Kill Points", (1986, 918), (2212, 968), True, False, 0),
86
     ("T5 Kill Points", (1986, 918), (2212, 968), True, False, 0),
87
-    ("Previous Kills", (1626, 985), (2228, 1039), False, False, -385)
87
+    ("Previous Kills", (1626, 985), (2228, 1039), True, False, -385)
88
 ]
88
 ]
89
 
89
 
90
 # ----
90
 # ----
116
         image = Image.open(file)
116
         image = Image.open(file)
117
         rgbImage = Image.new("RGB", image.size, (255, 255, 255))
117
         rgbImage = Image.new("RGB", image.size, (255, 255, 255))
118
         rgbImage.paste(image, mask = image.split()[3])
118
         rgbImage.paste(image, mask = image.split()[3])
119
+
120
+        if arguments.debug:
121
+            debugImage = Image.new("RGB", rgbImage.size, (255, 255, 255))
122
+            debugImage.paste(rgbImage, mask = image.split()[3])
123
+            draw = ImageDraw.Draw(debugImage)
119
         image.close()
124
         image.close()
120
 
125
 
121
         # Get data
126
         # Get data
123
         debugOutput = ""
128
         debugOutput = ""
124
         for i, target in enumerate(targets):
129
         for i, target in enumerate(targets):
125
             debugFile = os.path.splitext(filename)[0] + "_" + str(i) + ".png"
130
             debugFile = os.path.splitext(filename)[0] + "_" + str(i) + ".png"
126
-
127
             string = read_string_from_image(rgbImage, Box(target[1][0], target[1][1], target[2][0], target[2][1]), target[3], target[4], target[5], debugFolder + debugFile)
131
             string = read_string_from_image(rgbImage, Box(target[1][0], target[1][1], target[2][0], target[2][1]), target[3], target[4], target[5], debugFolder + debugFile)
128
             debugOutput = debugOutput + "  " + target[0] + ": " + string + "\n"
132
             debugOutput = debugOutput + "  " + target[0] + ": " + string + "\n"
129
             if i:
133
             if i:
130
                 outputLine = outputLine + "\t"
134
                 outputLine = outputLine + "\t"
131
             outputLine = outputLine + string
135
             outputLine = outputLine + string
136
+            if arguments.debug: draw.rectangle([target[1][0], target[1][1], target[2][0], target[2][1]], outline="rgb(255,0,0)")
137
+        if arguments.debug: debugImage.save(debugFolder + os.path.splitext(filename)[0] + "_debug" + ".png")
132
         return (fileNumber, filename, debugOutput, outputPath, outputLine, isDuplicate)
138
         return (fileNumber, filename, debugOutput, outputPath, outputLine, isDuplicate)
133
     else:
139
     else:
134
         return (fileNumber, filename, "", outputPath, "", isDuplicate)
140
         return (fileNumber, filename, "", outputPath, "", isDuplicate)
135
 
141
 
136
 # Read text from a section of an image using Tesseract
142
 # Read text from a section of an image using Tesseract
137
 def read_string_from_image(rgbImage, box, is_number, inv, bonusRightTrim, debugFilePath):
143
 def read_string_from_image(rgbImage, box, is_number, inv, bonusRightTrim, debugFilePath):
138
-        # Crop to correct dimensions
144
+        # Crop to correct dimensions and invert if flagged
139
         rgbImage = rgbImage.crop((box.x, box.y, box.x2, box.y2))
145
         rgbImage = rgbImage.crop((box.x, box.y, box.x2, box.y2))
140
-
141
-        # Invert if flagged
142
         if inv: rgbImage = invert(rgbImage)
146
         if inv: rgbImage = invert(rgbImage)
143
 
147
 
144
         # Apply Pillow filters to cut off artifacts
148
         # Apply Pillow filters to cut off artifacts
145
         rgbImage = autocontrast(rgbImage, cutoff=(0, 50))
149
         rgbImage = autocontrast(rgbImage, cutoff=(0, 50))
146
 
150
 
147
-        # Convert to OpenCV
148
-        npImage=numpy.array(rgbImage)
149
-
150
-        # Set colors to grayscale
151
-        npImage=cv2.cvtColor(npImage, cv2.COLOR_BGR2GRAY)
151
+        # Convert to OpenCV and swap to grayscale
152
+        npImage = numpy.array(rgbImage)
153
+        npImage = cv2.cvtColor(npImage, cv2.COLOR_BGR2GRAY)
152
 
154
 
153
         # Apply OpenCV Filters
155
         # Apply OpenCV Filters
154
         npImage = cv2.medianBlur(npImage, 3)
156
         npImage = cv2.medianBlur(npImage, 3)
155
         _, npImage = cv2.threshold(npImage, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
157
         _, npImage = cv2.threshold(npImage, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
156
         # npImage = cv2.adaptiveThreshold(npImage, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
158
         # npImage = cv2.adaptiveThreshold(npImage, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
157
 
159
 
158
-        # Convert to Pillow
160
+        # Convert back to Pillow and crop to content
159
         rgbImage = Image.fromarray(npImage)
161
         rgbImage = Image.fromarray(npImage)
160
-
161
-        # Crop to content
162
         bbox = invert(rgbImage).getbbox()
162
         bbox = invert(rgbImage).getbbox()
163
         if bbox: rgbImage = rgbImage.crop((bbox[0], bbox[1], bbox[2] + bonusRightTrim, bbox[3]))
163
         if bbox: rgbImage = rgbImage.crop((bbox[0], bbox[1], bbox[2] + bonusRightTrim, bbox[3]))
164
         if bbox: rgbImage = rgbImage.crop(invert(rgbImage).getbbox())
164
         if bbox: rgbImage = rgbImage.crop(invert(rgbImage).getbbox())
165
         rgbImage = expand(rgbImage, border=10, fill=255)
165
         rgbImage = expand(rgbImage, border=10, fill=255)
166
 
166
 
167
-        if arguments.debug:
168
-            rgbImage.save(debugFilePath)
167
+        # Save debug image
168
+        if arguments.debug: rgbImage.save(debugFilePath)
169
 
169
 
170
+        # Run Tesseract (depends whether it is a number or not)
170
         if is_number:
171
         if is_number:
171
             return pytesseract.image_to_string(rgbImage, config="--psm 6 -c tessedit_char_whitelist=0123456789,").strip().replace('\n', ' ').replace('\r', '').replace('.', '').replace(',', '').replace('\t', ' ').replace(' ', '')
172
             return pytesseract.image_to_string(rgbImage, config="--psm 6 -c tessedit_char_whitelist=0123456789,").strip().replace('\n', ' ').replace('\r', '').replace('.', '').replace(',', '').replace('\t', ' ').replace(' ', '')
172
         else: 
173
         else: 

Loading…
Cancel
Save