diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..600d2d3 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +.vscode \ No newline at end of file diff --git a/ffmpeg_video.py b/ffmpeg_video.py index 1edf8bd..3c61788 100644 --- a/ffmpeg_video.py +++ b/ffmpeg_video.py @@ -1,67 +1,73 @@ -import subprocess -import os -import shutil -import datetime -import time -import calendar - -import numpy as np -from sortedcontainers import SortedDict -import tqdm -from PIL import Image - -path_to_images = r"..\composite\goes18" -output_file = r"..\goes18_2023.mp4" -interp_file = r"..\goes18_2023_interp.mp4" -ffmpeg_path = r"..\ffmpeg.exe" -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()) -min_file_size = 350000 # Detect and remove corrupted files by filtering by file size -max_file_size = 450000 -encoding_crf = 16 -max_frame_interp = 120 - -# Set up ffmpeg to stream images from an input pipe -command_line = f'{ffmpeg_path} -y -f image2pipe -framerate 60 -i - -c:v libx264 -crf {encoding_crf} -preset veryfast {output_file}' -print(command_line) -p = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE, stdin=subprocess.PIPE) - -black_image = Image.fromarray(np.zeros((1080, 1920, 3), dtype='uint8')) - -files_by_time = SortedDict() -for root, dirs, files in tqdm.tqdm(os.walk(path_to_images), desc="Finding and filtering files."): - for f in files: - if f.endswith('.jpg'): - fpath = os.path.join(root, f) - fsize = os.path.getsize(fpath) - time_chunk = f.replace('-', '.').replace('_', '.').split('.')[1] - ftime = int(time_chunk) - if (fsize > min_file_size) and (fsize < max_file_size) and (ftime > starttime) and (ftime < stoptime): - files_by_time[ftime] = fpath - -difftimes = np.diff(files_by_time.keys()) -unique, counts = np.unique(difftimes, return_counts=True) -interval = unique[0] -assert np.sum((unique % interval) > 0) == 0 # Ensure all our timestamps align perfectly with our interval - -prevtime = files_by_time.peekitem(0)[0] - interval -for t, f in tqdm.tqdm(files_by_time.items(), desc="FFMPEG encoding images", total = len(files_by_time)): - framejump = (t - prevtime) // interval - if framejump < max_frame_interp: - for i in range(framejump - 1): # Fill black frames to fill gaps, these will later be interpolated with ffmpeg - black_image.save(p.stdin, 'jpeg', quality = 95) - else: # Skip past intervals that are too large to fill reasonably - print(f"Detected a frame gap of: {framejump}!") - with open(f, 'rb') as fh: - p.stdin.write(fh.read()) - prevtime = t - -p.stdin.close() # Close the ffmpeg input pipe -p.wait() # Wait for ffmpeg to finish encoding - -command_line = f'{ffmpeg_path} -y -i {output_file} -vf blackframe=0,metadata=select:key=lavfi.blackframe.pblack:value=95:function=less,minterpolate=mi_mode=mci:mc_mode=aobmc:me_mode=bidir:vsbmc=1:me=fss -c:a copy -c:v libx264 -crf {encoding_crf} -preset veryfast {interp_file}' -print(command_line) - -pipe = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE).stdout -output = pipe.read().decode() -pipe.close() \ No newline at end of file +import subprocess +import os +import shutil +import datetime +import time +import calendar + +import numpy as np +from sortedcontainers import SortedDict +import tqdm +from PIL import Image + +years = [2024] +sources = ["goes16", "goes18"] +ffmpeg_path = r"..\ffmpeg.exe" + +for year in years: + for source in sources: + path_to_images = f"..\\composite\\{source}" + output_file = f"..\\{source}_{year}_nofilt.mp4" + interp_file = f"..\\{source}_{year}_interp_nofilt.mp4" + + starttime = calendar.timegm(datetime.datetime(year, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) + stoptime = calendar.timegm(datetime.datetime(year+1, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) + min_file_size = 350000 # Detect and remove corrupted files by filtering by file size + max_file_size = 450000 + encoding_crf = 16 + max_frame_interp = 120 + + # Set up ffmpeg to stream images from an input pipe + command_line = f'{ffmpeg_path} -y -f image2pipe -framerate 60 -i - -c:v libx264 -crf {encoding_crf} -preset veryfast {output_file}' + print(command_line) + p = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE, stdin=subprocess.PIPE) + + black_image = Image.fromarray(np.zeros((1080, 1920, 3), dtype='uint8')) + + files_by_time = SortedDict() + for root, dirs, files in tqdm.tqdm(os.walk(path_to_images), desc="Finding and filtering files."): + for f in files: + if f.endswith('.jpg'): + fpath = os.path.join(root, f) + fsize = os.path.getsize(fpath) + time_chunk = f.replace('-', '.').replace('_', '.').split('.')[1] + ftime = int(time_chunk) + if (fsize > min_file_size) and (fsize < max_file_size) and (ftime > starttime) and (ftime < stoptime): + files_by_time[ftime] = fpath + + difftimes = np.diff(files_by_time.keys()) + unique, counts = np.unique(difftimes, return_counts=True) + interval = unique[0] + assert np.sum((unique % interval) > 0) == 0 # Ensure all our timestamps align perfectly with our interval + + prevtime = files_by_time.peekitem(0)[0] - interval + for t, f in tqdm.tqdm(files_by_time.items(), desc="FFMPEG encoding images", total = len(files_by_time)): + framejump = (t - prevtime) // interval + if framejump < max_frame_interp: + for i in range(framejump - 1): # Fill black frames to fill gaps, these will later be interpolated with ffmpeg + black_image.save(p.stdin, 'jpeg', quality = 95) + else: # Skip past intervals that are too large to fill reasonably + print(f"Detected a frame gap of: {framejump}!") + with open(f, 'rb') as fh: + p.stdin.write(fh.read()) + prevtime = t + + p.stdin.close() # Close the ffmpeg input pipe + p.wait() # Wait for ffmpeg to finish encoding + + command_line = f'{ffmpeg_path} -y -i {output_file} -vf blackframe=0,metadata=select:key=lavfi.blackframe.pblack:value=95:function=less,minterpolate=mi_mode=mci:mc_mode=aobmc:me_mode=bidir:vsbmc=1:me=fss -c:a copy -c:v libx264 -crf {encoding_crf} -preset veryfast {interp_file}' + print(command_line) + + pipe = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE).stdout + output = pipe.read().decode() + pipe.close() \ No newline at end of file diff --git a/filter_FITS.py b/filter_FITS.py index 192078b..742ebf1 100644 --- a/filter_FITS.py +++ b/filter_FITS.py @@ -1,281 +1,281 @@ -import os -from collections import defaultdict -from multiprocessing import Queue, Process -import re -import traceback -import datetime -import calendar - -import tqdm -import numpy as np -from matplotlib import pyplot as plt -from astropy.io import fits -from skimage.transform import hough_circle, hough_circle_peaks -from skimage.feature import canny -from skimage.filters import gaussian -from skimage.morphology import skeletonize -import cv2 as cv - -def lowpriority(): - """ Set the priority of the process to lowest possible.""" - - import sys - try: - sys.getwindowsversion() - except AttributeError: - isWindows = False - else: - isWindows = True - - if isWindows: - import win32api,win32process,win32con # pywin32 - - pid = win32api.GetCurrentProcessId() - phandle = win32api.OpenProcess(win32con.PROCESS_ALL_ACCESS, True, pid) - win32process.SetPriorityClass(phandle, win32process.THREAD_PRIORITY_LOWEST) - # win32process.SetPriorityClass(phandle, win32process.IDLE_PRIORITY_CLASS) - # tid = win32api.GetCurrentThreadId() - # thandle = win32api.OpenThread(win32con.PROCESS_ALL_ACCESS, True, tid) - # win32process.SetThreadPriority(thandle, win32process.THREAD_MODE_BACKGROUND_BEGIN) - else: - import os - - os.nice(19) - -measurement_names = ["094", "131", "171", "195", "284", "304"] -thresholds = [0.050, 0.10 , 1.00, 1.40, 1.0, 2.50] -max_center_skew = 7 -expected_dims = 1280 -half_dims = expected_dims // 2 -valid_radii = [383, 394] -avg_radius = (valid_radii[0] + valid_radii[1]) // 2 -max_radius_error = 80 - -ratio_above_thresh_max = 0.4 -ratio_above_thresh_min = 0.1 - -def filter_fits(work_queue): - lowpriority() - idealcircle = np.zeros((expected_dims // 2, expected_dims // 2)) - cv.circle(idealcircle, (expected_dims // 4, expected_dims // 4), avg_radius // 2, 1, -1) - idealcircle_axis = np.average(idealcircle, 0) - while True: - try: - job = work_queue.get() - if job is None: - return - measurement = job.split("dr_suvi-l2-")[1].split("_")[0][2:] - idx = measurement_names.index(measurement) - try: - data = fits.getdata(job) - except IndexError: # This can happen with a blank or corrupt HDU in the .fits file - new_name = job.split(".fits")[0] + "_e.fits" - os.rename(job, new_name) - continue - assert data.shape[0] == expected_dims - assert data.shape[1] == expected_dims - data = cv.resize(data, dsize=(expected_dims // 2, expected_dims // 2), interpolation=cv.INTER_LINEAR) - filtered_data = np.copy(data) - above_thresh_indexes = data > thresholds[idx] - filtered_data[above_thresh_indexes] = thresholds[idx] - filtered_data /= thresholds[idx] - ratio_above_thresh = np.count_nonzero(above_thresh_indexes) / data.shape[0] / data.shape[1] - - if ratio_above_thresh > ratio_above_thresh_max: - print(f"Exceeded ratio_above_thresh_max, possible data corruption in file: {job}") - new_name = job.split(".fits")[0] + "_e.fits" - os.rename(job, new_name) - - plt.figure(figsize=[10.24, 7.68]) - plt.title(f"Exceeded ratio_above_thresh_max [{ratio_above_thresh:0.2f}]") - plt.imshow(filtered_data, cmap='jet') - plt.savefig(job.split(".fits")[0] + "_e.jpg") - plt.close('all') - - continue - - if ratio_above_thresh < ratio_above_thresh_min: - print(f"Below ratio_above_thresh_min, possible bad data in file: {job}") - new_name = job.split(".fits")[0] + "_e.fits" - os.rename(job, new_name) - - plt.figure(figsize=[10.24, 7.68]) - plt.title(f"Below ratio_above_thresh_min [{ratio_above_thresh:0.2f}]") - plt.imshow(filtered_data, cmap='jet') - plt.savefig(job.split(".fits")[0] + "_e.jpg") - plt.close('all') - - continue - - # Decide whether the data is valid based on solar disc placement and levels. - xavg = np.average(filtered_data, 0) - yavg = np.average(filtered_data, 1) - axis_indices = list(range(half_dims)) - - # Center Estimation - weighted_avg_center_x = np.average(axis_indices, 0, xavg) - weighted_avg_center_y = np.average(axis_indices, 0, yavg) - - good_center = True - if (half_dims // 2 - weighted_avg_center_x > max_center_skew) or (half_dims // 2 - weighted_avg_center_y > max_center_skew): - print(f"Could not find centered solar disc in file: {job}") - good_center = False - - # Radius Estimation - edge_thresh = 0.98 - high_edge_x = np.argmax(np.cumsum(xavg) > (np.sum(xavg) * edge_thresh)) - low_edge_x = len(xavg) - np.argmax(np.cumsum(np.flip(xavg)) > (np.sum(xavg) * edge_thresh)) - high_edge_y = np.argmax(np.cumsum(yavg) > (np.sum(yavg) * edge_thresh)) - low_edge_y = len(yavg) - np.argmax(np.cumsum(np.flip(yavg)) > (np.sum(yavg) * edge_thresh)) - - good_radius = True - calc_radius_x = (high_edge_x - low_edge_x) # We do not divide by two because our image is downsized by half in each dimension - calc_radius_y = (high_edge_y - low_edge_y) - calc_radius = (calc_radius_x + calc_radius_y) / 2.0 - if abs(calc_radius - avg_radius) > max_radius_error: - print(f"Could not find correct solar radius in file: {job}") - good_radius = False - - # Goodness of fit estimation vs perfect disc - gof_x = np.sum(np.abs(xavg - idealcircle_axis)) / half_dims - gof_y = np.sum(np.abs(yavg - idealcircle_axis)) / half_dims - - good_fit = True - if gof_x > 0.2 or gof_y > 0.2: - print(f"Could not find valid solar disc in file: {job}") - good_fit = False - - - if good_center and good_radius and good_fit: - # We have validated this file, rename it appropriately - new_name = job.split(".fits")[0] + "_f.fits" - os.rename(job, new_name) - else: - new_name = job.split(".fits")[0] + "_e.fits" - os.rename(job, new_name) - - # Plot Results - fig, axes = plt.subplots(2, 2, width_ratios=(0.2, 1), height_ratios=(0.2, 1), gridspec_kw={"hspace":0.0, "wspace": 0.0}, figsize = [8, 8]) - - axes[0][0].axis('off') - - axes[0][1].plot(axis_indices, xavg) - axes[0][1].plot(axis_indices, idealcircle_axis) - # axes[0][1].fill_between(axis_indices, xavg, idealcircle_axis, hatch="//", edgecolor="red", facecolor="none") - - axes[1][0].plot(yavg, axis_indices) - axes[1][0].plot(idealcircle_axis, axis_indices) - axes[1][0].invert_xaxis() - axes[1][0].invert_yaxis() - - axes[1][1].imshow(filtered_data, aspect='auto', vmin=0, vmax=1) - axes[1][1].vlines([low_edge_x, high_edge_x], 0, half_dims) - axes[1][1].hlines([low_edge_y, high_edge_y], 0, half_dims) - axes[1][1].scatter(half_dims // 2, half_dims // 2, marker="o", linewidths = 1, alpha = 0.75, edgecolors="black", facecolors = "none") - axes[1][1].scatter(weighted_avg_center_x, weighted_avg_center_y, marker = "x", alpha = 0.75, facecolors='red') - - # Set shared x axis between imshow and xdata plots - axes[0][1].xaxis.set_ticks_position("top") - axes[0][1].yaxis.set_ticks_position("right") - axes[0][1].sharex(axes[1][1]) - axes[1][0].sharey(axes[1][1]) - axes[1][1].yaxis.set_ticks_position("right") - - plt.savefig(job.split(".fits")[0] + "_e.jpg") - plt.close('all') - - - except KeyboardInterrupt: - return - except Exception as e: - print(f"Error on file: {job} - {e}") - traceback.print_exception(e) - - -if __name__ == "__main__": - stored_fits_dirs = [r"..\Data\goes16\l2\data\suvi-l2-ci094\2024", - r"..\Data\goes16\l2\data\suvi-l2-ci131\2024", - r"..\Data\goes16\l2\data\suvi-l2-ci171\2024", - r"..\Data\goes16\l2\data\suvi-l2-ci195\2024", - r"..\Data\goes16\l2\data\suvi-l2-ci284\2024", - r"..\Data\goes16\l2\data\suvi-l2-ci304\2024", - r"..\Data\goes18\l2\data\suvi-l2-ci094\2024", - r"..\Data\goes18\l2\data\suvi-l2-ci131\2024", - r"..\Data\goes18\l2\data\suvi-l2-ci171\2024", - r"..\Data\goes18\l2\data\suvi-l2-ci195\2024", - r"..\Data\goes18\l2\data\suvi-l2-ci284\2024", - r"..\Data\goes18\l2\data\suvi-l2-ci304\2024",] - - # stored_fits_dirs = [r"..\Data\goes18\l2\data\suvi-l2-ci284\2023\01\10"] - - starttime = calendar.timegm(datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) - stoptime = calendar.timegm(datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) - - reprocess_errors = False - nworkers = 16 - - regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*\.fits$" - - work_queue = Queue(maxsize = nworkers) - workers = [] - for i in range(nworkers): - p = Process(target = filter_fits, args = (work_queue,), daemon=True) - p.start() - workers.append(p) - - lowpriority() - files_by_timestamp = defaultdict(list) - try: - for stored_fits_dir in stored_fits_dirs: - filename_tester = re.compile(regex_filename) - - 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("_") - file_name_end = file_parts[-1].split(".")[0] - 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): - if file_name_end == "f": - if file_parts[-2] == "f": # We have an accidentally double filtered file... - new_file_name = "_".join(file_parts[:-1]) + ".fits" - os.rename(os.path.join(root,f), os.path.join(root,new_file_name)) - continue # Already filtered from a previous run - elif file_name_end == "e": - if reprocess_errors: - # Rename file to remove error designation - new_file_name = "_".join(file_parts[:-1]) + ".fits" - os.rename(os.path.join(root,f), os.path.join(root,new_file_name)) - # Check if error image exists and delete if needed - img_name = new_file_name.split(".fits")[0] + "_e.jpg" - try: - os.remove(os.path.join(root, img_name)) - except: - pass - # Add job to queue - files_by_timestamp[measure_end_time].append(os.path.join(root, new_file_name)) - else: - continue - elif file_name_end.startswith("v1-0-"): # This is the normal case for unprocessed data - files_by_timestamp[measure_end_time].append(os.path.join(root,f)) - else: - # print(f"Error - Unexpected FITS file name: {f}") - pass - sorted_times = sorted(list(files_by_timestamp.keys())) - for st in tqdm.tqdm(sorted_times, desc="Filtering files"): - for f in files_by_timestamp[st]: - work_queue.put(f) - 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() +import os +from collections import defaultdict +from multiprocessing import Queue, Process +import re +import traceback +import datetime +import calendar + +import tqdm +import numpy as np +from matplotlib import pyplot as plt +from astropy.io import fits +from skimage.transform import hough_circle, hough_circle_peaks +from skimage.feature import canny +from skimage.filters import gaussian +from skimage.morphology import skeletonize +import cv2 as cv + +def lowpriority(): + """ Set the priority of the process to lowest possible.""" + + import sys + try: + sys.getwindowsversion() + except AttributeError: + isWindows = False + else: + isWindows = True + + if isWindows: + import win32api,win32process,win32con # pywin32 + + pid = win32api.GetCurrentProcessId() + phandle = win32api.OpenProcess(win32con.PROCESS_ALL_ACCESS, True, pid) + win32process.SetPriorityClass(phandle, win32process.THREAD_PRIORITY_LOWEST) + # win32process.SetPriorityClass(phandle, win32process.IDLE_PRIORITY_CLASS) + # tid = win32api.GetCurrentThreadId() + # thandle = win32api.OpenThread(win32con.PROCESS_ALL_ACCESS, True, tid) + # win32process.SetThreadPriority(thandle, win32process.THREAD_MODE_BACKGROUND_BEGIN) + else: + import os + + os.nice(19) + +measurement_names = ["094", "131", "171", "195", "284", "304"] +thresholds = [0.050, 0.10 , 1.00, 1.40, 1.0, 2.50] +max_center_skew = 7 +expected_dims = 1280 +half_dims = expected_dims // 2 +valid_radii = [383, 394] +avg_radius = (valid_radii[0] + valid_radii[1]) // 2 +max_radius_error = 80 + +ratio_above_thresh_max = 0.4 +ratio_above_thresh_min = 0.07 + +def filter_fits(work_queue): + lowpriority() + idealcircle = np.zeros((expected_dims // 2, expected_dims // 2)) + cv.circle(idealcircle, (expected_dims // 4, expected_dims // 4), avg_radius // 2, 1, -1) + idealcircle_axis = np.average(idealcircle, 0) + while True: + try: + job = work_queue.get() + if job is None: + return + measurement = job.split("dr_suvi-l2-")[1].split("_")[0][2:] + idx = measurement_names.index(measurement) + try: + data = fits.getdata(job) + except IndexError: # This can happen with a blank or corrupt HDU in the .fits file + new_name = job.split(".fits")[0] + "_e.fits" + os.rename(job, new_name) + continue + assert data.shape[0] == expected_dims + assert data.shape[1] == expected_dims + data = cv.resize(data, dsize=(expected_dims // 2, expected_dims // 2), interpolation=cv.INTER_LINEAR) + filtered_data = np.copy(data) + above_thresh_indexes = data > thresholds[idx] + filtered_data[above_thresh_indexes] = thresholds[idx] + filtered_data /= thresholds[idx] + ratio_above_thresh = np.count_nonzero(above_thresh_indexes) / data.shape[0] / data.shape[1] + + if ratio_above_thresh > ratio_above_thresh_max: + print(f"Exceeded ratio_above_thresh_max, possible data corruption in file: {job}") + new_name = job.split(".fits")[0] + "_e.fits" + os.rename(job, new_name) + + plt.figure(figsize=[10.24, 7.68]) + plt.title(f"Exceeded ratio_above_thresh_max [{ratio_above_thresh:0.2f}]") + plt.imshow(filtered_data, cmap='jet') + plt.savefig(job.split(".fits")[0] + "_e.jpg") + plt.close('all') + + continue + + if ratio_above_thresh < ratio_above_thresh_min: + print(f"Below ratio_above_thresh_min, possible bad data in file: {job}") + new_name = job.split(".fits")[0] + "_e.fits" + os.rename(job, new_name) + + plt.figure(figsize=[10.24, 7.68]) + plt.title(f"Below ratio_above_thresh_min [{ratio_above_thresh:0.2f}]") + plt.imshow(filtered_data, cmap='jet') + plt.savefig(job.split(".fits")[0] + "_e.jpg") + plt.close('all') + + continue + + # Decide whether the data is valid based on solar disc placement and levels. + xavg = np.average(filtered_data, 0) + yavg = np.average(filtered_data, 1) + axis_indices = list(range(half_dims)) + + # Center Estimation + weighted_avg_center_x = np.average(axis_indices, 0, xavg) + weighted_avg_center_y = np.average(axis_indices, 0, yavg) + + good_center = True + if (half_dims // 2 - weighted_avg_center_x > max_center_skew) or (half_dims // 2 - weighted_avg_center_y > max_center_skew): + print(f"Could not find centered solar disc in file: {job}") + good_center = False + + # Radius Estimation + edge_thresh = 0.98 + high_edge_x = np.argmax(np.cumsum(xavg) > (np.sum(xavg) * edge_thresh)) + low_edge_x = len(xavg) - np.argmax(np.cumsum(np.flip(xavg)) > (np.sum(xavg) * edge_thresh)) + high_edge_y = np.argmax(np.cumsum(yavg) > (np.sum(yavg) * edge_thresh)) + low_edge_y = len(yavg) - np.argmax(np.cumsum(np.flip(yavg)) > (np.sum(yavg) * edge_thresh)) + + good_radius = True + calc_radius_x = (high_edge_x - low_edge_x) # We do not divide by two because our image is downsized by half in each dimension + calc_radius_y = (high_edge_y - low_edge_y) + calc_radius = (calc_radius_x + calc_radius_y) / 2.0 + if abs(calc_radius - avg_radius) > max_radius_error: + print(f"Could not find correct solar radius in file: {job}") + good_radius = False + + # Goodness of fit estimation vs perfect disc + gof_x = np.sum(np.abs(xavg - idealcircle_axis)) / half_dims + gof_y = np.sum(np.abs(yavg - idealcircle_axis)) / half_dims + + good_fit = True + if gof_x > 0.2 or gof_y > 0.2: + print(f"Could not find valid solar disc in file: {job}") + good_fit = False + + + if good_center and good_radius and good_fit: + # We have validated this file, rename it appropriately + new_name = job.split(".fits")[0] + "_f.fits" + os.rename(job, new_name) + else: + new_name = job.split(".fits")[0] + "_e.fits" + os.rename(job, new_name) + + # Plot Results + fig, axes = plt.subplots(2, 2, width_ratios=(0.2, 1), height_ratios=(0.2, 1), gridspec_kw={"hspace":0.0, "wspace": 0.0}, figsize = [8, 8]) + + axes[0][0].axis('off') + + axes[0][1].plot(axis_indices, xavg) + axes[0][1].plot(axis_indices, idealcircle_axis) + # axes[0][1].fill_between(axis_indices, xavg, idealcircle_axis, hatch="//", edgecolor="red", facecolor="none") + + axes[1][0].plot(yavg, axis_indices) + axes[1][0].plot(idealcircle_axis, axis_indices) + axes[1][0].invert_xaxis() + axes[1][0].invert_yaxis() + + axes[1][1].imshow(filtered_data, aspect='auto', vmin=0, vmax=1) + axes[1][1].vlines([low_edge_x, high_edge_x], 0, half_dims) + axes[1][1].hlines([low_edge_y, high_edge_y], 0, half_dims) + axes[1][1].scatter(half_dims // 2, half_dims // 2, marker="o", linewidths = 1, alpha = 0.75, edgecolors="black", facecolors = "none") + axes[1][1].scatter(weighted_avg_center_x, weighted_avg_center_y, marker = "x", alpha = 0.75, facecolors='red') + + # Set shared x axis between imshow and xdata plots + axes[0][1].xaxis.set_ticks_position("top") + axes[0][1].yaxis.set_ticks_position("right") + axes[0][1].sharex(axes[1][1]) + axes[1][0].sharey(axes[1][1]) + axes[1][1].yaxis.set_ticks_position("right") + + plt.savefig(job.split(".fits")[0] + "_e.jpg") + plt.close('all') + + + except KeyboardInterrupt: + return + except Exception as e: + print(f"Error on file: {job} - {e}") + traceback.print_exception(e) + + +if __name__ == "__main__": + stored_fits_dirs = [r"..\Data\goes16\l2\data\suvi-l2-ci094\2024", + r"..\Data\goes16\l2\data\suvi-l2-ci131\2024", + r"..\Data\goes16\l2\data\suvi-l2-ci171\2024", + r"..\Data\goes16\l2\data\suvi-l2-ci195\2024", + r"..\Data\goes16\l2\data\suvi-l2-ci284\2024", + r"..\Data\goes16\l2\data\suvi-l2-ci304\2024", + r"..\Data\goes18\l2\data\suvi-l2-ci094\2024", + r"..\Data\goes18\l2\data\suvi-l2-ci131\2024", + r"..\Data\goes18\l2\data\suvi-l2-ci171\2024", + r"..\Data\goes18\l2\data\suvi-l2-ci195\2024", + r"..\Data\goes18\l2\data\suvi-l2-ci284\2024", + r"..\Data\goes18\l2\data\suvi-l2-ci304\2024",] + + # stored_fits_dirs = [r"..\Data\goes18\l2\data\suvi-l2-ci284\2023\01\10"] + + starttime = calendar.timegm(datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) + stoptime = calendar.timegm(datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) + + reprocess_errors = True + nworkers = 16 + + regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*\.fits$" + + work_queue = Queue(maxsize = nworkers) + workers = [] + for i in range(nworkers): + p = Process(target = filter_fits, args = (work_queue,), daemon=True) + p.start() + workers.append(p) + + lowpriority() + files_by_timestamp = defaultdict(list) + try: + for stored_fits_dir in stored_fits_dirs: + filename_tester = re.compile(regex_filename) + + 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("_") + file_name_end = file_parts[-1].split(".")[0] + 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): + if file_name_end == "f": + if file_parts[-2] == "f": # We have an accidentally double filtered file... + new_file_name = "_".join(file_parts[:-1]) + ".fits" + os.rename(os.path.join(root,f), os.path.join(root,new_file_name)) + continue # Already filtered from a previous run + elif file_name_end == "e": + if reprocess_errors: + # Rename file to remove error designation + new_file_name = "_".join(file_parts[:-1]) + ".fits" + os.rename(os.path.join(root,f), os.path.join(root,new_file_name)) + # Check if error image exists and delete if needed + img_name = new_file_name.split(".fits")[0] + "_e.jpg" + try: + os.remove(os.path.join(root, img_name)) + except: + pass + # Add job to queue + files_by_timestamp[measure_end_time].append(os.path.join(root, new_file_name)) + else: + continue + elif file_name_end.startswith("v1-0-"): # This is the normal case for unprocessed data + files_by_timestamp[measure_end_time].append(os.path.join(root,f)) + else: + # print(f"Error - Unexpected FITS file name: {f}") + pass + sorted_times = sorted(list(files_by_timestamp.keys())) + for st in tqdm.tqdm(sorted_times, desc="Filtering files"): + for f in files_by_timestamp[st]: + work_queue.put(f) + 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() diff --git a/merger_FITS.py b/merger_FITS.py index 833578d..9d99fc4 100644 --- a/merger_FITS.py +++ b/merger_FITS.py @@ -1,574 +1,579 @@ -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(2024, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) - stoptime = calendar.timegm(datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) - - regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*\_f.fits$" - nworkers = 16 - max_time_gap = 3 - fill_missing_data = True - - 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} | {datetime.datetime.fromtimestamp(timestamp, datetime.UTC).strftime('%Y-%m-%d %H:%M:%S %Z')} - {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} | {datetime.datetime.fromtimestamp(timestamp, datetime.UTC).strftime('%Y-%m-%d %H:%M:%S %Z')}") - 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() - +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(2024, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) + stoptime = calendar.timegm(datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) + + regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*\_f.fits$" + + nworkers = 16 + max_time_gap = 3 + fill_missing_data = False + ignore_errors = True + + if ignore_errors: + regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*.fits$" + + 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} | {datetime.datetime.fromtimestamp(timestamp, datetime.UTC).strftime('%Y-%m-%d %H:%M:%S %Z')} - {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} | {datetime.datetime.fromtimestamp(timestamp, datetime.UTC).strftime('%Y-%m-%d %H:%M:%S %Z')}") + 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, ignore_errors)) + + 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}") \ No newline at end of file diff --git a/puller_fits.py b/puller_fits.py index 6d3ba4c..0687a88 100644 --- a/puller_fits.py +++ b/puller_fits.py @@ -1,258 +1,273 @@ -import os -import urllib.request -import urllib.parse -import re -import time -import random -import datetime -import json -from threading import Thread -import queue -import math -from functools import partial -from queue import Empty - -import tqdm - -directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/" -stored_images_dir = r"..\Data" -# directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/goes16/l2/data/ephe-l2-orb1m/2017/12/" -# stored_images_dir = r"Z:\NOAA GOES Data\Data\goes16/l2/data/ephe-l2-orb1m/2017/12/" -ignore_folder_names = ["l1b", "goes17", "2017", "2018", "2019", "2020", "2021", "2022", "2023"] -file_database_path = r"..\file_database.json" -fetch_interval = 0 # 60*60 -nfetchworkers = 10 # Be nice to the servers, this value is how many threads will be asking for links and file info at the same time -ndownloadworkers = 4 # Be nice to the servers, this value is how many threads will be downloading files at the same time -randomize_order = False - -links_regex_pattern = r'(?<=.*\d{4}-\d{2}-\d{2} \d{2}:\d{2})' # Find href links that do not contain question marks or whitespace -times_regex_pattern = r'(?<=)(\d{4}-\d{2}-\d{2} \d{2}:\d{2})(?= )' # Find timestamps in UTC in the format YYYY-MM-DD HH:mm -sizes_regex_pattern = r'(?<=\d{4}-\d{2}-\d{2} \d{2}:\d{2} <\/td>)([ .\d]{1,3})([KMG]|- |0 | )(?=<\/td>)' # Find file size markers in one of many possible formats -links_matcher = re.compile(links_regex_pattern) -times_matcher = re.compile(times_regex_pattern) -sizes_matcher = re.compile(sizes_regex_pattern) - -def find_links_worker(query_work_queue, query_result_queue, attempt_count = 3, randomize_order = True): - while True: - job = query_work_queue.get() - if job == None: - return - url= job - - attempts = 0 - html_content = None - while attempts < attempt_count: - try: - with urllib.request.urlopen(url) as response: - html_content = response.read().decode('utf-8') - break - except Exception as e: - print(f"Exception while fetching links: {e}") - time.sleep(1 + random.random()) - attempts += 1 - if (attempt_count > 1) and (attempts == attempt_count): - print(f'\nAfter {attempt_count} retries, could not fetch: {url}') - return - - links = links_matcher.findall(html_content) - times = times_matcher.findall(html_content) - sizes = sizes_matcher.findall(html_content) - for i in range(len(times)): - dt = datetime.datetime.strptime(times[i].strip(), "%Y-%m-%d %H:%M") - dt.replace(tzinfo=datetime.timezone.utc) - times[i] = time.mktime(dt.timetuple()) - for i in range(len(sizes)): - match sizes[i][1].strip(): - case '-': - sizes[i] = 0 - case 'K': - sizes[i] = int(float(sizes[i][0].strip()) * 1024) - case 'M': - sizes[i] = int(float(sizes[i][0].strip()) * 1024 * 1024) - case 'G': - sizes[i] = int(float(sizes[i][0].strip()) * 1024 * 1024 * 1024) - case "0": - sizes[i] = 0 - case "": - sizes[i] = int(float(sizes[i][0].strip())) - case _: - raise(ValueError(f"Unexpected symbol while parsing links page: {_}")) - - - if (len(links) != len(times)) or (len(times) != len(sizes)): - raise(ValueError("Links parsing error, mismatched numbers of links, times, or sizes!")) - - results = list(zip(links, times, sizes)) - if randomize_order: - random.shuffle(results) - for _link, _time, _size in results: - if _link.endswith("/"): - if _link.split(r"/")[-2] in ignore_folder_names: - continue - else: - query_work_queue.put(urllib.parse.urljoin(url,_link)) - else: - query_result_queue.put((urllib.parse.urljoin(url,_link), _time, _size)) - - -# Fetch file from url and store it to path, retrying on failure -def file_download_worker(download_work_queue, download_result_queue, attempt_count = 2): - while True: - job = download_work_queue.get() - if job == None: - return - url, path, t, s = job - - attempts = 0 - while attempts < attempt_count: - try: - req = urllib.request.Request(url, data=None) - image_data = urllib.request.urlopen(req).read() - if math.isclose(len(image_data), s, rel_tol=0.05): - os.makedirs(os.path.split(path)[0], exist_ok=True) - open(path, 'wb').write(image_data) - download_result_queue.put((True, url, t)) - break - else: - raise ValueError("Downloaded file is the wrong size!") - except Exception as e: - if hasattr(e, "code") and e.code == 404: # This is expected if the file has been removed from the site (at least for swpc.noaa.gov) - attempts += attempt_count - elif attempts == 0: - print(f"\nA problem occurred on image: {url} | {e}") - time.sleep(1 + random.random()) - attempts += 1 - if (attempt_count > 1) and (attempts == attempt_count): - print(f'\nAfter {attempt_count} retries, could not fetch: {url}') - download_result_queue.put((False, url, t)) - break - - -if __name__ == "__main__": - file_info_cache = {} - try: - print(f"Attempting to load file records from cache: {file_database_path}") - with open(file_database_path, 'r') as f: - file_info_cache = json.loads(f.read()) - print(f"File records loaded from cache: {len(file_info_cache)} records found.") - except Exception as e: - print(f"Load failed, starting with empty cache") - file_info_cache = {} - - query_work_queue = queue.Queue() - query_result_queue = queue.Queue() - download_work_queue = queue.Queue() - download_result_queue = queue.Queue() - workers = [] - for _ in range(nfetchworkers): - t = Thread(target=find_links_worker, args=(query_work_queue, query_result_queue, 3, randomize_order), daemon=True) - t.start() - workers.append(t) - for _ in range(ndownloadworkers): - t = Thread(target=file_download_worker, args=(download_work_queue, download_result_queue), daemon=True) - t.start() - workers.append(t) - - try: - while True: - fetched_image_count = 0 - already_had_image_count = 0 - failed_image_count = 0 - urllen = len(directory_url) - - query_work_queue.put(directory_url) - for l, t, s in tqdm.tqdm(iter(partial(query_result_queue.get, timeout=30.0), None), desc="Downloading files"): - # Collect completed jobs and record completion status - while True: - try: - r_success, r_url, r_t = download_result_queue.get_nowait() - if r_success: - fetched_image_count += 1 - file_info_cache[r_url] = r_t - else: - failed_image_count += 1 - except queue.Empty: - break - - file_portion_of_link = l[urllen:] - filepath = os.path.join(stored_images_dir, file_portion_of_link.replace("/", os.sep)) - # If we dont have the file or the file at the link is newer than the one we previously fetched - if (not (l in file_info_cache)) or t > file_info_cache[l]: - if filepath.endswith(".fits"): - filepath2 = filepath.split(".fits")[0] + "_f.fits" # also look for the filtered version of the file - filepath3 = filepath.split(".fits")[0] + "_e.fits" # also look for the error version of the file - if os.path.exists(filepath2) or os.path.exists(filepath3): # If the unfiltered filename exists, that case will be handled in the alternative code path in if os.path.exists(filepath): - if l in file_info_cache: # If we have record of this file, it must be out of date, delete it and download the new version. - try: - os.remove(filepath2) - os.remove(filepath3) - except: - pass - else: # If we have no record of this file, update the file info cache and don't redownload - file_info_cache[l] = t - already_had_image_count += 1 - continue - if os.path.exists(filepath): - if l in file_info_cache: # If we have record of this file, it must be out of date, delete it and download the new version. - if os.path.exists(filepath): - os.remove(filepath) - else: # If we have no record of this file, update the file info cache and don't redownload if fsize is right - fsize = os.path.getsize(filepath) - if math.isclose(fsize, s, rel_tol=0.05): - file_info_cache[l] = t - already_had_image_count += 1 - continue - else: - print(f'Found a mismatched size on file: {filepath} Redownloading!') - download_work_queue.put((l, filepath, t, s)) - else: - # We have a download record, confirm the file actually exists on disk - if os.path.exists(filepath): - already_had_image_count += 1 - elif filepath.endswith(".fits"): - filepath2 = filepath.split(".fits")[0] + "_f.fits" # also look for the filtered version of the file - filepath3 = filepath.split(".fits")[0] + "_e.fits" # also look for the error version of the file - if os.path.exists(filepath2) or os.path.exists(filepath3): - already_had_image_count += 1 - else: # We could not find the file on disk, queue for redownload - download_work_queue.put((l, filepath, t)) - - with open(file_database_path, 'w') as f: - f.write(json.dumps(file_info_cache)) - print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}") - if fetch_interval > 0: - print(f"Run complete!, sleeping for {fetch_interval} seconds.") - with open(file_database_path, 'w') as f: - f.write(json.dumps(file_info_cache)) - time.sleep(fetch_interval) - else: - print("Run complete!, exiting...") - break - except KeyboardInterrupt: - print("Saving file database and shutting down.") - except Empty: - print("Work Complete, Shutting down...") - except Exception as e: - print(f"Unhandled Exception during run: {e}") - print("Shutting down") - - with open(file_database_path, 'w') as f: - f.write(json.dumps(file_info_cache)) - - for _ in range(nfetchworkers): - try: - query_work_queue.put(None, timeout=5.0) - except: - break - - time.sleep(1) - print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}") - - for _ in range(ndownloadworkers): - try: - download_work_queue.put(None, timeout=5.0) - except: - break - - for w in workers: +import os +import urllib.request +import urllib.parse +import re +import time +import random +import datetime +import json +from threading import Thread +import queue +import math +from functools import partial +from queue import Empty + +import tqdm + +directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/" +stored_images_dir = os.path.abspath(os.path.join("..", "Data")) +# directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/goes16/l2/data/ephe-l2-orb1m/2017/12/" +# stored_images_dir = r"Z:\NOAA GOES Data\Data\goes16/l2/data/ephe-l2-orb1m/2017/12/" +ignore_folder_names = ["Parent Directory", "l1b", "goes17", "2017", "2018", "2019", "2020", "2021", "2022", "2023"] +file_database_path = os.path.abspath(os.path.join("..", "file_database.json")) +fetch_interval = 0 # 60*60 +nfetchworkers = 2 # Be nice to the servers, this value is how many threads will be asking for links and file info at the same time +ndownloadworkers = 3 # Be nice to the servers, this value is how many threads will be downloading files at the same time +randomize_order = False + +links_regex_pattern = r'(?<=.*\d{4}-\d{2}-\d{2} \d{2}:\d{2})' # Find href links that do not contain question marks or whitespace +times_regex_pattern = r'(?<=)(\d{4}-\d{2}-\d{2} \d{2}:\d{2})(?= )' # Find timestamps in UTC in the format YYYY-MM-DD HH:mm +sizes_regex_pattern = r'(?<=\d{4}-\d{2}-\d{2} \d{2}:\d{2} <\/td>)([ .\d]{1,3})([KMG]|- |0 | )(?=<\/td>)' # Find file size markers in one of many possible formats +links_matcher = re.compile(links_regex_pattern) +times_matcher = re.compile(times_regex_pattern) +sizes_matcher = re.compile(sizes_regex_pattern) + +def find_links_worker(query_work_queue, query_result_queue, attempt_count = 3, randomize_order = True): + while True: + job = query_work_queue.get() + if job == None: + return + url= job + + attempts = 0 + html_content = None + while attempts < attempt_count: + try: + with urllib.request.urlopen(url) as response: + html_content = response.read().decode('utf-8') + break + except Exception as e: + tqdm.tqdm.write(f"Exception while fetching links: {e}") + time.sleep(1 + random.random()) + attempts += 1 + if (attempt_count > 1) and (attempts == attempt_count): + tqdm.tqdm.write(f'\nAfter {attempt_count} retries, could not fetch: {url}') + + if html_content is None: # We failed to fetch the link for some reason, continue to the next job + continue + else: + links = links_matcher.findall(html_content) + times = times_matcher.findall(html_content) + sizes = sizes_matcher.findall(html_content) + for i in range(len(times)): + dt = datetime.datetime.strptime(times[i].strip(), "%Y-%m-%d %H:%M") + dt.replace(tzinfo=datetime.timezone.utc) + times[i] = time.mktime(dt.timetuple()) + for i in range(len(sizes)): + match sizes[i][1].strip(): + case '-': + sizes[i] = 0 + case 'K': + sizes[i] = int(float(sizes[i][0].strip()) * 1024) + case 'M': + sizes[i] = int(float(sizes[i][0].strip()) * 1024 * 1024) + case 'G': + sizes[i] = int(float(sizes[i][0].strip()) * 1024 * 1024 * 1024) + case "0": + sizes[i] = 0 + case "": + sizes[i] = int(float(sizes[i][0].strip())) + case _: + raise(ValueError(f"Unexpected symbol while parsing links page: {_}")) + + + if (len(links) != len(times)) or (len(times) != len(sizes)): + raise(ValueError("Links parsing error, mismatched numbers of links, times, or sizes!")) + + results = list(zip(links, times, sizes)) + if randomize_order: + random.shuffle(results) + for _link, _time, _size in results: + if _link.endswith("/"): + if _link.split(r"/")[-2] in ignore_folder_names: + continue + else: + # print(f"Queueing folder: {urllib.parse.urljoin(url,_link)}") + query_work_queue.put(urllib.parse.urljoin(url,_link)) + else: + # print(f"Queueing file: {urllib.parse.urljoin(url,_link)}") + query_result_queue.put((urllib.parse.urljoin(url,_link), _time, _size)) + + +# Fetch file from url and store it to path, retrying on failure +def file_download_worker(download_work_queue, download_result_queue, attempt_count = 2): + while True: + job = download_work_queue.get() + if job == None: + return + url, path, t, s = job + + # print(f"Next job: {url}, {path}, {t}, {s}") + attempts = 0 + while attempts < attempt_count: + try: + starttime = time.time() + req = urllib.request.Request(url, data=None) + image_data = urllib.request.urlopen(req, timeout=10.0).read() + endtime = time.time() + tqdm.tqdm.write(f"Downloaded {url} in {endtime-starttime:0.2f} s | {len(image_data)/1024/1024/(endtime-starttime):0.2f} MB/s") + if math.isclose(len(image_data), s, rel_tol=0.05): + os.makedirs(os.path.split(path)[0], exist_ok=True) + open(path, 'wb').write(image_data) + download_result_queue.put((True, url, t)) + break + else: + raise ValueError("Downloaded file is the wrong size!") + except Exception as e: + if hasattr(e, "code") and e.code == 404: # This is expected if the file has been removed from the site (at least for swpc.noaa.gov) + attempts += attempt_count + elif (0 < attempts) and (attempts < attempt_count): + # tqdm.tqdm.write(f"\nA problem occurred on file: {url} | {e}") + time.sleep(1 + random.random()) + attempts += 1 + if (attempt_count > 1) and (attempts == attempt_count): + tqdm.tqdm.write(f'\nAfter {attempt_count} retries, could not fetch: {url}') + tqdm.tqdm.write(f"Exception: {e}") + download_result_queue.put((False, url, t)) + break + + +if __name__ == "__main__": + file_info_cache = {} + try: + print(f"Attempting to load file records from cache: {file_database_path}") + with open(file_database_path, 'r') as f: + file_info_cache = json.loads(f.read()) + print(f"File records loaded from cache: {len(file_info_cache)} records found.") + except Exception as e: + print(f"Load failed, starting with empty cache") + file_info_cache = {} + + query_work_queue = queue.Queue() + query_result_queue = queue.Queue() + download_work_queue = queue.Queue(maxsize=ndownloadworkers) + download_result_queue = queue.Queue(maxsize=ndownloadworkers) + workers = [] + for _ in range(nfetchworkers): + t = Thread(target=find_links_worker, args=(query_work_queue, query_result_queue, 3, randomize_order), daemon=True) + t.start() + workers.append(t) + for _ in range(ndownloadworkers): + t = Thread(target=file_download_worker, args=(download_work_queue, download_result_queue), daemon=True) + t.start() + workers.append(t) + + try: + while True: + fetched_image_count = 0 + already_had_image_count = 0 + failed_image_count = 0 + urllen = len(directory_url) + + query_work_queue.put(directory_url) + for l, t, s in tqdm.tqdm(iter(partial(query_result_queue.get, timeout=30.0), None), desc="Downloading files"): + # Collect completed jobs and record completion status + while True: + try: + r_success, r_url, r_t = download_result_queue.get_nowait() + if r_success: + fetched_image_count += 1 + file_info_cache[r_url] = r_t + else: + failed_image_count += 1 + except queue.Empty: + break + + file_portion_of_link = l[urllen:] + filepath = os.path.join(stored_images_dir, file_portion_of_link.replace("/", os.sep)) + # If we dont have the file or the file at the link is newer than the one we previously fetched + if (not (l in file_info_cache)) or t > file_info_cache[l]: + if filepath.endswith(".fits"): + filepath2 = filepath.split(".fits")[0] + "_f.fits" # also look for the filtered version of the file + filepath3 = filepath.split(".fits")[0] + "_e.fits" # also look for the error version of the file + if os.path.exists(filepath2) or os.path.exists(filepath3): # If the unfiltered filename exists, that case will be handled in the alternative code path in if os.path.exists(filepath): + if l in file_info_cache: # If we have record of this file, it must be out of date, delete it and download the new version. + try: + os.remove(filepath2) + os.remove(filepath3) + except: + pass + else: # If we have no record of this file, update the file info cache and don't redownload + file_info_cache[l] = t + already_had_image_count += 1 + continue + if os.path.exists(filepath): + if l in file_info_cache: # If we have record of this file, it must be out of date, delete it and download the new version. + if os.path.exists(filepath): + os.remove(filepath) + else: # If we have no record of this file, update the file info cache and don't redownload if fsize is right + fsize = os.path.getsize(filepath) + if math.isclose(fsize, s, rel_tol=0.05): + file_info_cache[l] = t + already_had_image_count += 1 + continue + else: + tqdm.tqdm.write(f'Found a mismatched size on file: {filepath} Redownloading!') + download_work_queue.put((l, filepath, t, s)) + else: + # We have a download record, confirm the file actually exists on disk + if os.path.exists(filepath): + already_had_image_count += 1 + elif filepath.endswith(".fits"): + filepath2 = filepath.split(".fits")[0] + "_f.fits" # also look for the filtered version of the file + filepath3 = filepath.split(".fits")[0] + "_e.fits" # also look for the error version of the file + if os.path.exists(filepath2) or os.path.exists(filepath3): + already_had_image_count += 1 + else: # We could not find the file on disk, queue for redownload + download_work_queue.put((l, filepath, t)) + + with open(file_database_path, 'w') as f: + f.write(json.dumps(file_info_cache)) + print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}") + if fetch_interval > 0: + print(f"Run complete!, sleeping for {fetch_interval} seconds.") + with open(file_database_path, 'w') as f: + f.write(json.dumps(file_info_cache)) + time.sleep(fetch_interval) + else: + print("Run complete!, exiting...") + break + except KeyboardInterrupt: + print("Saving file database and shutting down.") + except Empty: + print("Work Complete, Shutting down...") + except Exception as e: + print(f"Unhandled Exception during run: {e}") + print("Shutting down") + + for _ in range(nfetchworkers): + try: + query_work_queue.put(None) + except: + break + + time.sleep(1) + print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}") + + while len(download_work_queue.queue) > 0: + print(f"Waiting for {len(download_work_queue.queue)} downloads in queue...") + time.sleep(1) + + for _ in range(ndownloadworkers): + try: + download_work_queue.put(None, timeout=5.0) + except: + break + + with open(file_database_path, 'w') as f: + f.write(json.dumps(file_info_cache)) + + print("Waiting for workers to shutdown...") + + for w in workers: w.join(5.0) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..3fc1061 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,8 @@ +astropy +scikit-image +tqdm +numpy +opencv-python +palettable +matplotlib +pillow \ No newline at end of file