Przeglądaj źródła

Add bounding box debug image

main
Ashton Charbonneau 3 lat temu
rodzic
commit
99c74903b7
1 zmienionych plików z 26 dodań i 25 usunięć
  1. 26
    25
      rok-reader.py

+ 26
- 25
rok-reader.py Wyświetl plik

@@ -4,8 +4,8 @@ import glob
4 4
 import os
5 5
 import multiprocessing
6 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 9
 import pytesseract
10 10
 import signal
11 11
 import sys
@@ -49,10 +49,10 @@ OUTPUT_PATH_KILLS = "output-kills.csv"
49 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 54
 PROFILE_TARGETS = [
55
-    ("ID", (1246, 375), (1445, 430), True, True, -10),
55
+    ("ID", (1242, 375), (1445, 430), True, True, -10),
56 56
     ("Power", (1435, 585), (1733, 634), True, True, 0),
57 57
     ("Kill Points", (1806, 585), (2112, 633), True, True, 0),
58 58
     ("Alliance", (1025, 584), (1427, 637), False, True, 0),
@@ -61,7 +61,7 @@ PROFILE_TARGETS = [
61 61
 
62 62
 MOREINFO_TARGETS = [
63 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 65
     ("Highest Power", (1815, 416), (2105, 483), True, True, 0),
66 66
     ("Victories", (1815, 515), (2105, 580), True, True, 0),
67 67
     ("Defeats", (1815, 613), (2105, 675), True, True, 0),
@@ -74,17 +74,17 @@ MOREINFO_TARGETS = [
74 74
 
75 75
 KILLS_TARGETS = [
76 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 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 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 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 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 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,6 +116,11 @@ def read_file(fileTuple):
116 116
         image = Image.open(file)
117 117
         rgbImage = Image.new("RGB", image.size, (255, 255, 255))
118 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 124
         image.close()
120 125
 
121 126
         # Get data
@@ -123,50 +128,46 @@ def read_file(fileTuple):
123 128
         debugOutput = ""
124 129
         for i, target in enumerate(targets):
125 130
             debugFile = os.path.splitext(filename)[0] + "_" + str(i) + ".png"
126
-
127 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 132
             debugOutput = debugOutput + "  " + target[0] + ": " + string + "\n"
129 133
             if i:
130 134
                 outputLine = outputLine + "\t"
131 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 138
         return (fileNumber, filename, debugOutput, outputPath, outputLine, isDuplicate)
133 139
     else:
134 140
         return (fileNumber, filename, "", outputPath, "", isDuplicate)
135 141
 
136 142
 # Read text from a section of an image using Tesseract
137 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 145
         rgbImage = rgbImage.crop((box.x, box.y, box.x2, box.y2))
140
-
141
-        # Invert if flagged
142 146
         if inv: rgbImage = invert(rgbImage)
143 147
 
144 148
         # Apply Pillow filters to cut off artifacts
145 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 155
         # Apply OpenCV Filters
154 156
         npImage = cv2.medianBlur(npImage, 3)
155 157
         _, npImage = cv2.threshold(npImage, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
156 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 161
         rgbImage = Image.fromarray(npImage)
160
-
161
-        # Crop to content
162 162
         bbox = invert(rgbImage).getbbox()
163 163
         if bbox: rgbImage = rgbImage.crop((bbox[0], bbox[1], bbox[2] + bonusRightTrim, bbox[3]))
164 164
         if bbox: rgbImage = rgbImage.crop(invert(rgbImage).getbbox())
165 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 171
         if is_number:
171 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 173
         else: 

Ładowanie…
Anuluj
Zapisz