Add parallel processing

Runs in parallel, corresponding the number of CPU cores on the computer.
Each process handles a single image file.

rok-reader.py

@@ -1,9 +1,12 @@
 import argparse
 import glob
 import os
+import multiprocessing
 from PIL import Image, ImageFilter
 from PIL.ImageOps import autocontrast, invert, grayscale, contain
 import pytesseract
+import signal
+import sys
 
 # ----
 # Program
@@ -85,35 +88,84 @@ KILLS_TARGETS = [
 # Functions
 # ----
 
+# Read an image file
+def read_file(fileTuple):
+    fileNumber = fileTuple[0]
+    file = fileTuple[1]
+    isDuplicate = fileTuple[2]
+
+    filename = os.path.basename(file)
+
+    if "profile" in filename:
+        targets = PROFILE_TARGETS
+        outputPath = OUTPUT_PATH_PROFILE
+    elif "more" in filename:
+        targets = MOREINFO_TARGETS
+        outputPath = OUTPUT_PATH_MOREINFO
+    elif "kills" in filename:
+        targets = KILLS_TARGETS
+        outputPath = OUTPUT_PATH_KILLS
+    else:
+        sys.exit("File name doesn't contain type") # TODO: fix
+
+    if not isDuplicate or arguments.debug:
+        # Open image and swap to RGB
+        image = Image.open(file)
+        rgbImage = Image.new("RGB", image.size, (255, 255, 255))
+        rgbImage.paste(image, mask = image.split()[3])
+        image.close()
+
+        # Get data
+        outputLine = filename + "\t"
+        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
+        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(file, box, is_number, inv, bonusRightTrim, debugFilePath):
-    with Image.open(file) as image:
+def read_string_from_image(rgbImage, box, is_number, inv, bonusRightTrim, debugFilePath):
         # Crop to correct dimentions
-        image = image.crop((box.x, box.y, box.x2, box.y2))
-
-        # Switch to RGB mode
-        rgbimage = Image.new("RGB", image.size, (255, 255, 255))
-        rgbimage.paste(image, mask = image.split()[3])
+        rgbImage = rgbImage.crop((box.x, box.y, box.x2, box.y2))
 
         # Invert if flagged
-        if inv: rgbimage = invert(rgbimage)
+        if inv: rgbImage = invert(rgbImage)
 
         # Apply filters
-        rgbimage = grayscale(rgbimage)
-        rgbimage = autocontrast(rgbimage, cutoff=(0, 75))
-        bbox = autocontrast(invert(rgbimage), cutoff=(0, 90)).getbbox()
-        if bbox: rgbimage = rgbimage.crop((bbox[0], bbox[1], bbox[2] + bonusRightTrim, bbox[3]))
-        rgbimage = contain(rgbimage, (800, 800), method=1)
-        rgbimage = rgbimage.filter(ImageFilter.EDGE_ENHANCE_MORE)
-        rgbimage = rgbimage.filter(ImageFilter.SHARPEN)
+        rgbImage = grayscale(rgbImage)
+        rgbImage = autocontrast(rgbImage, cutoff=(0, 75))
+
+        # Crop to content
+        bbox = autocontrast(invert(rgbImage), cutoff=(0, 90)).getbbox()
+        if bbox: rgbImage = rgbImage.crop((bbox[0], bbox[1], bbox[2] + bonusRightTrim, bbox[3]))
+
+        # Resize and sharpen
+        rgbImage = contain(rgbImage, (800, 800), method=1)
+        rgbImage = rgbImage.filter(ImageFilter.EDGE_ENHANCE_MORE)
+        rgbImage = rgbImage.filter(ImageFilter.SHARPEN)
 
         if arguments.debug:
-            rgbimage.save(debugFilePath)
+            rgbImage.save(debugFilePath)
 
         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(' ', '')
+            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: 
-            return pytesseract.image_to_string(rgbimage, config="--psm 6").strip().replace('\n', ' ').replace('\r', '').replace('\t', ' ')
+            return pytesseract.image_to_string(rgbImage, config="--psm 6").strip().replace('\n', ' ').replace('\r', '').replace('\t', ' ')
+
+# Write to output file
+def write_file():
+    return
+
+# Initialize child processes (ignore SIGINT)
+def mpInitializer():
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
 
 # ----
 # Arguments
@@ -131,52 +183,58 @@ arguments = parser.parse_args()
 # Program
 # ----
 
+# TODO: remove globals
+debugFolder = "debug" + "/"
+
 if __name__ == '__main__':
     # Create project folder
     projectFolder = "output/" + arguments.project + "/"
     if not os.path.exists(projectFolder):
         os.makedirs(projectFolder)
 
-    debugFolder = "debug" + "/"
+    # Create debug folder
     if arguments.debug:
         if not os.path.exists(debugFolder):
             os.makedirs(debugFolder)
 
-    # Get files
+    # Get files to read
     screenshots_to_read = glob.glob(arguments.file, recursive=True)
     screenshot_count = len(screenshots_to_read)
     if screenshot_count < 1: sys.exit("No files found.")
 
-    # Scrape
-    if arguments.verbose: print("Scraping", screenshot_count, "files")
+    # Get all previously scraped data # TODO: limit to filenames only
+    alreadyScraped = ""
+    for outputPath in [OUTPUT_PATH_PROFILE, OUTPUT_PATH_MOREINFO, OUTPUT_PATH_KILLS]:
+        with open(projectFolder + outputPath, "r+", newline='', encoding='utf-8') as outputFile:
+            alreadyScraped = alreadyScraped + outputFile.read()
+
+    # Mark as duplicates
     for i, file in enumerate(screenshots_to_read):
-        filename = os.path.basename(file)
-
-        if arguments.verbose: print(i+1, "/", screenshot_count, ": ", filename, sep="")
-
-        if "profile" in filename:
-            targets = PROFILE_TARGETS
-            output = OUTPUT_PATH_PROFILE
-        elif "more" in filename:
-            targets = MOREINFO_TARGETS
-            output = OUTPUT_PATH_MOREINFO
-        elif "kills" in filename:
-            targets = KILLS_TARGETS
-            output = OUTPUT_PATH_KILLS
+        if os.path.basename(file) in alreadyScraped:
+            screenshots_to_read[i] = (i, file, True)
         else:
-            sys.exit("File name doesn't contain type")
-
-        # TODO: bad time complexity
-        exists = False
-        with open(projectFolder + output, "a+", newline='', encoding='utf-8') as output_file:
-            if filename not in output_file.read():
-                if not arguments.debug: output_file.write(filename + "\t")
-                for i, target in enumerate(targets):
-                    debugFile = os.path.splitext(filename)[0] + "_" + str(i) + ".png"
-                    string = read_string_from_image(file, Box(target[1][0], target[1][1], target[2][0], target[2][1]), target[3], target[4], target[5], debugFolder + debugFile)
-                    if arguments.verbose: print("  ", target[0], ": ", string, sep="")
-                    if i and not arguments.debug: output_file.write("\t")
-                    if not arguments.debug: output_file.write(string)
-                if not arguments.debug: output_file.write("\n")
-            else:
+            screenshots_to_read[i] = (i, file, False)
+
+    # Scrape
+    if arguments.verbose: print("Scraping", screenshot_count, "files")
+
+    cpuCount = multiprocessing.cpu_count()
+    mpPool = multiprocessing.Pool(cpuCount, initializer=mpInitializer)
+
+    try:
+        # Returns: (fileNumber, filename, debugOutput, outputPath, outputLine, isDuplicate)
+        for result in mpPool.imap(read_file, screenshots_to_read):
+            if arguments.verbose: print(result[0]+1, "/", screenshot_count, ": ", result[1], sep="")
+            if result[5] and not arguments.debug:
                 if arguments.verbose: print("  ", "already scraped.")
+            else:
+                if arguments.verbose: print(result[2])
+                if not arguments.debug:
+                    with open(projectFolder + result[3], "a+", newline='', encoding='utf-8') as outputFile:
+                        outputFile.write(result[4] + "\n")
+    except KeyboardInterrupt:
+        print("Exiting...")
+        mpPool.terminate()
+    finally:
+        mpPool.terminate()
+        mpPool.join()