TAFG good result

tool/dump_tensorboard.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python
 
+# edit from https://gist.github.com/hysts/81a0d30ac4f33dfa0c8859383aec42c2
+
 import argparse
 import pathlib
 
@@ -8,33 +10,36 @@ import numpy as np
 from tensorboard.backend.event_processing import event_accumulator
 
 
+def save(outdir: pathlib.Path, tag, event_acc):
+    events = event_acc.Images(tag)
+
+    for index, event in enumerate(events):
+        s = np.frombuffer(event.encoded_image_string, dtype=np.uint8)
+        image = cv2.imdecode(s, cv2.IMREAD_COLOR)
+        outpath = outdir / f"{tag.replace('/', '_')}@{index}.png"
+        cv2.imwrite(outpath.as_posix(), image)
+
+
 # ffmpeg -framerate 1 -i ./tmp/test_b/%04d.jpg -vcodec mpeg4 test_b.mp4
-# https://gist.github.com/hysts/81a0d30ac4f33dfa0c8859383aec42c2
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--path', type=str, required=True)
     parser.add_argument('--outdir', type=str, required=True)
+    parser.add_argument("--tag", type=str, required=False)
     args = parser.parse_args()
 
-    event_acc = event_accumulator.EventAccumulator(
-        args.path, size_guidance={'images': 0})
+    event_acc = event_accumulator.EventAccumulator(args.path, size_guidance={'images': 0})
     event_acc.Reload()
 
     outdir = pathlib.Path(args.outdir)
     outdir.mkdir(exist_ok=True, parents=True)
 
-    for tag in event_acc.Tags()['images']:
-        events = event_acc.Images(tag)
-
-        tag_name = tag.replace('/', '_')
-        dirpath = outdir / tag_name
-        dirpath.mkdir(exist_ok=True, parents=True)
-
-        for index, event in enumerate(events):
-            s = np.frombuffer(event.encoded_image_string, dtype=np.uint8)
-            image = cv2.imdecode(s, cv2.IMREAD_COLOR)
-            outpath = dirpath / '{:04}.jpg'.format(index)
-            cv2.imwrite(outpath.as_posix(), image)
+    if args.tag is None:
+        for tag in event_acc.Tags()['images']:
+            save(outdir, tag, event_acc)
+    else:
+        assert args.tag in event_acc.Tags()['images'], f"{args.tag} not in {event_acc.Tags()['images']}"
+        save(outdir, args.tag, event_acc)
 
 
 if __name__ == '__main__':