noaa-goes-visualization/merger_FITS.py

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(2023, 1, 1, tzinfo=datetime.timezone.utc).timetuple())
    stoptime = calendar.timegm(datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc).timetuple())

    regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*\_f.fits$"
    nworkers = 20
    max_time_gap = 3
    fill_missing_data = False

    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} - {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}")
                                    break
                        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()

    print(f"Created {ncreated} | Already had {nexists} | Failed {nfailed}")