import os import time import calendar import datetime from collections import defaultdict from multiprocessing import Queue, Process import re import warnings import queue import traceback import tqdm from PIL import Image, ImageDraw, ImageFont 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 # pywin32 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() 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] trimx = 64 trimy = 100 fnt1 = ImageFont.truetype("OpenSans-Regular.ttf", size = 24) fnt2 = ImageFont.truetype("OpenSans-Regular.ttf", size = 16) while True: try: job = work_queue.get() if job is None: result_queue.cancel_join_thread() return files_this_timestamp, timestamp, processed_images_dir, synthetic_data = job synthetic_filepath = os.path.join(processed_images_dir, f"Composite-{int(timestamp)}_s.jpg") normal_filepath = os.path.join(processed_images_dir, f"Composite-{int(timestamp)}.jpg") filepath = None if synthetic_data: if os.path.isfile(synthetic_filepath): result_queue.put(("Exists", timestamp)) continue filepath = synthetic_filepath else: if os.path.isfile(synthetic_filepath): os.remove(synthetic_filepath) if os.path.isfile(normal_filepath): result_queue.put(("Exists", timestamp)) continue filepath = normal_filepath base_imgs = [] for i in range(6): raw_data = np.flip(fits.getdata(files_this_timestamp[i]), 0)[trimy:-trimy,trimx:-trimx] # Trim and reorient the image to match our final desired dimensions 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])[:,:,:3]) # 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)**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)**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() # Now shrink the component images and assemble them alongside the composite. new_dimx = composite_image_data.shape[1] // 3 new_dimy = 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_dimx, new_dimx),(0,0))) for i in range(6): img = base_imgs[i] img = bin_ndarray(img, (new_dimy, new_dimx, 3)) # Shrink down to 1/3 for assembly img = contrast(img, 1.25, 0.0) xdimoff = i%2 * (composite_image_data.shape[1] - new_dimx) ydimoff = i//2*new_dimy composite_image_data[ydimoff:ydimoff+new_dimy, xdimoff:xdimoff+new_dimx, :] = 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((602, 15), f"NOAA GOES Satellite SUVI Composite - {timestring} UTC",(255,255,255), font = fnt1) 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_dimy + new_dimy / 2.0 - 14 ImageDraw.Draw(img).text((xdimtxtoff, ydimtxtoff), image_names[i], font = fnt2) img.save(filepath, quality = 95) result_queue.put(("Created", timestamp)) except KeyboardInterrupt: return except Exception as e: traceback.print_exception(e) result_queue.put((e, timestamp)) image_names = ["094A", "131A", "171A", "195A", "284A", "304A"] 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"] starttime = calendar.timegm(datetime.datetime(2023, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) stoptime = calendar.timegm(datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*\_f.fits$" nworkers = 20 max_time_gap = 3 fill_missing_data = False file_prefixes = ["dr_suvi-l2-ci" + n[:-1] for n in image_names] # Testing # stored_fits_dirs = [r"..\fits_test_2024"] 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 black_image = Image.fromarray(np.zeros((1080, 1920, 3), dtype='uint8')) try: for stored_fits_dir, processed_images_dir in zip(stored_fits_dirs, processed_images_dirs): files_by_timestamp = defaultdict(list) num_found_files = 0 os.makedirs(processed_images_dir, exist_ok=True) filename_tester = re.compile(regex_filename) 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 = int(datetime.datetime.strptime(file_parts[4][1:16] + " +0000", "%Y%m%dT%H%M%S %z").timestamp()) if (measure_end_time >= starttime) and (measure_end_time < stoptime): files_by_timestamp[measure_end_time].append(os.path.join(root,f)) num_found_files += 1 print(f"Found {num_found_files} FITS files. Starting conversion.") if num_found_files == 0: exit(3) sorted_times = sorted(list(files_by_timestamp.keys())) min_time = sorted_times[0] max_time = sorted_times[-1] diff_times = np.diff(sorted_times) unique, counts = np.unique(diff_times, return_counts=True) interval = unique[0] # This is the amount of time between each sample in seconds assert np.sum((unique % interval) > 0) == 0 # Ensure all our timestamps align perfectly with our interval last_good_files = None last_good_file_times = None # Lets find the first and last timestamps in the sorted_times from our files which actually have a full set of 6/6 images available # This check will prevent partially downloaded sets of data from generating composite images which have "filled in" data from detected gaps # which would have been later filled with downloaded imagery. f = None l = None for i, timestamp in enumerate(sorted_times): if len(files_by_timestamp[timestamp]) == 6: f = i break for i, timestamp in enumerate(reversed(sorted_times)): if len(files_by_timestamp[timestamp]) == 6: l = len(sorted_times) - 1 - i break assert f is not None # Check to make sure we found valid indices assert l is not None assert f != l sorted_times = sorted_times[f:l] # Limit our composite image generation to only files within the valid range last_good_files = files_by_timestamp[sorted_times[0]] for timestamp in tqdm.tqdm(sorted_times, desc="Creating Composite Solar Images"): if timestamp < min_time or timestamp > max_time: continue # 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_by_timestamp[timestamp] files_this_timestamp = sorted(files_this_timestamp) if (not len(files_this_timestamp) == 6): if fill_missing_data: print(f"Invalid or incomplete sensor records for {timestamp} - {len(files_this_timestamp)}/6 filling from last good data.") files_for_job = [] for i, prefix in enumerate(file_prefixes): found = False for f in files_this_timestamp: filename = os.path.split(f)[-1] if filename.startswith(prefix): files_for_job.append(f) last_good_files[i] = f last_good_file_times[i] = timestamp found = True break if not found: # We did not find this prefix, use the last good file time_gap = (timestamp - last_good_file_times[i]) // interval if time_gap <= max_time_gap: files_for_job.append(last_good_files[i]) else: print(f"Detected a gap of {time_gap} frames at {timestamp}") break if len(files_for_job) == 6: work_queue.put((files_for_job, timestamp, processed_images_dir, True)) else: # We have a complete file set, update the last_good_files last_good_files = files_this_timestamp last_good_file_times = [timestamp for _ in last_good_files] work_queue.put((files_this_timestamp, timestamp, processed_images_dir, False)) 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) w.join() print(f"Created {ncreated} | Already had {nexists} | Failed {nfailed}")