import os import sys import calendar import datetime from collections import defaultdict from multiprocessing import Queue, Process 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 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) 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) #: Band labels drawn on the composite, and the order `composite_from_arrays` #: expects its six arrays in. image_names = ["094A", "131A", "171A", "195A", "284A", "304A"] # --- Composite rendering ------------------------------------------------------- # # The blending below is the product's visual identity and is deliberately # untouched. It was extracted from the worker loop so that callers holding six # bands *in memory* -- the bench, rendering frames that were reconstructed rather # than read from disk -- can produce the same image without first writing FITS. #: Radiance-to-colour mapping, per band, in the order of `image_names`. 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] #: Trim applied to each 1280x1280 band to reach the final aspect ratio. TRIM_X = 64 TRIM_Y = 100 _render_cache = {} def _render_assets(): """Colormaps and fonts, built once per process. The first four colormaps are mutated to start at true black; that mutation is global to the palettable/matplotlib objects, so it must happen exactly once. """ if not _render_cache: 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')] # Start these maps at perfect black rather than a very dark colour. 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 font_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "OpenSans-Regular.ttf") _render_cache.update( cmaps=cmaps, fnt1=ImageFont.truetype(font_path, size=24), fnt2=ImageFont.truetype(font_path, size=16), ) return _render_cache def composite_from_arrays(arrays, timestamp): """Render one composite from six in-memory bands. Returns a PIL Image. `arrays` are the six 1280x1280 radiance arrays in `image_names` order. """ assets = _render_assets() cmaps, fnt1, fnt2 = assets["cmaps"], assets["fnt1"], assets["fnt2"] base_imgs = [] for i in range(6): raw_data = np.flip(np.asarray(arrays[i], dtype=np.float32), 0)[TRIM_Y:-TRIM_Y, TRIM_X:-TRIM_X] 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 = 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) return img def generate_composite(work_queue, result_queue): """Worker: read six bands from disk and write their composite.""" lowpriority() 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 arrays = [fits.getdata(f) for f in files_this_timestamp] img = composite_from_arrays(arrays, timestamp) 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)) if __name__ == "__main__": import argparse from suvi import db, paths parser = argparse.ArgumentParser( description="Render SUVI composite images from indexed FITS frames.", ) parser.add_argument("--db", default=None, help="SQLite index (default: $SUVI_DB)") parser.add_argument("--root", default=None, help="archive root (default: $SUVI_DATA_ROOT)") parser.add_argument("--out", default=None, help="composite output root") parser.add_argument("--satellites", default="16,18") parser.add_argument("--start", default="2024-01-01", help="UTC date, inclusive") parser.add_argument("--stop", default="2025-01-01", help="UTC date, exclusive") parser.add_argument( "--detector", default="header_v1", help="detector run whose 'good' verdicts select frames, or 'none' to use every " "indexed frame. Default header_v1: it is conservative, catching genuine " "dropouts without discarding good data. Do not default this to geometry_v1 " "-- over sampled days of 2024 it rejects 36.9%% of frames, entire days at a " "time, because its centre test tracks the brightness centroid rather than " "the disc. Run bench.py before changing this.", ) parser.add_argument("--workers", type=int, default=16) parser.add_argument("--max-time-gap", type=int, default=3, help="slots a missing band may be carried forward when filling") parser.add_argument("--fill-missing", action="store_true", help="carry the last good frame into incomplete timestamps") args = parser.parse_args() def parse_day(text): return calendar.timegm( datetime.datetime.strptime(text, "%Y-%m-%d") .replace(tzinfo=datetime.timezone.utc) .timetuple() ) archive_root = args.root or paths.data_root() composite_root = args.out or os.path.join(os.path.dirname(archive_root), "composite") satellites = [int(part) for part in args.satellites.split(",")] starttime, stoptime = parse_day(args.start), parse_day(args.stop) nworkers = args.workers max_time_gap = args.max_time_gap fill_missing_data = args.fill_missing conn = db.connect(args.db, readonly=True) run_id = None if args.detector != "none": run_id = db.latest_run_id(conn, args.detector) if run_id is None: raise SystemExit( f"No detector run named {args.detector!r} in the index. " f"Run `filter_FITS.py detect --detectors {args.detector}` first, " f"or pass --detector none to use every indexed frame." ) file_prefixes = ["dr_suvi-l2-ci" + n[:-1] for n in image_names] 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 satellite in satellites: processed_images_dir = os.path.join(composite_root, f"goes{satellite}") os.makedirs(processed_images_dir, exist_ok=True) # Frame selection is a query against the index rather than a filename # regex. The old `_f.fits$` pattern is gone along with the suffixes it # matched; see migrate_unrename.py. if run_id is None: rows = conn.execute( "SELECT path, wavelength, t_start FROM frame " "WHERE satellite = ? AND t_start >= ? AND t_start < ? " "ORDER BY t_start, wavelength", (satellite, starttime, stoptime), ).fetchall() else: rows = db.good_slots(conn, run_id, starttime, stoptime, satellite) files_by_timestamp = defaultdict(list) for row in rows: files_by_timestamp[row["t_start"]].append( paths.abspath(row["path"], archive_root) ) num_found_files = sum(len(v) for v in files_by_timestamp.values()) print(f"goes{satellite}: {num_found_files} frames selected " f"({args.detector}) across {len(files_by_timestamp)} timestamps") if num_found_files == 0: continue sorted_times = sorted(files_by_timestamp) interval = paths.CADENCE # Trim to the first and last timestamps holding a full set of six bands, so # a partially downloaded tail does not get permanently filled from a gap. first = next((i for i, t in enumerate(sorted_times) if len(files_by_timestamp[t]) == 6), None) last = next((len(sorted_times) - 1 - i for i, t in enumerate(reversed(sorted_times)) if len(files_by_timestamp[t]) == 6), None) if first is None or last is None or first >= last: print(f" no complete six-band timestamps for goes{satellite}; skipping") continue sorted_times = sorted_times[first:last] last_good_files = list(files_by_timestamp[sorted_times[0]]) last_good_file_times = [sorted_times[0]] * len(last_good_files) for timestamp in tqdm.tqdm(sorted_times, desc=f"goes{satellite} composites"): 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 failed on job {result[1]}: {result[0]}") nfailed += 1 except queue.Empty: break files_this_timestamp = sorted(files_by_timestamp[timestamp]) if len(files_this_timestamp) == 6: last_good_files = files_this_timestamp last_good_file_times = [timestamp] * 6 work_queue.put((files_this_timestamp, timestamp, processed_images_dir, False)) elif fill_missing_data: files_for_job = [] for i, prefix in enumerate(file_prefixes): found = next( (f for f in files_this_timestamp if os.path.basename(f).startswith(prefix)), None, ) if found is not None: files_for_job.append(found) last_good_files[i] = found last_good_file_times[i] = timestamp elif (timestamp - last_good_file_times[i]) // interval <= max_time_gap: files_for_job.append(last_good_files[i]) if len(files_for_job) == 6: work_queue.put((files_for_job, timestamp, processed_images_dir, True)) except KeyboardInterrupt: print("Finishing current jobs and exiting") for _ in range(nworkers): try: work_queue.put(None, timeout=10.0) except Exception: break for w in workers: w.join() conn.close() print(f"Created {ncreated} | Already had {nexists} | Failed {nfailed}")