CVPR/lama-example

1#!/usr/bin/env python32 3import glob4import os5import re6 7import tensorflow as tf8from torch.utils.tensorboard import SummaryWriter9 10 11GROUPING_RULES = [12    re.compile(r'^(?P<group>train|test|val|extra_val_.*?(256|512))_(?P<title>.*)', re.I)13]14 15 16DROP_RULES = [17    re.compile(r'_std$', re.I)18]19 20 21def need_drop(tag):22    for rule in DROP_RULES:23        if rule.search(tag):24            return True25    return False26 27 28def get_group_and_title(tag):29    for rule in GROUPING_RULES:30        match = rule.search(tag)31        if match is None:32            continue33        return match.group('group'), match.group('title')34    return None, None35 36 37def main(args):38    os.makedirs(args.outdir, exist_ok=True)39 40    ignored_events = set()41 42    for orig_fname in glob.glob(args.inglob):43        cur_dirpath = os.path.dirname(orig_fname)  # remove filename, this should point to "version_0" directory44        subdirname = os.path.basename(cur_dirpath)  # == "version_0" most of time45        exp_root_path = os.path.dirname(cur_dirpath)  # remove "version_0"46        exp_name = os.path.basename(exp_root_path)47 48        writers_by_group = {}49 50        for e in tf.compat.v1.train.summary_iterator(orig_fname):51            for v in e.summary.value:52                if need_drop(v.tag):53                    continue54 55                cur_group, cur_title = get_group_and_title(v.tag)56                if cur_group is None:57                    if v.tag not in ignored_events:58                        print(f'WARNING: Could not detect group for {v.tag}, ignoring it')59                        ignored_events.add(v.tag)60                    continue61 62                cur_writer = writers_by_group.get(cur_group, None)63                if cur_writer is None:64                    if args.include_version:65                        cur_outdir = os.path.join(args.outdir, exp_name, f'{subdirname}_{cur_group}')66                    else:67                        cur_outdir = os.path.join(args.outdir, exp_name, cur_group)68                    cur_writer = SummaryWriter(cur_outdir)69                    writers_by_group[cur_group] = cur_writer70 71                cur_writer.add_scalar(cur_title, v.simple_value, global_step=e.step, walltime=e.wall_time)72 73 74if __name__ == '__main__':75    import argparse76 77    aparser = argparse.ArgumentParser()78    aparser.add_argument('inglob', type=str)79    aparser.add_argument('outdir', type=str)80    aparser.add_argument('--include-version', action='store_true',81                         help='Include subdirectory name e.g. "version_0" into output path')82 83    main(aparser.parse_args())84