502 lines
19 KiB
Python
502 lines
19 KiB
Python
|
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):
|
||
|
|
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.cm.get_cmap('gist_heat'),
|
||
|
|
plt.cm.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 = 14
|
||
|
|
|
||
|
|
# Testing
|
||
|
|
# stored_fits_dirs = [r"C:\NOAA_SOLAR_IMAGES\fits_test_2024"]
|
||
|
|
# processed_images_dirs = [r"C:\NOAA_SOLAR_IMAGES\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}")
|