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")