Created
February 11, 2025 22:30
-
-
Save mosure/1dc89a4052ca4a275d75b4c3ef7b211c to your computer and use it in GitHub Desktop.
tensorboard_images_to_videos.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import glob | |
| import os | |
| import argparse | |
| import cv2 | |
| import numpy as np | |
| import tensorflow as tf | |
| from tqdm import tqdm | |
| def process_event_file(event_file, output_dir, tag_prefix=None, fps=24): | |
| video_writers = {} | |
| frame_counts = {} | |
| with tqdm(desc=f"Processing {os.path.basename(event_file)}", unit="evt", | |
| miniters=100, mininterval=1.0, dynamic_ncols=True) as pbar: | |
| for event in tf.compat.v1.train.summary_iterator(event_file): | |
| if not hasattr(event, "summary"): | |
| pbar.update(1) | |
| continue | |
| for value in event.summary.value: | |
| tag = value.tag | |
| if tag_prefix and not tag.startswith(tag_prefix): | |
| continue | |
| if not value.HasField("image"): | |
| continue | |
| encoded = value.image.encoded_image_string | |
| if not encoded: | |
| continue | |
| nparr = np.frombuffer(encoded, np.uint8) | |
| img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) | |
| if img is None: | |
| continue | |
| if tag not in video_writers: | |
| height, width, _ = img.shape | |
| safe_tag = tag.replace('/', '_') | |
| output_filename = f"{safe_tag}.mp4" | |
| output_filepath = os.path.join(output_dir, output_filename) | |
| fourcc = cv2.VideoWriter_fourcc(*'mp4v') | |
| writer = cv2.VideoWriter(output_filepath, fourcc, fps, (width, height)) | |
| if not writer.isOpened(): | |
| print(f"Error opening video writer for {output_filepath}") | |
| continue | |
| video_writers[tag] = writer | |
| frame_counts[tag] = 0 | |
| video_writers[tag].write(img) | |
| frame_counts[tag] += 1 | |
| pbar.update(1) | |
| for tag, writer in video_writers.items(): | |
| writer.release() | |
| print(f"Saved video for tag '{tag}' with {frame_counts[tag]} frames from file '{event_file}'") | |
| def main(): | |
| parser = argparse.ArgumentParser( | |
| description="Stream image events from TensorBoard event files to MP4 videos." | |
| ) | |
| parser.add_argument("--logdir", required=True, | |
| help="Path to TensorBoard log directory (will search recursively for event files)") | |
| parser.add_argument("--output_dir", required=True, | |
| help="Directory to output the generated MP4 videos") | |
| parser.add_argument("--tag_prefix", help="Only process image events with tags that start with this prefix") | |
| parser.add_argument("--fps", type=int, default=24, | |
| help="Frames per second for the output video (default: 24)") | |
| args = parser.parse_args() | |
| if not os.path.exists(args.output_dir): | |
| os.makedirs(args.output_dir) | |
| event_files = glob.glob(os.path.join(args.logdir, "**", "events.out.tfevents.*"), recursive=True) | |
| if not event_files: | |
| print("No event files found in the specified logdir.") | |
| return | |
| for event_file in event_files: | |
| process_event_file(event_file, args.output_dir, args.tag_prefix, args.fps) | |
| if __name__ == "__main__": | |
| tf.compat.v1.disable_eager_execution() | |
| main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment