Rewrote trimming for 1920x1080 and updated to use new filtered filenames

merger_FITS.py

@@ -6,6 +6,7 @@ from multiprocessing import Queue, Process
 import re
 import warnings
 import queue
+import traceback
 
 import tqdm
 from PIL import Image, ImageDraw
@@ -29,7 +30,7 @@ def lowpriority():
         # Based on:
         #   "Recipe 496767: Set Process Priority In Windows" on ActiveState
         #   http://code.activestate.com/recipes/496767/
-        import win32api,win32process,win32con
+        import win32api,win32process,win32con # pywin32
 
         pid = win32api.GetCurrentProcessId()
         handle = win32api.OpenProcess(win32con.PROCESS_ALL_ACCESS, True, pid)
@@ -266,10 +267,8 @@ def hsl_to_rgb(img):
     else:
         return np.stack([r, g, b], -1)
 
-
 def generate_composite(work_queue, result_queue):
     lowpriority()
-    image_names = ["094A", "131A", "171A", "195A", "284A", "304A"]
     cmaps =[palettable.cmocean.sequential.Ice_5.mpl_colormap, 
             palettable.cmocean.sequential.Ice_20.mpl_colormap, 
             palettable.cmocean.sequential.Turbid_5_r.mpl_colormap, 
@@ -285,14 +284,19 @@ def generate_composite(work_queue, result_queue):
     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
     while True:
         try:
             job = work_queue.get()
             if job is None:
                 result_queue.cancel_join_thread()
                 return
-            files_this_timestamp, timestamp, processed_images_dir = job
+            files_this_timestamp, timestamp, processed_images_dir, synthetic_data = job
-            filename = f"Composite-{int(timestamp)}.jpg"
+            if synthetic_data:
+                filename = f"Composite-{int(timestamp)}_s.jpg"
+            else:
+                filename = f"Composite-{int(timestamp)}.jpg"
             filepath = os.path.join(processed_images_dir, filename)
             if os.path.isfile(filepath):
                 result_queue.put(("Exists", timestamp))
@@ -300,16 +304,15 @@ def generate_composite(work_queue, result_queue):
 
             base_imgs = []
             for i in range(6):
-                raw_data = np.flip(fits.getdata(files_this_timestamp[i]), 0)
+                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]))
+                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)
+            composite_image_data = composite_alpha_over(base_imgs[4], base_imgs[5], 0.2)**0.5
-            composite_image_data = composite_image_data**0.5
             # plt.imshow(composite_image_data)
 
             # plt.figure("Linear Burn with 304")
@@ -318,8 +321,7 @@ def generate_composite(work_queue, result_queue):
 
             # plt.figure("Light Ops with mid bands")
             mix_img = composite_alpha_over(exclusion(base_imgs[3], composite_image_data), composite_image_data, 0.95)
-            mix_img = composite_alpha_over(linear_light(base_imgs[2], mix_img), mix_img, 1.0)
+            mix_img = composite_alpha_over(linear_light(base_imgs[2], mix_img), mix_img, 1.0)**0.75
-            mix_img = mix_img**0.75
             composite_image_data = composite_alpha_over(mix_img, composite_image_data, 0.5)
             # plt.imshow(composite_image_data)
 
@@ -385,50 +387,62 @@ def generate_composite(work_queue, result_queue):
             # We have our final image
             # plt.show()
 
-            # Trim edges of final image so it fits nicely in 1920
-            composite_image_data = composite_image_data[64:-64,64:-64,:3]
-
             # Now shrink the component images and assemble them alongside the composite.
-            new_dim = composite_image_data.shape[0] // 3
+            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_dim, new_dim),(0,0)))
+            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][64:-64,64:-64,:3] # Trim edges of image data to fit nicely
+                img = base_imgs[i]
-                img = bin_ndarray(img, (new_dim, new_dim, 3)) # Shrink down to 1/3 for assembly
+                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_dim)
+                xdimoff = i%2 * (composite_image_data.shape[1] - new_dimx)
-                ydimoff = i//2*new_dim
+                ydimoff = i//2*new_dimy
-                composite_image_data[ydimoff:ydimoff+new_dim, xdimoff:xdimoff+new_dim, :] = img
+                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((655, 15), f"NOAA GOES Sattelite SUVI Composite - {timestring} UTC",(255,255,255), font_size = 24)
+            ImageDraw.Draw(img).text((605, 15), f"NOAA GOES Satellite SUVI Composite - {timestring} UTC",(255,255,255), font_size = 24)
             for i in range(6): # Draw component angstrom labels
                 if i%2 == 0:
                     xdimtxtoff = 5
                 else:
                     xdimtxtoff = composite_image_data.shape[1] - 44
-                ydimtxtoff = i//2*new_dim + new_dim / 2.0 - 8
+                ydimtxtoff = i//2*new_dimy + new_dimy / 2.0 - 6
                 ImageDraw.Draw(img).text((xdimtxtoff, ydimtxtoff), image_names[i], font_size = 16)
-            img.save(filepath, quality = 90)
+            img.save(filepath, quality = 95)
             result_queue.put(("Created", timestamp))
         
         except KeyboardInterrupt:
             return
         except Exception as e:
+            traceback.print_exception(e)
             result_queue.put((e, timestamp))
 
 
-if __name__ == "__main__":
+image_names = ["094A", "131A", "171A", "195A", "284A", "304A"]
-    stored_fits_dirs = [r"..\Data\goes16\l2\data", r"..\Data\goes18\l2\data"]
-    processed_images_dirs = [r"..\composite\goes16", r"..\composite\goes18"]
 
-    regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*\.fits"
+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"]
+    stored_fits_dirs = [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",]
+    processed_images_dir = r"..\composite\goes18"
+    starttime = time.mktime(datetime.datetime(2024, 1, 1).timetuple())
+    stoptime = time.mktime(datetime.datetime(2025, 1, 1).timetuple())
+
+    regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*\_f.fits$"
     nworkers = 20
+    max_time_gap = 10
+
+    file_prefixes = ["dr_suvi-l2-ci" + n[:-1] for n in image_names]
     
     # Testing
     # stored_fits_dirs = [r"..\fits_test_2024"]
-    # processed_images_dirs = [r"..\composite"]
 
     work_queue = Queue(maxsize = nworkers)
     result_queue = Queue()
@@ -443,12 +457,11 @@ if __name__ == "__main__":
     nexists = 0
     nfailed = 0
     try:
-        for stored_fits_dir, processed_images_dir in zip(stored_fits_dirs, processed_images_dirs):
+        files_sorted_by_timestamp = defaultdict(list)
+        found_files = 0
+        for stored_fits_dir in stored_fits_dirs:
             os.makedirs(processed_images_dir, exist_ok=True)
             filename_tester = re.compile(regex_filename)
-
-            files_sorted_by_timestamp = defaultdict(list)
-            found_files = 0
             print(f"Searching for FITS files in: {stored_fits_dir}")
             for root, dirs, files in tqdm.tqdm(os.walk(stored_fits_dir), desc="Searching"):
                 for f in files:
@@ -458,35 +471,79 @@ if __name__ == "__main__":
                         sattelite = file_parts[2]
                         measure_end_time = datetime.datetime.strptime(file_parts[4][1:16], "%Y%m%dT%H%M%S")
                         measure_end_time.replace(tzinfo=datetime.timezone.utc)
-                        measure_end_time = time.mktime(measure_end_time.timetuple())
+                        measure_end_time = int(time.mktime(measure_end_time.timetuple()))
                         files_sorted_by_timestamp[measure_end_time].append(os.path.join(root,f))
                         found_files += 1
                     
-            print(f"Found {found_files} FITS files. Starting conversion.")
+        print(f"Found {found_files} FITS files. Starting conversion.")
 
-            for timestamp in tqdm.tqdm(files_sorted_by_timestamp, desc="Creating Composite Solar Images"):
+        sorted_times = sorted(list(files_sorted_by_timestamp.keys()))
-                # Collect completed jobs and record completion status
+        min_time = sorted_times[0]
-                while True:
+        max_time = sorted_times[-1]
-                    try:
+        diff_times = np.diff(sorted_times)
-                        result = result_queue.get_nowait()
+        unique, counts = np.unique(diff_times, return_counts=True)
-                        if result[0] == "Exists":
+        interval = unique[0] # This is the amount of time between each sample in seconds
-                            nexists += 1
+        assert np.sum((unique % interval) > 0) == 0 # Ensure all our timestamps align perfectly with our interval
-                        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
+        last_good_files = None
-                files_this_timestamp = files_sorted_by_timestamp[timestamp]
+        last_good_file_times = None
-                files_this_timestamp = sorted(files_this_timestamp)
+        for timestamp in tqdm.tqdm(sorted_times, desc="Creating Composite Solar Images"):
-                if not len(files_this_timestamp) == 6:
+            if timestamp < min_time or timestamp > max_time:
-                    print(f"Invalid or incomplete sensor records for: {timestamp}")
+                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_sorted_by_timestamp[timestamp]
+            files_this_timestamp = sorted(files_this_timestamp)
+            synthetic_data = False
+            if (not len(files_this_timestamp) == 6):
+                if (not last_good_files):
+                    print(f"Invalid or incomplete sensor records for {timestamp} - {len(files_this_timestamp)}/6 and no last-good data.")
                     continue
-                
+                else:
-                work_queue.put((files_this_timestamp, timestamp, processed_images_dir))
+                    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:
+                                synthetic_data = True
+                                files_for_job.append(last_good_files[i])
+                            else:
+                                print(f"Detected a gap of {time_gap} frames at {timestamp}, skipping.")
+                                continue
+                    if len(files_for_job) == 6:
+                        work_queue.put((files_for_job, timestamp, processed_images_dir, synthetic_data))
+                    else:
+                        # We did not get a full file set to process, because we were missing one or more files and also exceeeded time_gap limits
+                        pass
+            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]
+                if len(files_this_timestamp) != 6:
+                    print("!!!!!")
+                work_queue.put((files_this_timestamp, timestamp, processed_images_dir, synthetic_data))
 
     except KeyboardInterrupt:
         print("Finishing current jobs and exiting")
@@ -498,6 +555,7 @@ if __name__ == "__main__":
             break
 
     for w in workers:
-        w.join(10.0)
+        # w.join(10.0)
+        w.join()
 
     print(f"Created {ncreated} | Already had {nexists} | Failed {nfailed}")