diff --git a/merger_FITS.py b/merger_FITS.py index fc703f8..0c66553 100644 --- a/merger_FITS.py +++ b/merger_FITS.py @@ -433,16 +433,15 @@ def generate_composite(work_queue, result_queue): 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"] - stored_fits_dirs = [r"..\Data\goes16",] - processed_images_dir = r"..\composite\goes16" - 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()) + 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] @@ -463,9 +462,9 @@ if __name__ == "__main__": nfailed = 0 black_image = Image.fromarray(np.zeros((1080, 1920, 3), dtype='uint8')) try: - files_sorted_by_timestamp = defaultdict(list) - num_found_files = 0 - for stored_fits_dir in stored_fits_dirs: + 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}") @@ -477,98 +476,88 @@ if __name__ == "__main__": # 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_sorted_by_timestamp[measure_end_time].append(os.path.join(root,f)) + 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) + print(f"Found {num_found_files} FITS files. Starting conversion.") + if num_found_files == 0: + exit(3) - sorted_times = sorted(list(files_sorted_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 + 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_sorted_by_timestamp[timestamp]) == 6: - f = i - break - for i, timestamp in enumerate(reversed(sorted_times)): - if len(files_sorted_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_sorted_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: + 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 - - # 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): - 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]) + 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"Detected a gap of {time_gap} frames at {timestamp}, inserting black frames.") - filename = f"Composite-{int(timestamp)}_b.jpg" - filepath = os.path.join(processed_images_dir, filename) - if not os.path.isfile(filepath): - black_image.save(filepath, quality = 95) - break - 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)) + 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")