269 lines
No EOL
9.9 KiB
Python
269 lines
No EOL
9.9 KiB
Python
import os
|
|
import time
|
|
import datetime
|
|
from collections import defaultdict
|
|
from multiprocessing import Queue, Process
|
|
|
|
import tqdm
|
|
from PIL import Image, ImageFont, ImageDraw
|
|
import numpy as np
|
|
from matplotlib import pyplot as plt
|
|
|
|
def bin_ndarray(ndarray, new_shape, operation='mean'):
|
|
"""
|
|
Bins an ndarray in all axes based on the target shape, by summing or
|
|
averaging.
|
|
|
|
Number of output dimensions must match number of input dimensions and
|
|
new axes must divide old ones.
|
|
|
|
Example
|
|
-------
|
|
>>> m = np.arange(0,100,1).reshape((10,10))
|
|
>>> n = bin_ndarray(m, new_shape=(5,5), operation='sum')
|
|
>>> print(n)
|
|
|
|
[[ 22 30 38 46 54]
|
|
[102 110 118 126 134]
|
|
[182 190 198 206 214]
|
|
[262 270 278 286 294]
|
|
[342 350 358 366 374]]
|
|
|
|
"""
|
|
operation = operation.lower()
|
|
if not operation in ['sum', 'mean']:
|
|
raise ValueError("Operation not supported.")
|
|
if ndarray.ndim != len(new_shape):
|
|
raise ValueError("Shape mismatch: {} -> {}".format(ndarray.shape,
|
|
new_shape))
|
|
compression_pairs = [(d, c//d) for d,c in zip(new_shape,
|
|
ndarray.shape)]
|
|
flattened = [l for p in compression_pairs for l in p]
|
|
ndarray = ndarray.reshape(flattened)
|
|
for i in range(len(new_shape)):
|
|
op = getattr(ndarray, operation)
|
|
ndarray = op(-1*(i+1))
|
|
return ndarray
|
|
|
|
def gamma_correct(fun):
|
|
def wrapper(*args, **kwargs):
|
|
args = list(args)
|
|
args[0] = np.power(args[0], 2.2)
|
|
args[1] = np.power(args[1], 2.2)
|
|
args = tuple(args)
|
|
result = fun(*args, **kwargs)
|
|
return np.power(result, 1/2.2)
|
|
return wrapper
|
|
|
|
def clip_color(fun):
|
|
def wrapper(*args, **kwargs):
|
|
return np.clip(fun(*args, **kwargs), 0.0, 1.0)
|
|
return wrapper
|
|
|
|
# linear_srgb_matrix = np.array([[0.4124, 0.3576, 0.1805],
|
|
# [0.2126, 0.7152, 0.0722],
|
|
# [0.0193, 0.1192, 0.9505]])
|
|
|
|
# linear_srgb_matrix_inv = np.array([[ 3.2406, -1.5372, -0.4986],
|
|
# [-0.9689, 1.8758, 0.0415],
|
|
# [ 0.0557, -0.2040, 1.0570]])
|
|
|
|
# def linear_color_correction(fun):
|
|
# def wrapper(*args, **kwargs):
|
|
# args = list(args)
|
|
# inds = args[0] <= 0.04045
|
|
# ninds = args[0] > 0.04045
|
|
# for i in range(2):
|
|
# args[i][inds] = args[i][inds] / 12.92
|
|
# args[i][ninds] = np.power((args[i][ninds] + 0.055) / 1.055, 2.4)
|
|
# for x in range(args[i].shape[0]):
|
|
# for y in range(args[i].shape[1]):
|
|
# args[i][x,y,:] = np.matmul(linear_srgb_matrix, args[i][x,y,:])
|
|
# args = tuple(args)
|
|
# result = fun(*args, **kwargs)
|
|
|
|
# for x in range(result.shape[0]):
|
|
# for y in range(result.shape[1]):
|
|
# result[x,y,:] = np.matmul(linear_srgb_matrix_inv, result[x,y,:])
|
|
# inds = result <= 0.0031308
|
|
# ninds = result > 0.0031308
|
|
# result[inds] = result[inds] * 12.92
|
|
# result[ninds] = np.power(result[ninds], 1.0/2.4) * 1.055 - 0.055
|
|
# return result
|
|
# return wrapper
|
|
|
|
|
|
@gamma_correct
|
|
def composite_alpha_over(F, B, alpha_F, alpha_B = 1):
|
|
return (F*alpha_F + B*alpha_B*(1-alpha_F)) / (alpha_F + alpha_B*(1-alpha_F))
|
|
|
|
def composite_alpha_blend(F, B, alpha):
|
|
return F*alpha + B*(1-alpha)
|
|
|
|
def linear_burn(F, B):
|
|
burn = F + B - 1
|
|
burn[burn < 0.0] = 0.0
|
|
return burn
|
|
|
|
def difference(F, B):
|
|
return np.abs(F - B)
|
|
|
|
@clip_color
|
|
def linear_light(F, B):
|
|
result = np.zeros_like(F)
|
|
inds = F <= 0.5
|
|
ninds = F > 0.5
|
|
result[inds] = B[inds] + 2.0 * F[inds] - 1
|
|
result[ninds] = 2.0 * (F[ninds] - 0.5) + B[ninds]
|
|
return result
|
|
|
|
@clip_color
|
|
def hard_light(F, B):
|
|
result = np.zeros_like(F)
|
|
inds = B < 0.5
|
|
ninds = B >= 0.5
|
|
result[inds] = 2 * F[inds] * B[inds]
|
|
result[ninds] = 1 - (2*(1 - F[ninds])*(1 - B[ninds]))
|
|
return result
|
|
|
|
@clip_color
|
|
def color_dodge(F, B):
|
|
return B / (1.000001 - F)
|
|
|
|
@clip_color
|
|
def exclusion(F, B):
|
|
return F + B - 2*F*B
|
|
|
|
@clip_color
|
|
def saturation(img, R, G, B):
|
|
img[:,:,0] *= R
|
|
img[:,:,1] *= G
|
|
img[:,:,2] *= B
|
|
return img
|
|
|
|
@clip_color
|
|
def contrast(img, c, b):
|
|
return (img - 0.5) * c + 0.5 + b*c
|
|
|
|
def generate_composite(work_queue, result_queue):
|
|
while True:
|
|
try:
|
|
job = work_queue.get()
|
|
if job is None:
|
|
break
|
|
files_this_timestamp, timestamp, processed_images_dir = job
|
|
filename = f"Composite-{int(timestamp)}.jpg"
|
|
filepath = os.path.join(processed_images_dir, filename)
|
|
if os.path.isfile(filepath):
|
|
result_queue.put(("Exists", timestamp))
|
|
continue
|
|
|
|
# image_names = ["094Å", "131Å", "171Å", "195Å", "284Å", "304Å"]
|
|
data = []
|
|
for i in range(6):
|
|
img = Image.open(files_this_timestamp[i])
|
|
trimmed_img_data = np.array(img)[40:-40,40:-40,:3] # Trim off edges of image to remove text
|
|
normalized_img_data = trimmed_img_data / 255.0 # Normalize to float 0.0-1.0 instead of uint8
|
|
srgb_img_data = np.power(normalized_img_data, 2.2) # Gamma correct to sRGB color space
|
|
data.append(normalized_img_data)
|
|
|
|
# Assemble composite image
|
|
composite_image_data = data[4] # Start with (284Å)
|
|
# Do a linear burn with 304Å at 95% alpha
|
|
composite_image_data = composite_alpha_over(linear_burn(data[5], composite_image_data), composite_image_data, 0.95)
|
|
# Do a difference operation with 195Å at 95% alpha
|
|
composite_image_data = composite_alpha_over(exclusion(data[3], composite_image_data), composite_image_data, 0.90)
|
|
# Do a linear_light layer op with 171Å
|
|
composite_image_data = linear_light(data[2], composite_image_data)
|
|
# Do a hard_light layer op with 131Å
|
|
composite_image_data = composite_alpha_over(hard_light(data[1], composite_image_data), composite_image_data, 0.20)
|
|
# Do a color_dodge layer op with 094Å
|
|
composite_image_data = composite_alpha_over(color_dodge(data[0], composite_image_data), composite_image_data, 0.25)
|
|
# Do an exclusion layer op with 094Å
|
|
composite_image_data = composite_alpha_over(exclusion(data[0], composite_image_data), composite_image_data, 0.80)
|
|
# Tweak the colors a little
|
|
composite_image_data = saturation(composite_image_data, 1.0, 0.95, 1.15)
|
|
# Boost contrast
|
|
composite_image_data = contrast(composite_image_data, 1.5, 0.15)
|
|
|
|
# Now shrink the component images and assemble them alongside the composite.
|
|
new_dim = composite_image_data.shape[0] // 3
|
|
# Enlarge the composite to fit the new images
|
|
composite_image_data = np.pad(composite_image_data, ((0,0),(new_dim, new_dim),(0,0)))
|
|
for i in range(6):
|
|
resized = bin_ndarray(data[i], (new_dim, new_dim, 3))
|
|
if i < 3:
|
|
composite_image_data[i*new_dim:(i+1)*new_dim, :new_dim, :] = resized
|
|
else:
|
|
composite_image_data[(i-3)*new_dim:(i-2)*new_dim, -new_dim:, :] = resized
|
|
|
|
img = Image.fromarray((255 * composite_image_data).astype('uint8'))
|
|
timestring = datetime.datetime.fromtimestamp(timestamp, tz = datetime.UTC).strftime('%Y-%m-%d %H:%M:%S')
|
|
ImageDraw.Draw(img).text((655, 15), f"NOAA GOES Sattelite SUVI Composite - {timestring} UTC",(255,255,255), font_size = 24)
|
|
img.save(filepath, quality = 90)
|
|
result_queue.put(("Created", timestamp))
|
|
|
|
# plt.figure("Composite")
|
|
# plt.imshow(composite_image_data)
|
|
# plt.show()
|
|
# plt.close('all')
|
|
except Exception as e:
|
|
result_queue.put((e, timestamp))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
stored_images_dir = r"..\comp"
|
|
processed_images_dir = r"..\composite"
|
|
nworkers = 8
|
|
os.makedirs(processed_images_dir, exist_ok=True)
|
|
|
|
files_sorted_by_timestamp = defaultdict(list)
|
|
for root, dirs, files in os.walk(stored_images_dir):
|
|
for f in files:
|
|
if f.endswith(".png"):
|
|
file_parts = f.split("_")
|
|
measurement = file_parts[1]
|
|
sattelite = file_parts[2]
|
|
measure_end_time = datetime.datetime.strptime(file_parts[4][1:16], "%Y%m%dT%H%M%S")
|
|
measure_end_time.replace(tzinfo=datetime.timezone.utc)
|
|
measure_end_time = time.mktime(measure_end_time.timetuple())
|
|
files_sorted_by_timestamp[measure_end_time].append(os.path.join(root,f))
|
|
|
|
work_queue = Queue(maxsize = 3)
|
|
result_queue = Queue()
|
|
workers = []
|
|
for i in range(nworkers):
|
|
p = Process(target = generate_composite, args = (work_queue, result_queue), daemon=True)
|
|
p.start()
|
|
workers.append(p)
|
|
|
|
for timestamp in tqdm.tqdm(files_sorted_by_timestamp, desc="Creating Composite Solar Images"):
|
|
files_this_timestamp = files_sorted_by_timestamp[timestamp]
|
|
files_this_timestamp = sorted(files_this_timestamp)
|
|
if not len(files_this_timestamp) == 7:
|
|
print(f"Detected a file gap at: {timestamp}")
|
|
continue
|
|
|
|
work_queue.put((files_this_timestamp, timestamp, processed_images_dir))
|
|
|
|
ncreated = 0
|
|
nexists = 0
|
|
for _ in range(len(files_sorted_by_timestamp)):
|
|
result = result_queue.get(5.0)
|
|
if result[0] == "Exists":
|
|
nexists += 1
|
|
elif result[0] == "Created":
|
|
ncreated += 1
|
|
else:
|
|
print(f"A worker encountered an exception on job {result[1]}: {result[0]}")
|
|
|
|
for _ in range(nworkers):
|
|
try:
|
|
work_queue.put(None, timeout=1.0)
|
|
except:
|
|
break
|
|
|
|
for w in workers:
|
|
w.join(5.0)
|
|
|
|
print("Done") |