657 lines
25 KiB
Python
657 lines
25 KiB
Python
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}")
|