Add bounding box debug image

rok-reader.py

@@ -4,8 +4,8 @@ import glob
 import os
 import multiprocessing
 import numpy
-from PIL import Image, ImageFilter
-from PIL.ImageOps import autocontrast, invert, grayscale, contain, expand
+from PIL import Image, ImageDraw, ImageFilter
+from PIL.ImageOps import autocontrast, contain, expand, grayscale, invert
 import pytesseract
 import signal
 import sys
@@ -49,10 +49,10 @@ OUTPUT_PATH_KILLS = "output-kills.csv"
 # Coordinates
 # ----
 
-# Name, Top Left, Bottom Right, Number, Invert, BonusRightTrim
+# Name, Top Left, Bottom Right, IsNumber, Invert, BonusRightTrim
 
 PROFILE_TARGETS = [
-    ("ID", (1246, 375), (1445, 430), True, True, -10),
+    ("ID", (1242, 375), (1445, 430), True, True, -10),
     ("Power", (1435, 585), (1733, 634), True, True, 0),
     ("Kill Points", (1806, 585), (2112, 633), True, True, 0),
     ("Alliance", (1025, 584), (1427, 637), False, True, 0),
@@ -61,7 +61,7 @@ PROFILE_TARGETS = [
 
 MOREINFO_TARGETS = [
     ("Power", (1305, 223), (1540, 274), True, True, 0),
-    ("Kill Points", (1931, 222), (2188, 276), True, True, 0),
+    ("Kill Points", (1933, 222), (2188, 276), True, True, 0),
     ("Highest Power", (1815, 416), (2105, 483), True, True, 0),
     ("Victories", (1815, 515), (2105, 580), True, True, 0),
     ("Defeats", (1815, 613), (2105, 675), True, True, 0),
@@ -74,17 +74,17 @@ MOREINFO_TARGETS = [
 
 KILLS_TARGETS = [
     ("Kill Points", (1418, 312), (1694, 352), True, False, 0),
-    ("T1 Kills", (1325, 637), (1538, 684), True, False, 0),
+    ("T1 Kills", (1323, 637), (1538, 684), True, False, 0),
     ("T1 Kill Points", (1986, 637), (2212, 684), True, False, 0),
-    ("T2 Kills", (1325, 702), (1538, 755), True, False, 0),
+    ("T2 Kills", (1323, 702), (1538, 755), True, False, 0),
     ("T2 Kill Points", (1986, 702), (2212, 755), True, False, 0),
-    ("T3 Kills", (1325, 770), (1538, 824), True, False, 0),
+    ("T3 Kills", (1323, 770), (1538, 824), True, False, 0),
     ("T3 Kill Points", (1986, 770), (2212, 824), True, False, 0),
-    ("T4 Kills", (1325, 847), (1538, 897), True, False, 0),
+    ("T4 Kills", (1323, 847), (1538, 897), True, False, 0),
     ("T4 Kill Points", (1986, 847), (2212, 897), True, False, 0),
-    ("T5 Kills", (1325, 918), (1538, 968), True, False, 0),
+    ("T5 Kills", (1323, 918), (1538, 968), True, False, 0),
     ("T5 Kill Points", (1986, 918), (2212, 968), True, False, 0),
-    ("Previous Kills", (1626, 985), (2228, 1039), False, False, -385)
+    ("Previous Kills", (1626, 985), (2228, 1039), True, False, -385)
 ]
 
 # ----
@@ -116,6 +116,11 @@ def read_file(fileTuple):
         image = Image.open(file)
         rgbImage = Image.new("RGB", image.size, (255, 255, 255))
         rgbImage.paste(image, mask = image.split()[3])
+
+        if arguments.debug:
+            debugImage = Image.new("RGB", rgbImage.size, (255, 255, 255))
+            debugImage.paste(rgbImage, mask = image.split()[3])
+            draw = ImageDraw.Draw(debugImage)
         image.close()
 
         # Get data
@@ -123,50 +128,46 @@ def read_file(fileTuple):
         debugOutput = ""
         for i, target in enumerate(targets):
             debugFile = os.path.splitext(filename)[0] + "_" + str(i) + ".png"
-
             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)
             debugOutput = debugOutput + "  " + target[0] + ": " + string + "\n"
             if i:
                 outputLine = outputLine + "\t"
             outputLine = outputLine + string
+            if arguments.debug: draw.rectangle([target[1][0], target[1][1], target[2][0], target[2][1]], outline="rgb(255,0,0)")
+        if arguments.debug: debugImage.save(debugFolder + os.path.splitext(filename)[0] + "_debug" + ".png")
         return (fileNumber, filename, debugOutput, outputPath, outputLine, isDuplicate)
     else:
         return (fileNumber, filename, "", outputPath, "", isDuplicate)
 
 # Read text from a section of an image using Tesseract
 def read_string_from_image(rgbImage, box, is_number, inv, bonusRightTrim, debugFilePath):
-        # Crop to correct dimensions
+        # Crop to correct dimensions and invert if flagged
         rgbImage = rgbImage.crop((box.x, box.y, box.x2, box.y2))
-
-        # Invert if flagged
         if inv: rgbImage = invert(rgbImage)
 
         # Apply Pillow filters to cut off artifacts
         rgbImage = autocontrast(rgbImage, cutoff=(0, 50))
 
-        # Convert to OpenCV
-        npImage=numpy.array(rgbImage)
-
-        # Set colors to grayscale
-        npImage=cv2.cvtColor(npImage, cv2.COLOR_BGR2GRAY)
+        # Convert to OpenCV and swap to grayscale
+        npImage = numpy.array(rgbImage)
+        npImage = cv2.cvtColor(npImage, cv2.COLOR_BGR2GRAY)
 
         # Apply OpenCV Filters
         npImage = cv2.medianBlur(npImage, 3)
         _, npImage = cv2.threshold(npImage, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
         # npImage = cv2.adaptiveThreshold(npImage, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
 
-        # Convert to Pillow
+        # Convert back to Pillow and crop to content
         rgbImage = Image.fromarray(npImage)
-
-        # Crop to content
         bbox = invert(rgbImage).getbbox()
         if bbox: rgbImage = rgbImage.crop((bbox[0], bbox[1], bbox[2] + bonusRightTrim, bbox[3]))
         if bbox: rgbImage = rgbImage.crop(invert(rgbImage).getbbox())
         rgbImage = expand(rgbImage, border=10, fill=255)
 
-        if arguments.debug:
-            rgbImage.save(debugFilePath)
+        # Save debug image
+        if arguments.debug: rgbImage.save(debugFilePath)
 
+        # Run Tesseract (depends whether it is a number or not)
         if is_number:
             return pytesseract.image_to_string(rgbImage, config="--psm 6 -c tessedit_char_whitelist=0123456789,").strip().replace('\n', ' ').replace('\r', '').replace('.', '').replace(',', '').replace('\t', ' ').replace(' ', '')
         else: