diff --git a/rectify.py b/rectify.py index f7ed429..d25581e 100755 --- a/rectify.py +++ b/rectify.py @@ -1,10 +1,14 @@ #!/usr/bin/python3 +from multiprocessing import Pool, cpu_count import sys import re +import numpy from math import * from PIL import Image + + EARTH_RADIUS = 6371.0 SAT_HEIGHT = 822.5 SAT_ORBIT_RADIUS = EARTH_RADIUS + SAT_HEIGHT @@ -38,36 +42,14 @@ def correction_factor(theta_c): 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) - - - -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)) - + +# Worker thread +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) @@ -84,7 +66,7 @@ def theta_center(img_size, x): interp_g = int((start_px[1]*(delta-i) + end_px[1]*i) / delta) interp_b = int((start_px[2]*(delta-i) + end_px[2]*i) / delta) - rectified_pixels[cur_col,row] = (interp_r, interp_g, interp_b) + rectified_pixels[cur_col, row] = (interp_r, interp_g, interp_b) cur_col += 1 start_px = end_px @@ -105,12 +87,80 @@ def theta_center(img_size, x): interp_g = int((start_px[1]*(delta-i) + end_px[1]*i) / delta) interp_b = int((start_px[2]*(delta-i) + end_px[2]*i) / delta) - rectified_pixels[cur_col,row] = (interp_r, interp_g, interp_b) + rectified_pixels[cur_col, row] = (interp_r, interp_g, interp_b) cur_col -= 1 start_px = end_px + + # Convert to a numpy array so STUPID !#$&ING PICKLE WILL WORK + out = numpy.array(working_img) + # 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: + # Let's make a float! + oset = "0." + str(data["offs"]) + oset = float(oset) + # Get the temp image from numpy array (PICKLE IS STILL STUPID) + working_img = Image.fromarray(data["pixels"]) + # Crop the portion we worked on + slice = working_img.crop(box=(0, data["offs"], rectified_width, data["offe"])) + # Write it to the new image in the right place + rectified_img.paste(slice, 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() + + print("Writing rectified image to {}".format(out_fname + ".png")) rectified_img.save(out_fname + ".png", "PNG") -