import os import time import datetime from collections import defaultdict from multiprocessing import Queue, Process import re import warnings import queue import tqdm from PIL import Image, ImageDraw import numpy as np from matplotlib import pyplot as plt from astropy.io import fits import palettable def lowpriority(): """ Set the priority of the process to below-normal.""" import sys try: sys.getwindowsversion() except AttributeError: isWindows = False else: isWindows = True if isWindows: # Based on: # "Recipe 496767: Set Process Priority In Windows" on ActiveState # http://code.activestate.com/recipes/496767/ import win32api,win32process,win32con pid = win32api.GetCurrentProcessId() handle = win32api.OpenProcess(win32con.PROCESS_ALL_ACCESS, True, pid) win32process.SetPriorityClass(handle, win32process.BELOW_NORMAL_PRIORITY_CLASS) else: import os os.nice(1) 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): d = np.abs(F - B) if d.shape[2] == 4: # Preserve alpha of base image d[:,:,3] = B[:,:,3] return d @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): d = F + B - 2*F*B if d.shape[2] == 4: # Preserve alpha of base image d[:,:,3] = B[:,:,3] return d @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 rgb_to_hsl(img): r = img[:,:,0] g = img[:,:,1] b = img[:,:,2] cmax = np.copy(r) cmax[g > cmax] = g[g > cmax] cmax[b > cmax] = b[b > cmax] cmin = np.copy(r) cmin[g < cmin] = g[g < cmin] cmin[b < cmin] = b[b < cmin] delta = cmax - cmin # Calc hue hue = np.zeros_like(r) inds = cmax == r with warnings.catch_warnings(): warnings.filterwarnings('ignore') hue[inds] = 60 * np.mod((g[inds]-b[inds])/delta[inds], 6) inds = cmax == g hue[inds] = 60 * ((b[inds]-r[inds])/delta[inds] + 2) inds = cmax == b hue[inds] = 60 * ((r[inds]-g[inds])/delta[inds] + 4) hue[np.isnan(hue)] = 0 hue[hue < 0] = hue[hue < 0] + 360 # Make negative hue values positive behind 360 # Calc lightness / luminance luminance = (cmax + cmin) / 2.0 # Calc saturation saturation = np.zeros_like(r) inds = delta != 0 saturation[inds] = delta[inds] / (1.0 - np.abs(2 * luminance[inds] - 1.0)) # Multiply luminance and saturation by 100 to scale them to the appropriate range (0-100) saturation *= 100.0 luminance *= 100.0 if img.shape[2] == 4: # Preserve alpha of original image return np.stack([hue, saturation, luminance, img[:,:,3]], -1) else: return np.stack([hue, saturation, luminance], -1) def hsl_to_rgb(img): hue = img[:,:,0] saturation = img[:,:,1] luminance = img[:,:,2] saturation /= 100.0 luminance /= 100.0 c = (1.0 - np.abs(2.0 * luminance - 1.0)) * saturation x = c * (1.0 - np.abs((hue/60.0) % 2.0 - 1.0)) m = luminance - c / 2.0 r = np.zeros_like(hue) g = np.zeros_like(hue) b = np.zeros_like(hue) inds1 = np.logical_and(0.0 <= hue, hue < 60.0) r[inds1] = c[inds1] g[inds1] = x[inds1] inds2 = np.logical_and(60.0 <= hue, hue < 120.0) r[inds2] = x[inds2] g[inds2] = c[inds2] inds3 = np.logical_and(120.0 <= hue, hue < 180.0) g[inds3] = c[inds3] b[inds3] = x[inds3] inds4 = np.logical_and(180.0 <= hue, hue < 240.0) g[inds4] = x[inds4] b[inds4] = c[inds4] inds5 = np.logical_and(240.0 <= hue, hue < 300.0) r[inds5] = x[inds5] b[inds5] = c[inds5] inds6 = np.logical_and(240.0 <= hue, hue < 300.0) r[inds6] = c[inds6] b[inds6] = x[inds6] r += m g += m b += m if img.shape[2] == 4: # Preserve alpha of original image return np.stack([r, g, b, img[:,:,3]], -1) else: return np.stack([r, g, b], -1) def generate_composite(work_queue, result_queue): lowpriority() image_names = ["094A", "131A", "171A", "195A", "284A", "304A"] cmaps =[palettable.cmocean.sequential.Ice_5.mpl_colormap, palettable.cmocean.sequential.Ice_20.mpl_colormap, palettable.cmocean.sequential.Turbid_5_r.mpl_colormap, palettable.cmocean.sequential.Turbid_20_r.mpl_colormap, plt.colormaps.get_cmap('gist_heat'), plt.colormaps.get_cmap('afmhot')] # Modify colormaps to start at perfect black (when they otherwise start at very dark colors) for cmi in range(0,4): for c in ['red','green','blue']: for i in range(3): cmaps[cmi]._segmentdata[c][0][i] = 0.0 # These values are used to map floating point radiance values to colors using the above color maps. vmins = [0.050, 0.05, 00.100, 00.10, 00.100, 00.1] vmaxs = [8.000, 8.00, 20.000, 30.00, 40.000, 90.0] gammas = [0.375, 0.40, 00.425, 00.45, 00.475, 00.5] while True: try: job = work_queue.get() if job is None: result_queue.cancel_join_thread() return 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 base_imgs = [] for i in range(6): raw_data = np.flip(fits.getdata(files_this_timestamp[i]), 0) raw_data[raw_data < 0.0] = 0.0 # Remove non-zero data because it doesn't makes sense (supposed to be std Radiance) base_imgs.append(cmaps[i](np.clip((raw_data - vmins[i]) / vmaxs[i], 0, 1.0)**gammas[i])) # plt.figure(image_names[i]) # plt.imshow(base_imgs[-1]) # plt.figure("Initial Blend") composite_image_data = composite_alpha_over(base_imgs[4], base_imgs[5], 0.2) composite_image_data = composite_image_data**0.5 # plt.imshow(composite_image_data) # plt.figure("Linear Burn with 304") composite_image_data = composite_alpha_over(linear_burn(base_imgs[5], composite_image_data), composite_image_data, 0.60) # plt.imshow(composite_image_data) # plt.figure("Light Ops with mid bands") mix_img = composite_alpha_over(exclusion(base_imgs[3], composite_image_data), composite_image_data, 0.95) mix_img = composite_alpha_over(linear_light(base_imgs[2], mix_img), mix_img, 1.0) mix_img = mix_img**0.75 composite_image_data = composite_alpha_over(mix_img, composite_image_data, 0.5) # plt.imshow(composite_image_data) # plt.figure("Mix in mid 131") composite_image_data = composite_alpha_over(base_imgs[2], composite_image_data, 0.25) # Mix in a small amount of 131Å for the nice streamers # plt.imshow(composite_image_data) # plt.figure("Mid diff") mix_img = difference(base_imgs[3], base_imgs[2]) # 171Å - 131Å # plt.imshow(mix_img) # plt.figure("HSL ops with mid") comp_hsl = rgb_to_hsl(composite_image_data) mix_img_hsl = rgb_to_hsl(mix_img) del mix_img base_img2_hsl = rgb_to_hsl(base_imgs[2]) base_img3_hsl = rgb_to_hsl(base_imgs[3]) comp_hsl = np.copy(comp_hsl) comp_hsl[:,:,0] += 0.025*mix_img_hsl[:,:,0] # Rotate hue based on mix_1_hsl del mix_img_hsl comp_hsl[:,:,0][comp_hsl[:,:,0] > 360.0] -= 360.0 comp_hsl[:,:,1] -= 0.1*base_img3_hsl[:,:,1] # Reduce saturation based on base_img3 comp_hsl[:,:,1] = np.clip(comp_hsl[:,:,1], 0.0, 100.0) comp_hsl[:,:,2] += 0.5*base_img2_hsl[:,:,2] # Boost luminance based on base_img2 comp_hsl[:,:,1] = np.clip(comp_hsl[:,:,2], 0.0, 100.0) composite_image_data = hsl_to_rgb(comp_hsl) del comp_hsl composite_image_data = saturation(composite_image_data, 1.0, 0.8, 0.1) # Remove some blue and green # plt.imshow(composite_image_data) # plt.figure("Hue shifted 131") mix_img_hsl = rgb_to_hsl(base_imgs[1]) mix_img_hsl[:,:,0] -= 50 mix_img_hsl[:,:,0][mix_img_hsl[:,:,0] < 0.0] += 360.0 mix_img = hsl_to_rgb(mix_img_hsl) del mix_img_hsl # plt.imshow(mix_img) # plt.figure("Hard Light with hue shifted 131") composite_image_data = composite_alpha_over(hard_light(mix_img, composite_image_data), composite_image_data, 0.2) # plt.imshow(composite_image_data) # plt.figure("Adjusted 094") mix_img = saturation(base_imgs[0], 1.15, 1.2, 1.05) mix_img = contrast(mix_img, 1.5, 0.0) # plt.imshow(mix_img) # plt.figure("Color Dodge with adjusted 094") composite_image_data = composite_alpha_over(color_dodge(mix_img, composite_image_data), composite_image_data, 0.5) # plt.imshow(composite_image_data) # plt.figure("Exclusion with adjusted 094") composite_image_data = composite_alpha_over(exclusion(mix_img, composite_image_data), composite_image_data, 0.65) # plt.imshow(composite_image_data) del mix_img # plt.figure("Final Image") composite_image_data = saturation(composite_image_data, 1.0, 1.1, 1.2) composite_image_data = contrast(composite_image_data, 1.20, 0.00) # plt.imshow(composite_image_data) # We have our final image # plt.show() # Trim edges of final image so it fits nicely in 1920 composite_image_data = composite_image_data[64:-64,64:-64,:3] # 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): img = base_imgs[i][64:-64,64:-64,:3] # Trim edges of image data to fit nicely img = bin_ndarray(img, (new_dim, new_dim, 3)) # Shrink down to 1/3 for assembly img = contrast(img, 1.25, 0.0) xdimoff = i%2 * (composite_image_data.shape[1] - new_dim) ydimoff = i//2*new_dim composite_image_data[ydimoff:ydimoff+new_dim, xdimoff:xdimoff+new_dim, :] = img 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) for i in range(6): # Draw component angstrom labels if i%2 == 0: xdimtxtoff = 5 else: xdimtxtoff = composite_image_data.shape[1] - 44 ydimtxtoff = i//2*new_dim + new_dim / 2.0 - 8 ImageDraw.Draw(img).text((xdimtxtoff, ydimtxtoff), image_names[i], font_size = 16) img.save(filepath, quality = 90) result_queue.put(("Created", timestamp)) except KeyboardInterrupt: return except Exception as e: result_queue.put((e, timestamp)) if __name__ == "__main__": stored_fits_dirs = [r"..\Data\goes16\l2\data", r"..\Data\goes18\l2\data"] processed_images_dirs = [r"..\composite\goes16", r"..\composite\goes18"] regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*\.fits" nworkers = 20 # Testing # stored_fits_dirs = [r"..\fits_test_2024"] # processed_images_dirs = [r"..\composite"] work_queue = Queue(maxsize = nworkers) 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) lowpriority() ncreated = 0 nexists = 0 nfailed = 0 try: for stored_fits_dir, processed_images_dir in zip(stored_fits_dirs, processed_images_dirs): os.makedirs(processed_images_dir, exist_ok=True) filename_tester = re.compile(regex_filename) files_sorted_by_timestamp = defaultdict(list) found_files = 0 print(f"Searching for FITS files in: {stored_fits_dir}") for root, dirs, files in tqdm.tqdm(os.walk(stored_fits_dir), desc="Searching"): for f in files: if filename_tester.match(f): 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)) found_files += 1 print(f"Found {found_files} FITS files. Starting conversion.") for timestamp in tqdm.tqdm(files_sorted_by_timestamp, desc="Creating Composite Solar Images"): # Collect completed jobs and record completion status while True: try: result = result_queue.get_nowait() 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]}") nfailed += 1 except queue.Empty: break # Submit new jobs files_this_timestamp = files_sorted_by_timestamp[timestamp] files_this_timestamp = sorted(files_this_timestamp) if not len(files_this_timestamp) == 6: print(f"Invalid or incomplete sensor records for: {timestamp}") continue work_queue.put((files_this_timestamp, timestamp, processed_images_dir)) except KeyboardInterrupt: print("Finishing current jobs and exiting") for _ in range(nworkers): try: work_queue.put(None, timeout=10.0) except: break for w in workers: w.join(10.0) print(f"Created {ncreated} | Already had {nexists} | Failed {nfailed}")