25개 이상의 토픽을 선택하실 수 없습니다. Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

rectify.py 6.1 KiB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173
  1. #!/usr/bin/python3
  2. from multiprocessing import Pool, cpu_count
  3. import sys
  4. import re
  5. import numpy
  6. from math import atan,sin,cos,sqrt,tan,acos,ceil
  7. from PIL import Image
  8. EARTH_RADIUS = 6371.0
  9. SAT_HEIGHT = 822.5
  10. SAT_ORBIT_RADIUS = EARTH_RADIUS + SAT_HEIGHT
  11. SWATH_KM = 2800.0
  12. THETA_C = SWATH_KM / EARTH_RADIUS
  13. # Note: theta_s is the satellite viewing angle, theta_c is the angle between the projection of the satellite on the
  14. # Earth's surface and the point the satellite is looking at, measured at the center of the Earth
  15. # Compute the satellite angle of view given the center angle
  16. def theta_s(theta_c):
  17. return atan(EARTH_RADIUS * sin(theta_c)/(SAT_HEIGHT+EARTH_RADIUS*(1-cos(theta_c))))
  18. # Compute the inverse of the function above
  19. def theta_c(theta_s):
  20. delta_sqrt = sqrt(EARTH_RADIUS**2 + tan(theta_s)**2 *
  21. (EARTH_RADIUS**2-SAT_ORBIT_RADIUS**2))
  22. return acos((tan(theta_s)**2*SAT_ORBIT_RADIUS+delta_sqrt)/(EARTH_RADIUS*(tan(theta_s)**2+1)))
  23. # The nightmare fuel that is the correction factor function.
  24. # It is the reciprocal of d/d(theta_c) of theta_s(theta_c) a.k.a.
  25. # the derivative of the inverse of theta_s(theta_c)
  26. def correction_factor(theta_c):
  27. norm_factor = EARTH_RADIUS/SAT_HEIGHT
  28. tan_derivative_recip = (
  29. 1+(EARTH_RADIUS*sin(theta_c)/(SAT_HEIGHT+EARTH_RADIUS*(1-cos(theta_c))))**2)
  30. arg_derivative_recip = (SAT_HEIGHT+EARTH_RADIUS*(1-cos(theta_c)))**2/(EARTH_RADIUS*cos(
  31. theta_c)*(SAT_HEIGHT+EARTH_RADIUS*(1-cos(theta_c)))-EARTH_RADIUS**2*sin(theta_c)**2)
  32. return norm_factor * tan_derivative_recip * arg_derivative_recip
  33. # Radians position given the absolute x pixel position, assuming that the sensor samples the Earth
  34. # surface with a constant angular step
  35. def theta_center(img_size, x):
  36. ts = theta_s(THETA_C/2.0) * (abs(x-img_size/2.0) / (img_size/2.0))
  37. return theta_c(ts)
  38. # Worker thread
  39. def wthread(rectified_width, corr, endrow, startrow):
  40. # Make temporary working img to push pixels onto
  41. working_img = Image.new(img.mode, (rectified_width, img.size[1]))
  42. rectified_pixels = working_img.load()
  43. for row in range(startrow, endrow):
  44. # First pass: stretch from the center towards the right side of the image
  45. start_px = orig_pixels[img.size[0]/2, row]
  46. cur_col = int(rectified_width/2)
  47. target_col = cur_col
  48. for col in range(int(img.size[0]/2), img.size[0]):
  49. target_col += corr[col]
  50. end_px = orig_pixels[col, row]
  51. delta = int(target_col) - cur_col
  52. # Linearly interpolate
  53. for i in range(delta):
  54. interp_r = int((start_px[0]*(delta-i) + end_px[0]*i) / delta)
  55. interp_g = int((start_px[1]*(delta-i) + end_px[1]*i) / delta)
  56. interp_b = int((start_px[2]*(delta-i) + end_px[2]*i) / delta)
  57. rectified_pixels[cur_col, row] = (interp_r, interp_g, interp_b)
  58. cur_col += 1
  59. start_px = end_px
  60. # First pass: stretch from the center towards the left side of the image
  61. start_px = orig_pixels[img.size[0]/2, row]
  62. cur_col = int(rectified_width/2)
  63. target_col = cur_col
  64. for col in range(int(img.size[0]/2)-1, -1, -1):
  65. target_col -= corr[col]
  66. end_px = orig_pixels[col, row]
  67. delta = cur_col - int(target_col)
  68. # Linearly interpolate
  69. for i in range(delta):
  70. interp_r = int((start_px[0]*(delta-i) + end_px[0]*i) / delta)
  71. interp_g = int((start_px[1]*(delta-i) + end_px[1]*i) / delta)
  72. interp_b = int((start_px[2]*(delta-i) + end_px[2]*i) / delta)
  73. rectified_pixels[cur_col, row] = (interp_r, interp_g, interp_b)
  74. cur_col -= 1
  75. start_px = end_px
  76. # Crop the portion we worked on
  77. slice = working_img.crop(box=(0, startrow, rectified_width, endrow))
  78. # Convert to a numpy array so STUPID !#$&ING PICKLE WILL WORK
  79. out = numpy.array(slice)
  80. # Make dict of important values, return that.
  81. return {"offs": startrow, "offe": endrow, "pixels": out}
  82. if __name__ == "__main__":
  83. if len(sys.argv) < 2:
  84. print("Usage: {} <input file>".format(sys.argv[0]))
  85. sys.exit(1)
  86. out_fname = re.sub("\..*$", "-rectified", sys.argv[1])
  87. img = Image.open(sys.argv[1])
  88. print("Opened {}x{} image".format(img.size[0], img.size[1]))
  89. # Precompute the correction factors
  90. corr = []
  91. for i in range(img.size[0]):
  92. corr.append(correction_factor(theta_center(img.size[0], i)))
  93. # Estimate the width of the rectified image
  94. rectified_width = ceil(sum(corr))
  95. # Make new image
  96. rectified_img = Image.new(img.mode, (rectified_width, img.size[1]))
  97. # Get the pixel 2d arrays from the source image
  98. orig_pixels = img.load()
  99. # Callback function to modify the new image
  100. def modimage(data):
  101. if data:
  102. # Write slice to the new image in the right place
  103. rectified_img.paste(Image.fromarray(
  104. data["pixels"]), box=(0, data["offs"]))
  105. # Number of workers to be spawned - Probably best to not overdo this...
  106. numworkers = cpu_count()
  107. # Estimate the number of rows per worker
  108. wrows = ceil(img.size[1]/numworkers)
  109. # Initialize some starting data
  110. startrow = 0
  111. endrow = wrows
  112. # Make out process pool
  113. p = Pool(processes=numworkers)
  114. # Let's have a pool party! Only wnum workers are invited, though.
  115. for wnum in range(numworkers):
  116. # Make the workers with appropriate arguments, pass callback method to actually write data.
  117. p.apply_async(wthread, (rectified_width, corr,
  118. endrow, startrow), callback=modimage)
  119. # Aparrently ++ doesn't work?
  120. wnum = wnum+1
  121. # Beginning of next worker is the end of this one
  122. startrow = wrows*wnum
  123. # End of the worker is the specified number of rows past the beginning
  124. endrow = startrow + wrows
  125. # Show how many processes we're making!
  126. print("Spawning process ", wnum)
  127. # Pool's closed, boys
  128. p.close()
  129. # It's a dead pool now
  130. p.join()
  131. rectified_img.save(out_fname + ".jpg", "JPEG", quality=90)