126 lines
5.8 KiB
Python
126 lines
5.8 KiB
Python
import argparse
|
|
import subprocess
|
|
import os
|
|
import sys
|
|
import datetime
|
|
import shutil as _shutil
|
|
import calendar
|
|
|
|
import numpy as np
|
|
from sortedcontainers import SortedDict
|
|
import tqdm
|
|
from PIL import Image
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
from suvi import paths, vfs
|
|
|
|
|
|
def find_ffmpeg():
|
|
"""Locate ffmpeg: $FFMPEG first, then PATH, then the bundled Windows binary."""
|
|
override = os.environ.get("FFMPEG")
|
|
if override and os.path.exists(override):
|
|
return override
|
|
found = _shutil.which("ffmpeg")
|
|
if found:
|
|
return found
|
|
bundled = os.path.join(os.path.dirname(paths.data_root()), "ffmpeg.exe")
|
|
if os.path.exists(bundled):
|
|
return bundled
|
|
raise SystemExit("ffmpeg not found. Install it, or set $FFMPEG to its path.")
|
|
|
|
|
|
def encode_videos(args):
|
|
"""Encode composite JPGs into per-satellite videos.
|
|
|
|
Gaps are padded with black frames and smoothed by ffmpeg's minterpolate, except
|
|
where the gap exceeds `max_frame_interp`, which is skipped entirely. That
|
|
behaviour is what the bench's "today" variant reproduces.
|
|
"""
|
|
years = args.years
|
|
sources = args.sources
|
|
ffmpeg_path = find_ffmpeg()
|
|
project_root = args.out or os.path.dirname(paths.data_root())
|
|
os.makedirs(project_root, exist_ok=True)
|
|
|
|
for year in years:
|
|
for source in sources:
|
|
path_to_images = args.images or os.path.join(
|
|
os.path.dirname(paths.data_root()), "composite", source)
|
|
output_file = os.path.join(project_root, f"{source}_{year}_{args.suffix}.mp4")
|
|
interp_file = os.path.join(project_root, f"{source}_{year}_interp_{args.suffix}.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
|
|
reliever = vfs.Reliever(label=f"encode {source}")
|
|
for t, f in tqdm.tqdm(files_by_time.items(), desc="FFMPEG encoding images", total = len(files_by_time)):
|
|
reliever.tick()
|
|
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
|
|
pipe.read()
|
|
pipe.close()
|
|
|
|
def build_parser():
|
|
parser = argparse.ArgumentParser(
|
|
description="Encode SUVI composite images into videos.",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
)
|
|
parser.add_argument("--years", type=int, nargs="+", default=[2024])
|
|
parser.add_argument("--sources", nargs="+", default=["goes16", "goes18"])
|
|
parser.add_argument("--images", default=None,
|
|
help="directory of composites (default: <root>/composite/<source>)")
|
|
parser.add_argument("--out", default=None, help="directory for the video files")
|
|
parser.add_argument("--start", type=int, default=None, help="Unix time, inclusive")
|
|
parser.add_argument("--stop", type=int, default=None, help="Unix time, exclusive")
|
|
parser.add_argument("--suffix", default="nofilt", help="tag in the output filenames")
|
|
return parser
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Guarded: this module is imported for find_ffmpeg(), and without the guard that
|
|
# import kicked off a full-year encode as a side effect.
|
|
sys.exit(encode_videos(build_parser().parse_args()) or 0)
|