multi threading

rectify.py

@@ -1,7 +1,9 @@
 #!/usr/bin/python3
 
+from multiprocessing import Pool, cpu_count
 import sys
 import re
+import numpy
 from math import atan,sin,cos,sqrt,tan,acos,ceil
 from PIL import Image
 
@@ -50,34 +52,15 @@ def theta_center(img_size, x):
     ts = theta_s(THETA_C/2.0) * (abs(x-img_size/2.0) / (img_size/2.0))
     return theta_c(ts)
 
+# Worker thread
 
-if __name__ == "__main__":
-    if len(sys.argv) < 2:
-        print("Usage: {} <input file>".format(sys.argv[0]))
-        sys.exit(1)
-
-    out_fname = re.sub("\..*$", "-rectified", sys.argv[1])
 
-    img = Image.open(sys.argv[1])
-    print("Opened {}x{} image".format(img.size[0], img.size[1]))
-
-    # Precompute the correction factors
-    corr = []
-    for i in range(img.size[0]):
-        corr.append(correction_factor(theta_center(img.size[0], i)))
-
-    # Estimate the width of the rectified image
-    rectified_width = ceil(sum(corr))
-    rectified_img = Image.new(img.mode, (rectified_width, img.size[1]))
-
-    # Get the pixel 2d arrays from both the source image and the target image
-    orig_pixels = img.load()
-    rectified_pixels = rectified_img.load()
-
-    for row in range(img.size[1]):
-        if row % 20 == 0:
-            print("Row {}".format(row))
+def wthread(rectified_width, corr, endrow, startrow):
+    # Make temporary working img to push pixels onto
+    working_img = Image.new(img.mode, (rectified_width, img.size[1]))
+    rectified_pixels = working_img.load()
 
+    for row in range(startrow, endrow):
         # First pass: stretch from the center towards the right side of the image
         start_px = orig_pixels[img.size[0]/2, row]
         cur_col = int(rectified_width/2)
@@ -120,5 +103,71 @@ if __name__ == "__main__":
 
             start_px = end_px
 
-    print("Writing rectified image to {}".format(out_fname + ".jpg"))
+    # Crop the portion we worked on
+    slice = working_img.crop(box=(0, startrow, rectified_width, endrow))
+    # Convert to a numpy array so STUPID !#$&ING PICKLE WILL WORK
+    out = numpy.array(slice)
+    # Make dict of important values, return that.
+    return {"offs": startrow, "offe": endrow, "pixels": out}
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: {} <input file>".format(sys.argv[0]))
+        sys.exit(1)
+
+    out_fname = re.sub("\..*$", "-rectified", sys.argv[1])
+
+    img = Image.open(sys.argv[1])
+    print("Opened {}x{} image".format(img.size[0], img.size[1]))
+
+    # Precompute the correction factors
+    corr = []
+    for i in range(img.size[0]):
+        corr.append(correction_factor(theta_center(img.size[0], i)))
+
+    # Estimate the width of the rectified image
+    rectified_width = ceil(sum(corr))
+
+    # Make new image
+    rectified_img = Image.new(img.mode, (rectified_width, img.size[1]))
+
+    # Get the pixel 2d arrays from the source image
+    orig_pixels = img.load()
+
+    # Callback function to modify the new image
+    def modimage(data):
+        if data:
+            # Write slice to the new image in the right place
+            rectified_img.paste(Image.fromarray(
+                data["pixels"]), box=(0, data["offs"]))
+
+    # Number of workers to be spawned - Probably best to not overdo this...
+    numworkers = cpu_count()
+    # Estimate the number of rows per worker
+    wrows = ceil(img.size[1]/numworkers)
+    # Initialize some starting data
+    startrow = 0
+    endrow = wrows
+    # Make out process pool
+    p = Pool(processes=numworkers)
+
+    # Let's have a pool party! Only wnum workers are invited, though.
+    for wnum in range(numworkers):
+        # Make the workers with appropriate arguments, pass callback method to actually write data.
+        p.apply_async(wthread, (rectified_width, corr,
+                                endrow, startrow), callback=modimage)
+        # Aparrently ++ doesn't work?
+        wnum = wnum+1
+        # Beginning of next worker is the end of this one
+        startrow = wrows*wnum
+        # End of the worker is the specified number of rows past the beginning
+        endrow = startrow + wrows
+        # Show how many processes we're making!
+        print("Spawning process ", wnum)
+    # Pool's closed, boys
+    p.close()
+    # It's a dead pool now
+    p.join()
+
     rectified_img.save(out_fname + ".jpg", "JPEG", quality=90)