noaa-goes-visualization/filter_FITS.py

import os
from collections import defaultdict
from multiprocessing import Queue, Process
import re
import traceback

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 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)

measurement_names = ["094", "131", "171", "195", "284", "304"]
thresholds =        [0.050, 0.10 ,  1.00,  1.40,  0.90,  2.50]
max_center_skew = 7
circle_hough_thresh = 0.7
expected_dims = 1280
ratio_above_thresh_max = 0.5

def filter_fits(work_queue):
    lowpriority()
    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)
            data = fits.getdata(job)
            assert data.shape[0] == expected_dims
            assert data.shape[1] == expected_dims
            filtered_data = np.zeros_like(data)
            above_thresh_indexes = gaussian(data, 3) > thresholds[idx]
            filtered_data[above_thresh_indexes] = 1
            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(f"Data {measurement}")
                # plt.imshow(data, cmap='jet')
                # plt.figure(f"Circle")
                # plt.imshow(filtered_data, cmap='jet')
                # plt.show()

                continue

            hough_radii = np.arange(387, 394, 1)
            hough_res = hough_circle(filtered_data, hough_radii)

            vals, cxs, cys, rads = hough_circle_peaks(hough_res, hough_radii, threshold=circle_hough_thresh, total_num_peaks=10)
            found_circle = None
            for v, cx, cy, rad in zip(vals, cxs, cys, rads):
                xskew = abs(639 - cx)
                yskew = abs(639 - cy)
                if xskew < max_center_skew and yskew < max_center_skew:
                    found_circle = (cx, cy, rad)
                    break

            if not found_circle:
                print(f"Could not find valid solar disc in file: {job}")

                # plt.figure(f"Data {measurement}")
                # if found_circle: cv.circle(data, (int(found_circle[0]),int(found_circle[1])), int(found_circle[2]), float(np.max(np.max(data))), 1)
                # plt.imshow(data, cmap='jet')
                # plt.figure(f"Circle")
                # if found_circle: cv.circle(filtered_data, (int(found_circle[0]),int(found_circle[1])), int(found_circle[2]), 0.5, 1)
                # plt.imshow(filtered_data, cmap='jet')
                # plt.show()

                new_name = job.split(".fits")[0] + "_e.fits"
                os.rename(job, new_name)
                continue

            # We have validated this file, rename it appropriately
            new_name = job.split(".fits")[0] + "_f.fits"
            os.rename(job, new_name)


        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",
                        r"..\Data\goes16\l2\data\suvi-l2-ci131",
                        r"..\Data\goes16\l2\data\suvi-l2-ci171",
                        r"..\Data\goes16\l2\data\suvi-l2-ci195",
                        r"..\Data\goes16\l2\data\suvi-l2-ci284",
                        r"..\Data\goes16\l2\data\suvi-l2-ci304",
                        r"..\Data\goes18\l2\data\suvi-l2-ci094",
                        r"..\Data\goes18\l2\data\suvi-l2-ci131",
                        r"..\Data\goes18\l2\data\suvi-l2-ci171",
                        r"..\Data\goes18\l2\data\suvi-l2-ci195",
                        r"..\Data\goes18\l2\data\suvi-l2-ci284",
                        r"..\Data\goes18\l2\data\suvi-l2-ci304",]
    # stored_fits_dirs = [r"..\Data\goes16\l2\data"]
    # stored_fits_dirs = [r"Z:\NOAA GOES Data\fits_test_2024"]

    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_to_process = []
    try:
        for stored_fits_dir in stored_fits_dirs:
            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("_")
                        file_name_end = file_parts[-1].split(".")[0]
                        if file_name_end == "f":
                            continue # Already filtered from a previous run
                        elif file_name_end == "e":
                            if reprocess_errors:
                                new_file_name = "_".join(file_parts[:-1]) + ".fits"
                                os.rename(os.path.join(root,f), os.path.join(root,new_file_name))
                                files_to_process.append(os.path.abspath(os.path.join(root, new_file_name)))
                            else:
                                continue
                        elif file_name_end == "v1-0-2": # This is the normal case for unprocessed data
                            files_to_process.append(os.path.abspath(os.path.join(root,f)))
                        else:
                            # print(f"Error - Unexpected FITS file name: {f}")
                            pass
        for ftp in tqdm.tqdm(files_to_process, desc="Filtering files"):
            work_queue.put(ftp)
    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()