diff --git a/rectify.py b/rectify.py
index 3fd9897..1b95e82 100755
--- a/rectify.py
+++ b/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: {} ".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: {} ".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)