def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config', default=Path('./release/config.yaml'))
args = parser.parse_args()
with open(args.config, 'r') as f:
config = yaml.load(f)
output_root = Path(config['davis16']['output_dir'])
flow_output = output_root / 'flow'
flow_output.mkdir(exist_ok=True, parents=True)
common_setup(__file__, flow_output, args)
logging.debug('Config:\n%s', pprint.pformat(config))
split_dirs = link_splits(config)
for split in config['davis16']['splits']:
input_dir = split_dirs[split][0]
output_split = flow_output / split
msg("Computing flow on DAVIS 2016 %s set.")
compute_flow_helper(config,
input_dir,
output_split,
extensions=['.jpg'])
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config', default=Path('./release/config.yaml'))
parser.add_argument(
'--with-ytvos-train',
action='store_true',
help=("By default, we evaluate tracks from a model that was "
"not trained on YTVOS for fair evaluation of generalization. If "
"--with-ytvos-train is specified, infer with model that uses "
"YTVOS for training."))
args = parser.parse_args()
with open(args.config, 'r') as f:
config = yaml.load(f)
output_dir = Path(config['ytvos']['output_dir']) / 'evaluation' / (
'with_ytvos' if args.with_ytvos_train else 'without_ytvos')
output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, output_dir, args)
logging.info('Evaluating using proposed metric.')
evaluate_proposed(config, args.with_ytvos_train)
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config', default=Path('./release/config.yaml'))
args = parser.parse_args()
with open(args.config, 'r') as f:
config = yaml.load(f)
output_dir = Path(config['davis17']['output_dir']) / 'tracks'
output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, output_dir, args)
split_dirs = link_splits(config)
for split in config['davis17']['splits']:
image_dir = split_dirs[split][0]
init_detections = (Path(config['davis17']['output_dir']) /
'detections' / split)
output_split = Path(output_dir) / split
args = [
'--images-dir', image_dir, '--init-detections-dir',
init_detections, '--output-dir', output_split, '--save-numpy',
True, '--save-images', False, '--bidirectional',
'--score-init-min', 0.9, '--remove-continue-overlap', 0.1, '--fps',
30, '--filename-format', 'frame', '--save-video',
config['tracker']['visualize'], '--recursive'
]
cmd = ['python', 'tracker/two_detector_track.py'] + args
msg(f'Running tracker on DAVIS {split}')
subprocess_call(cmd)
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config', default=Path('./release/config.yaml'))
args = parser.parse_args()
with open(args.config, 'r') as f:
config = yaml.load(f)
output_dir = Path(config['fbms']['output_dir']) / 'flow'
output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, output_dir, args)
logging.debug('Config:\n%s', pprint.pformat(config))
for split in config['fbms']['splits']:
input_dir = config['fbms']['images_highres'][split]
output_split = output_dir / split
logging.info("\n\n###\n"
"Computing flow on FBMS %s set.\n"
"###\n\n", split)
compute_flow_helper(config,
input_dir,
output_split,
extensions=['.png'])
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('evaluation',
default='proposed',
choices=['official', 'proposed'])
parser.add_argument('--config', default=Path('./release/config.yaml'))
args = parser.parse_args()
with open(args.config, 'r') as f:
config = yaml.load(f)
output_dir = Path(config['fbms']['output_dir']) / 'evaluation'
output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, output_dir, args)
if args.evaluation == 'proposed':
logging.info('Evaluating using proposed metric.')
evaluate_proposed(config)
elif args.evaluation == 'official':
logging.info('Evaluating using official metric.')
evaluate_official(config)
else:
raise ValueError(f'Unknown evaluation: {args.evaluation}')
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config', default=Path('./release/config.yaml'))
parser.add_argument('--frames-dir', type=Path, required=True)
parser.add_argument('--output-dir', type=Path, required=True)
parser.add_argument('--extensions', nargs='*', default=IMG_EXTENSIONS)
args = parser.parse_args()
with open(args.config, 'r') as f:
config = yaml.load(f)
args.output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output_dir)
logging.debug('Args:\n%s', pprint.pformat(vars(args)))
logging.debug('Config:\n%s', pprint.pformat(config))
msg(f"Computing flow on {args.frames_dir}.")
compute_flow_helper(config,
args.frames_dir,
args.output_dir,
extensions=args.extensions,
recursive=True)
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config', default=Path('./release/config.yaml'))
parser.add_argument('--frames-dir', required=True, type=Path)
parser.add_argument('--detections-dir', required=True, type=Path)
parser.add_argument('--output-dir', required=True, type=Path)
parser.add_argument('--save-numpy', type=parse_bool, default=True)
parser.add_argument('--save-images', type=parse_bool, default=False)
parser.add_argument('--save-video', type=parse_bool, default=True)
parser.add_argument('--fps', default=30, type=float)
parser.add_argument('--model',
choices=['joint', 'appearance', 'motion'],
default='joint')
parser.add_argument(
'--filename-format',
choices=[
'frame', 'frameN', 'sequence_frame', 'sequence-frame', 'fbms'
],
required=True,
help=('Specifies how to get frame number from the filename. '
'"frame": the filename is the frame number, '
'"frameN": format <frame><number>, '
'"sequence_frame": frame number is separated by an underscore, '
'"sequence-frame": frame number is separated by a dash, '
'"fbms": assume fbms style frame numbers'))
args = parser.parse_args()
with open(args.config, 'r') as f:
config = yaml.load(f)
args.output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output_dir)
logging.debug('Args:\n%s', pprint.pformat(vars(args)))
if args.model == 'joint':
detections_args = ['--init-detections-dir', args.detections_dir]
script = 'tracker/two_detector_track.py'
extra_args = ['--remove-continue-overlap', 0.1, '--recursive']
else:
detections_args = ['--detections-dir', args.detections_dir]
script = 'tracker/track_multiple.py'
extra_args = []
args = ['--images-dir', args.frames_dir] + detections_args + [
'--output-dir', args.output_dir, '--save-numpy', args.save_numpy,
'--save-images', args.save_images, '--save-video', args.save_video,
'--bidirectional', '--score-init-min', 0.9, '--fps', args.fps,
'--filename-format', args.filename_format, '--quiet'
] + extra_args
cmd = ['python', script] + args
msg(f'Running tracker')
subprocess_call(cmd)
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config', default=Path('./release/config.yaml'))
args = parser.parse_args()
with open(args.config, 'r') as f:
config = yaml.load(f)
args = parser.parse_args()
output_dir = Path(config['davis16']['output_dir']) / 'detections'
output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, output_dir, args)
detectron_dir = (
Path(__file__).resolve().parent.parent.parent / 'detectron_pytorch')
split_dirs = link_splits(config)
for split in config['davis16']['splits']:
image_dir = split_dirs[split][0]
flow_dir = Path(config['davis16']['output_dir']) / 'flow' / split
joint_config, joint_model = get_config_ckpt(
config['model']['joint']['dir'], config['model']['joint']['step'])
appearance_config, appearance_model = get_config_ckpt(
config['model']['appearance']['dir'],
config['model']['appearance']['step'])
cmd = ['python', 'tools/infer_simple.py']
args = [
'--cfg', joint_config,
'--num_classes', 2,
'--load_ckpt', joint_model,
'--load_appearance_ckpt', appearance_model,
'--set', 'MODEL.MERGE_WITH_APPEARANCE.ENABLED', 'True',
'--image_dirs', image_dir, flow_dir,
'--input_type', 'rgb', 'flow',
'--save_images', False,
'--output_dir', output_dir / split,
'--quiet',
'--recursive'
]
subprocess_call(cmd + args, cwd=str(detectron_dir))
cmd = ['python', 'detectron_to_fgbg_masks.py']
args = [
'--detections-root', output_dir / split,
'--images-dir', image_dir,
'--output-dir', output_dir / split,
'--recursive',
'--extension', '.jpg',
'--threshold', '0.7'
]
subprocess_call(cmd + args)
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('root', type=Path)
parser.add_argument('--split',
required=True,
choices=['train', 'validation'])
args = parser.parse_args()
log_dir = args.root / 'logs'
log_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, log_dir, args)
ann_path = args.root / f'annotations/{args.split}.json'
with open(ann_path, 'r') as f:
tao = json.load(f)
checksums_path = (args.root /
f'annotations/checksums/{args.split}_checksums.json')
with open(checksums_path, 'r') as f:
checksums = json.load(f)
videos_by_dataset = defaultdict(list)
for video in tao['videos']:
videos_by_dataset[video['metadata']['dataset']].append(video)
status = {}
for dataset, videos in sorted(videos_by_dataset.items()):
status[dataset] = True
for video in tqdm(videos, desc=f'Verifying {dataset}'):
name = video['name']
frame_dir = args.root / 'frames' / name
if not are_tao_frames_dumped(
frame_dir, checksums[name], warn=True, allow_extra=False):
logging.warning(
f'Frames for {name} are not extracted properly. '
f'Skipping rest of dataset.')
status[dataset] = False
break
success = []
for dataset in sorted([d for d, v in status.items() if v]):
success.append(f'{dataset: <12}: Verified ✓✓✓')
failure = []
for dataset in sorted([d for d, v in status.items() if not v]):
failure.append(f'{dataset: <12}: FAILED ���')
if success:
logging.info('Success!\n' + ('\n'.join(success)))
if failure:
logging.warning('Some datasets were not properly extracted!\n' +
('\n'.join(failure)))
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--annotations', type=Path, required=True)
# We need the frames dir because the pickles contain boxes into the ordered
# list frames.
parser.add_argument('--frames-dir', type=Path, required=True)
parser.add_argument('--pickle-dir', type=Path, required=True)
parser.add_argument('--oracle-category', action='store_true')
parser.add_argument('--workers', type=int, default=8)
parser.add_argument('--threshold', default=0.5, type=float)
parser.add_argument('--vis-cats', nargs='*', type=str)
parser.add_argument('--videos', nargs='*')
parser.add_argument('--output-dir', type=Path, required=True)
args = parser.parse_args()
args.output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output_dir, args)
paths = list(args.pickle_dir.rglob('*.pkl'))
tao = Tao(args.annotations)
cats = tao.cats.copy()
for cat in cats.values():
if cat['name'] == 'baby':
cat['name'] = 'person'
tasks = []
for p in paths:
video_name = str(p.relative_to(args.pickle_dir)).split('.pkl')[0]
if args.videos is not None and video_name not in args.videos:
continue
output_video = args.output_dir / f'{video_name}.mp4'
if output_video.exists():
continue
tasks.append({
'pickle_path': p,
'video_name': video_name,
'frames_root': args.frames_dir,
'cats': cats,
'vis_cats': args.vis_cats,
'annotations_json': args.annotations,
'threshold': args.threshold,
'output_video': output_video
})
if args.workers == 0:
for task in tqdm(tasks):
visualize(**task)
else:
pool = Pool(args.workers)
list(tqdm(pool.imap_unordered(visualize_star, tasks),
total=len(tasks)))
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--numpy-dir', type=Path, required=True)
parser.add_argument('--output-dir', type=Path, required=True)
parser.add_argument('--fbms-groundtruth', type=Path, required=True)
parser.add_argument('--np-extensions', default=['.npy', '.npz'])
args = parser.parse_args()
args.output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output_dir, args)
inputs = []
for ext in args.np_extensions:
inputs.extend(args.numpy_dir.glob('*' + ext))
all_shot_paths = []
all_track_paths = []
for mask_file in tqdm(inputs):
sequence = mask_file.stem
output_file = args.output_dir / (sequence + '.dat')
groundtruth_dir = args.fbms_groundtruth / sequence / 'GroundTruth'
assert groundtruth_dir.exists(), (
f"Couldn't find groundtruth at {groundtruth_dir}")
groundtruth = fbms_utils.FbmsGroundtruth(groundtruth_dir)
all_shot_paths.append(groundtruth_dir / (sequence + 'Def.dat'))
all_track_paths.append(output_file)
if output_file.exists():
logging.info(f'Output {output_file} exists, skipping.')
continue
mask = np.load(mask_file)
if isinstance(mask, np.lib.npyio.NpzFile):
# Segmentation saved with savez_compressed; ensure there is only
# one item in the dict and retrieve it.
keys = list(mask.keys())
assert len(keys) == 1, (
'Numpy file (%s) contained dict with multiple items, not sure '
'which one to load.' % mask_file)
mask = mask[keys[0]]
mask = {t: mask[t] for t in groundtruth.frame_label_paths.keys()}
fbms_tracks = fbms_utils.masks_to_tracks(mask)
fbms_tracks_str = fbms_utils.get_tracks_text(fbms_tracks,
groundtruth.num_frames)
with open(output_file, 'w') as f:
f.write(fbms_tracks_str)
with open(args.output_dir / 'all_shots.txt', 'w') as f:
f.write(str(len(all_shot_paths)) + '\n')
f.write('\n'.join(str(x.resolve()) for x in all_shot_paths))
with open(args.output_dir / 'all_tracks.txt', 'w') as f:
f.write('\n'.join(str(x.resolve()) for x in all_track_paths))
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--frames-dir', type=Path, required=True)
parser.add_argument('--output-dir', type=Path, required=True)
parser.add_argument('--model',
default='joint',
choices=['joint', 'appearance', 'motion'])
parser.add_argument(
'--filename-format',
choices=[
'frame', 'frameN', 'sequence_frame', 'sequence-frame', 'fbms'
],
required=True,
help=('Specifies how to get frame number from the filename. '
'"frame": the filename is the frame number, '
'"frameN": format "frame<number>", '
'"sequence_frame": frame number is separated by an underscore, '
'"sequence-frame": frame number is separated by a dash, '
'"fbms": assume fbms style frame numbers'))
parser.add_argument('--config', default=Path('./release/config.yaml'))
args = parser.parse_args()
args.output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output_dir, args)
flow_dir = args.output_dir / 'flow'
if args.model != 'appearance':
subprocess_call([
'python', 'release/custom/compute_flow.py', '--frames-dir',
args.frames_dir, '--config', args.config, '--output-dir', flow_dir
])
output_dir = args.output_dir / args.model
detections_dir = output_dir / 'detections'
subprocess_call([
'python', 'release/custom/infer.py', '--frames-dir', args.frames_dir,
'--flow-dir', flow_dir, '--model', args.model, '--config', args.config,
'--output-dir', detections_dir
])
tracks_dir = output_dir / 'tracks'
subprocess_call([
'python', 'release/custom/track.py', '--frames-dir', args.frames_dir,
'--detections-dir', detections_dir, '--filename-format',
args.filename_format, '--config', args.config, '--model', args.model,
'--output-dir', tracks_dir
])
msg(f'Output results to: {tracks_dir}')
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('root', type=Path)
parser.add_argument('--split',
required=True,
choices=['train', 'val', 'test'])
optional = parser.add_argument_group('Optional')
optional.add_argument('--workers', default=8, type=int)
optional.add_argument(
'--movies-dir',
type=Path,
help=('Directory to save AVA movies to. If you have a copy '
'AVA locally, you can point to that directory to skip '
'downloading. NOTE: Any movies downloaded by this script will '
'be deleted after the script completes. Any movies that already '
'existed on disk will not be deleted.'))
args = parser.parse_args()
log_dir = args.root / 'logs'
log_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, log_dir, args)
ann_path = args.root / f'annotations/{args.split}.json'
with open(ann_path, 'r') as f:
tao = json.load(f)
checksums_path = (
args.root / f'annotations/checksums/{args.split}_checksums.json')
with open(checksums_path, 'r') as f:
checksums = json.load(f)
# checksums = {}
# for image in tao['images']:
# video = image['video']
# if video not in checksums:
# checksums[video] = {}
# name = image['file_name'].split('/')[-1].replace('.jpeg', '.jpg')
# checksums[video][name] = ''
download_ava(args.root,
tao,
checksums,
workers=args.workers,
movies_dir=args.movies_dir)
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config', default=Path('./release/config.yaml'))
args = parser.parse_args()
with open(args.config, 'r') as f:
config = yaml.load(f)
output_dir = Path(config['davis17']['output_dir']) / 'evaluation'
output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, output_dir, args)
logging.info('Evaluating using proposed metric.')
evaluate_proposed(config)
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--track-result', required=True, type=Path)
parser.add_argument('--annotations-json',
type=Path,
help='Annotations json')
parser.add_argument('--output-dir', required=True, type=Path)
args = parser.parse_args()
args.output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output_dir, args)
with open(args.annotations_json, 'r') as f:
groundtruth = json.load(f)
create_json(args.track_result, groundtruth, args.output_dir)
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--meta-json',
default=Path('/data/all/youtube-vos/train/meta.json'),
type=Path)
parser.add_argument('--output-dir', type=Path, required=True)
args = parser.parse_args()
args.output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output_dir, args)
with open(args.meta_json, 'r') as f:
data = json.load(f)
ytvos_categories = set([
y['category'] for x in data['videos'].values()
for y in x['objects'].values()
])
logging.info('Total YTVOS categories: %s', len(ytvos_categories))
common_categories = ytvos_categories & COCO_CATEGORIES
ytvos_only = ytvos_categories - COCO_CATEGORIES
logging.info('Common: %s', ', '.join(list(common_categories)))
logging.info('YTVOS only: %s', ', '.join(list(ytvos_only)))
with open(args.output_dir / 'ytvos_not_in_coco_exact.txt', 'w') as f:
f.write('\n'.join(sorted(ytvos_only)))
with open(args.output_dir / 'coco_labels.txt', 'w') as f:
f.write('\n'.join(sorted(COCO_CATEGORIES)))
with open(args.output_dir / 'exact_match_labels.json', 'w') as f:
json.dump({x: x for x in common_categories}, f, indent=True)
near_matches = {}
for category in ytvos_only:
near_matches[category] = sorted(
COCO_CATEGORIES, key=lambda l: levenshtein(category, l))[:5]
with open(args.output_dir / 'closest_matches.json', 'w') as f:
json.dump(near_matches, f, indent=True)
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--annotations', type=Path, required=True)
# We need the frames dir because the pickles contain boxes into the ordered
# list frames.
parser.add_argument('--frames-dir', type=Path, required=True)
parser.add_argument('--pickle-dir', type=Path, nargs='+', required=True)
parser.add_argument('--oracle-category', action='store_true')
parser.add_argument('--output-dir', type=Path, required=True)
args = parser.parse_args()
args.output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output_dir, args)
tao = Tao(args.annotations)
create_json(args.pickle_dir, tao, args.frames_dir, args.output_dir,
args.oracle_category)
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--frames-dir', type=Path, required=True)
parser.add_argument('--output-json', type=Path, required=True)
parser.add_argument('--tao-annotations', type=Path, required=True)
args = parser.parse_args()
output_dir = args.output_json.parent
output_dir.mkdir(exist_ok=True, parents=True)
common_setup(args.output_json.name, output_dir, args)
with open(args.tao_annotations, 'r') as f:
tao = json.load(f)
videos = [x['name'] for x in tao['videos']]
labeled_frames = defaultdict(set)
for frame in tao['images']:
video, frame_name = frame['file_name'].rsplit('/', 1)
labeled_frames[video].add(frame_name)
# videos = videos[:10]
hashes = {}
for video in tqdm(videos):
frames = fs.glob_ext(args.frames_dir / video, ('.jpg', '.jpeg'))
hashes[video] = {}
for i, frame in tqdm(enumerate(frames)):
if frame.name in labeled_frames[video]:
with open(frame, 'rb') as f:
hashes[video][frame.name] = md5(f.read()).hexdigest()
else:
hashes[video][frame.name] = ''
if all(x == '' for x in hashes[video].values()):
raise ValueError(f'Did not find any labeled frames for {video}')
with open(args.output_json, 'w') as f:
json.dump(hashes, f)
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('annotations', type=Path)
parser.add_argument('predictions', type=Path)
parser.add_argument('--output-dir', type=Path)
parser.add_argument('--config',
type=Path,
default=CONFIG_DIR / 'default.yaml')
parser.add_argument('--config-updates', nargs='*')
args = parser.parse_args()
if args.output_dir:
tensorboard_dir = args.output_dir / 'tensorboard'
if tensorboard_dir.exists():
raise ValueError(
f'Tensorboard dir already exists, not evaluating.')
args.output_dir.mkdir(exist_ok=True, parents=True)
log_path = common_setup(__file__, args.output_dir, args).name
else:
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(format='%(asctime)s.%(msecs).03d: %(message)s',
datefmt='%H:%M:%S')
logging.info('Args:\n%s', vars(args))
log_path = None
cfg = get_cfg_defaults()
merge_from_file_with_base(cfg, args.config)
if args.config_updates:
cfg.merge_from_list(args.config_updates)
cfg.freeze()
if args.output_dir:
with open(args.output_dir / 'config.yaml', 'w') as f:
f.write(cfg.dump())
tao_eval = evaluate(args.annotations, args.predictions, cfg)
log_eval(tao_eval, cfg, output_dir=args.output_dir, log_path=log_path)
help="path to config file")
parser.add_argument('--gpus', default=[0], nargs='+', type=int)
parser.add_argument(
"--opts",
help="Modify model config options using the command-line",
default=[],
nargs=argparse.REMAINDER)
parser.add_argument(
'--save-masks', default=False, action='store_true')
return parser
if __name__ == "__main__":
args = get_parser().parse_args()
Path(args.output).mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output, args)
# Prevent detectron from flooding terminal with messages.
logging.getLogger('detectron2.checkpoint.c2_model_loading').setLevel(
logging.WARNING)
logging.getLogger('fvcore.common.checkpoint').setLevel(
logging.WARNING)
logger = logging.root
cfg = setup_cfg(args)
threads_per_worker = 4
torch.set_num_threads(threads_per_worker)
os.environ['OMP_NUM_THREADS'] = str(threads_per_worker)
all_files = args.root.rglob('*.jpg')
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
'--ytvos-matched',
help='JSON mapping some YTVOS labels to COCO labels.',
default=Path(
'/data/achald/track/ytvos/compare-labels-with-coco/all_matched_labels.json'
),
type=Path)
parser.add_argument(
'--annotations-dir',
required=True,
type=Path,
help='Contains directory with annotations for each video')
parser.add_argument('--meta-json',
help='YTVOS train meta.json',
default=Path('/data/all/youtube-vos/train/meta.json'),
type=Path)
parser.add_argument('--output-dir', type=Path, required=True)
args = parser.parse_args()
args.output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output_dir, args)
with open(args.meta_json, 'r') as f:
all_labels = json.load(f)
labels = {
video: {x['category']
for x in info['objects'].values()}
for video, info in all_labels['videos'].items()
}
with open(args.ytvos_matched, 'r') as f:
known_labels = set(json.load(f).keys())
unknown_only = []
known_only = []
mixed = []
for video_dir in args.annotations_dir.iterdir():
if not video_dir.is_dir():
continue
video = video_dir.name
video_labels = labels[video]
if video_labels.isdisjoint(known_labels):
unknown_only.append(video)
elif video_labels.issubset(known_labels):
known_only.append(video)
else:
mixed.append(video)
logging.info('New category videos: %s', len(unknown_only))
logging.info('Known category videos: %s', len(known_only))
logging.info('Mixed new/known category videos: %s', len(mixed))
unknown_dir = args.output_dir / 'unknown'
unknown_dir.mkdir(exist_ok=True)
for video in unknown_only:
destination = unknown_dir / video
source = args.annotations_dir / video
destination.symlink_to(source.resolve())
known_dir = args.output_dir / 'known'
known_dir.mkdir(exist_ok=True)
for video in known_only:
destination = known_dir / video
source = args.annotations_dir / video
destination.symlink_to(source.resolve())
mixed_dir = args.output_dir / 'mixed'
mixed_dir.mkdir(exist_ok=True)
for video in mixed:
destination = mixed_dir / video
source = args.annotations_dir / video
destination.symlink_to(source.resolve())
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--annotations', type=Path, required=True)
parser.add_argument('--frames-dir', type=Path, required=True)
parser.add_argument('--init',
choices=['first', 'biggest'],
default='first')
parser.add_argument('--output-dir', type=Path, required=True)
parser.add_argument('--tracker',
choices=['pysot', 'pytrack', 'staple', 'srdcf'],
default='pysot')
pysot_args = parser.add_argument_group('pysot_params')
pysot_args.add_argument('--pysot-config-file', '--config-file', type=Path)
pysot_args.add_argument('--pysot-model-path', '--model-path', type=Path)
pytracking_args = parser.add_argument_group('pytracking_params')
pytracking_args.add_argument('--pytrack-name')
pytracking_args.add_argument('--pytrack-param')
pytracking_args.add_argument(
'--pytrack-model-path',
help=('Specify path to model, if different from the one implied by '
'--pytrack-param.'))
parser.add_argument('--gpus', default=[0, 1, 2, 3], nargs='*', type=int)
parser.add_argument('--tasks-per-gpu', default=1, type=int)
parser.add_argument('--visualize', default=False, type=misc.parse_bool)
args = parser.parse_args()
args.output_dir.mkdir(exist_ok=True, parents=True)
if args.init == 'first':
common_setup(__file__, args.output_dir, args)
else:
common_setup(f'{Path(__file__).stem}_{args.init}', args.output_dir,
args)
_num_threads = 4
torch.set_num_threads(_num_threads)
os.environ['OMP_NUM_THREADS'] = str(_num_threads)
if args.tracker == 'pysot':
assert args.pysot_config_file is not None
assert args.pysot_model_path is not None
elif args.tracker == 'pytrack':
assert args.pytrack_name is not None
assert args.pytrack_param is not None
elif args.tracker in ('staple', 'srdcf'):
pass
tao = Tao(args.annotations)
video_tracks = defaultdict(list)
for track_id, track in tao.tracks.items():
video_tracks[track['video_id']].append(track)
# List of kwargs passed to track_video().
track_video_tasks = []
pickle_output_dir = args.output_dir / "pickles"
pickle_output_dir.mkdir(exist_ok=True, parents=True)
for video_id, tracks in tqdm(video_tracks.items(),
desc='Collecting tasks'):
video_name = tao.vids[video_id]['name']
frames_dir = args.frames_dir / video_name
output = (pickle_output_dir / video_name).with_suffix('.pkl')
if output.exists():
logging.info(f'{output} already exists, skipping.')
continue
# Map track id to
# {'frame': name, 'init': [x0, y0, w, h]}
frames = natsorted(fs.glob_ext(frames_dir, fs.IMG_EXTENSIONS))
if not frames[0].exists():
# Just check the first frame for efficiency; usually, either all
# frames will be missing, or all will be available.
logging.info(f'Frame link {frames[0]} broken for {video_name} in '
f'{frames_dir}, skipping.')
continue
objects = {}
for track in tracks:
annotation = tao.get_single_object_init(track['id'], args.init)
frame_name = tao.imgs[annotation['image_id']]['file_name']
frame_indices = {
str(x.relative_to(args.frames_dir)): i
for i, x in enumerate(frames)
}
init_frame_index = frame_indices[frame_name]
if args.init == 'first':
first_frame_index = init_frame_index
else:
first_ann = tao.get_kth_annotation(track['id'], 0)
first_frame_name = tao.imgs[first_ann['image_id']]['file_name']
first_frame_index = frame_indices[first_frame_name]
objects[track['id']] = {
'first_annotated_frame': first_frame_index,
'sot_init_frame': init_frame_index,
'init': annotation['bbox'],
}
task = {
'objects': objects,
'output_pickle': output,
'frames_dir': frames_dir,
'visualize': args.visualize
}
track_video_tasks.append(task)
gpus = args.gpus * args.tasks_per_gpu
if args.tracker == 'pysot':
tracker_init = {
'config_file': args.pysot_config_file,
'model_path': args.pysot_model_path,
}
elif args.tracker == 'pytrack':
tracker_init = {
'tracker_name': args.pytrack_name,
'tracker_param': args.pytrack_param,
'model_path': args.pytrack_model_path
}
elif args.tracker in ('staple', 'srdcf'):
tracker_init = {}
if track_video_tasks:
gpus = gpus[:len(track_video_tasks)]
print(gpus)
if len(gpus) == 1:
context = {'gpu': gpus[0]}
init_tracker(
{
'tracker_init': tracker_init,
'tracker_type': args.tracker
}, context)
for task in tqdm(track_video_tasks):
task['show_progress'] = True
track_video_helper(task, context)
else:
pool = FixedGpuPool(gpus,
initializer=init_tracker,
initargs={
'tracker_init': tracker_init,
'tracker_type': args.tracker
})
list(
tqdm(pool.imap_unordered(track_video_helper,
track_video_tasks),
total=len(track_video_tasks)))
else:
logging.warning('No tasks found!')
create_json(pickle_output_dir,
tao,
args.frames_dir,
args.output_dir,
oracle_category=True)
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--lvis', type=Path, required=True)
parser.add_argument('--coco', type=Path, required=True)
parser.add_argument('--mapping',
type=Path,
default=ROOT / 'data/lvis_coco_to_synset.json')
parser.add_argument('--output-json',
type=Path,
required=True)
parser.add_argument(
'--iou-thresh',
default=0.7,
type=float,
help=('If a COCO annotation overlaps with an LVIS annotations with '
'IoU over this threshold, we use only the LVIS annotation.'))
args = parser.parse_args()
args.output_json.parent.mkdir(exist_ok=True, parents=True)
common_setup(args.output_json.name + '.log', args.output_json.parent, args)
coco = COCO(args.coco)
lvis = COCO(args.lvis)
synset_to_lvis_id = {x['synset']: x['id'] for x in lvis.cats.values()}
coco_to_lvis_category = {}
with open(args.mapping, 'r') as f:
name_mapping = json.load(f)
for category in coco.cats.values():
mapped = name_mapping[category['name']]
assert mapped['coco_cat_id'] == category['id']
synset = mapped['synset']
if synset not in synset_to_lvis_id:
logging.debug(
f'Found no LVIS category for "{category["name"]}" from COCO')
continue
coco_to_lvis_category[category['id']] = synset_to_lvis_id[synset]
for image_id, image in coco.imgs.items():
if image_id in lvis.imgs:
coco_name = coco.imgs[image_id]['file_name']
lvis_name = lvis.imgs[image_id]['file_name']
assert coco_name in lvis_name
else:
logging.info(
f'Image {image_id} in COCO, but not annotated in LVIS')
lvis_highest_id = max(x['id'] for x in lvis.anns.values())
ann_id_generator = itertools.count(lvis_highest_id + 1)
new_annotations = []
for image_id, lvis_anns in tqdm(lvis.imgToAnns.items()):
if image_id not in coco.imgToAnns:
logging.info(
f'Image {image_id} in LVIS, but not annotated in COCO')
continue
coco_anns = coco.imgToAnns[image_id]
# Compute IoU between coco_anns and lvis_anns
# Shape (num_coco_anns, num_lvis_anns)
mask_iou = mask_util.iou([coco.annToRLE(x) for x in coco_anns],
[lvis.annToRLE(x) for x in lvis_anns],
pyiscrowd=np.zeros(len(lvis_anns)))
does_overlap = mask_iou.max(axis=1) > args.iou_thresh
to_add = []
for i, ann in enumerate(coco_anns):
if does_overlap[i]:
continue
if ann['category_id'] not in coco_to_lvis_category:
continue
ann['category_id'] = coco_to_lvis_category[ann['category_id']]
ann['id'] = next(ann_id_generator)
to_add.append(ann)
new_annotations.extend(to_add)
with open(args.lvis, 'r') as f:
merged = json.load(f)
merged['annotations'].extend(new_annotations)
with open(args.output_json, 'w') as f:
json.dump(merged, f)
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
'--input-dir',
required=True,
help='Directory containing a subdir for every sequence.')
parser.add_argument('--output-dir', required=True)
parser.add_argument('--flow-type',
choices=['flownet2', 'liteflownet'],
required=True)
parser.add_argument(
'--recursive',
action='store_true',
help="""Search recursively in input-dir for sequences. Any directory
containing a file with extension specified by --extensions is
treated as a sequence directory. NOTE: Does not support
symlinked directories.""")
parser.add_argument('--extensions', nargs='*', default=IMG_EXTENSIONS)
parser.add_argument(
'--convert-to-angle-magnitude-png',
choices=['off', 'on', 'pavel-matlab'],
default='off',
help=('Convert flo files to angle/magnitude PNGs, and do not keep '
'.flo files around.'))
parser.add_argument('--gpus', default=[0, 1, 2, 3], nargs='*', type=int)
parser.add_argument(
'--num-workers',
default=-1,
type=int,
help=('Number of workers. By default, set to the number of GPUs. '
'Having more workers than GPUs allows some workers to process '
'CPU operations, like loading input/output lists, checking '
'image dimensions, and converting .flo to .png while other '
'workers use the GPU.'))
parser.add_argument('--quiet', action='store_true')
flownet2_parser = parser.add_argument_group('Flownet2 params')
flownet2_parser.add_argument('--flownet2-dir',
help='Path to flownet2 repo.')
flownet2_parser.add_argument('--flownet2-model',
default='kitti',
choices=['kitti', 'sintel', 'chairs-things'])
liteflownet_parser = parser.add_argument_group('Liteflownet Params')
liteflownet_parser.add_argument('--liteflownet-dir',
help='Path to liteflownet repo')
liteflownet_parser.add_argument(
'--liteflownet-model',
default='liteflownet-ft-kitti',
help=('Model to use for evaluation. chairs-things maps to the '
'`liteflownet` model, `sintel` maps to `liteflownet-ft-sintel` '
'and `kitti` maps to `liteflownet-ft-kitti`.'),
choices=['chairs-things', 'sintel', 'kitti'])
args = parser.parse_args()
input_root = Path(args.input_dir)
output_root = Path(args.output_dir)
output_root.mkdir(parents=True, exist_ok=True)
file_name = Path(__file__).stem
log_level = logging.ERROR if args.quiet else logging.INFO
file_logger = common_setup(__file__,
output_root,
args,
log_console_level=log_level)
if args.recursive:
sequences = sorted({
x.parent
for x in glob_ext(input_root, args.extensions, recursive=True)
})
else:
sequences = sorted(input_root.iterdir())
import multiprocessing as mp
manager = mp.Manager()
gpu_queue = manager.Queue()
if args.num_workers == -1:
args.num_workers = len(args.gpus)
pool = mp.Pool(args.num_workers)
for gpu in args.gpus:
gpu_queue.put(gpu)
if args.flow_type == 'flownet2':
flownet2_root = Path(args.flownet2_dir)
assert args.flownet2_model is not None
assert flownet2_root.exists()
flow_fn = compute_flownet2_flow
flow_args = {
'flownet_root': flownet2_root,
'cnn_model': args.flownet2_model,
'tmp_prefix': file_name
}
else:
liteflownet_root = Path(args.liteflownet_dir)
assert args.liteflownet_model is not None
assert liteflownet_root.exists()
flow_fn = compute_liteflownet_flow
flow_args = {
'liteflownet_root': liteflownet_root,
'cnn_model': args.liteflownet_model,
'tmp_prefix': file_name
}
tasks = []
convert_png = {
'off': False,
'on': 'default',
'pavel-matlab': 'pavel-matlab'
}[args.convert_to_angle_magnitude_png]
for sequence_path in sequences:
output_dir = output_root / (sequence_path.relative_to(input_root))
tasks.append({
'input_dir': sequence_path,
'output_dir': output_dir,
'flow_fn': flow_fn,
'flow_args': flow_args,
'gpu_queue': gpu_queue,
'logger_name': file_logger.name,
'convert_png': convert_png,
'remove_flo': bool(convert_png),
'extensions': args.extensions
})
list(
tqdm(pool.imap_unordered(compute_sequence_flow_gpu_helper, tasks),
total=len(tasks)))
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
'--rendered-masks',
type=Path,
default=Path('/data/achald/track/FBMS/khan-continuum_cvpr17/outputs/'))
parser.add_argument(
'--fbms-root',
help='Should contain TrainingSet and TestSet dirs.',
type=Path,
default=Path('/data/all/FBMS'))
parser.add_argument('--output-dir', type=Path, required=True)
args = parser.parse_args()
args.output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output_dir, args)
test_dir = args.fbms_root / 'TestSet'
train_dir = args.fbms_root / 'TrainingSet'
output_vis = args.output_dir / 'vis'
output_mask = args.output_dir / 'masks'
output_vis.mkdir(exist_ok=True)
output_mask.mkdir(exist_ok=True)
for split_dir in [train_dir, test_dir]:
for subdir in tqdm(list(split_dir.iterdir())):
if not subdir.is_dir():
continue
output_seq_vis = output_vis / subdir.stem
output_seq_vis.mkdir(exist_ok=True)
sequence = subdir.name
mask_dir = args.rendered_masks / sequence
frames = sorted(subdir.glob('*.jpg'),
key=lambda x: get_framenumber(x))
output_numpy = output_mask / (subdir.stem + '.npz')
# if output_numpy.exists():
# logging.info(f'{output_numpy} already exists, skipping')
# continue
segmentation = []
image_h, image_w = cv2.imread(str(frames[0])).shape[:2]
diffs = []
for i, frame in enumerate(frames):
if i == 0:
continue
frame_no = get_framenumber(frame) - 1
segmentation_path = (
mask_dir / f'segmentation{frame_no:04d}.png')
if not segmentation_path.exists():
diffs.append(None)
logging.error(
f'Missing segmentation at {segmentation_path}')
continue
diff = compute_diff(frame, segmentation_path)
cv2.imwrite(str(output_seq_vis / f'diff{i:04d}.png'),
diff)
diffs.append(diff)
try:
valid_diff = next(x for x in diffs if x is not None)
except StopIteration:
logging.error(f'No valid segmentation found in {subdir}')
continue
for i, x in enumerate(diffs):
if x is None:
diffs[i] = np.zeros_like(valid_diff)
quantized, segmentation = quantize_colors(np.stack(diffs))
resized_segmentation = []
for i in range(quantized.shape[0]):
resized_segmentation.append(
cv2.resize(segmentation[i], (image_w, image_h),
interpolation=cv2.INTER_NEAREST))
cv2.imwrite(str(output_seq_vis / f'mask{i+1:04d}.png'),
quantized[i])
np.savez_compressed(output_numpy, np.stack(resized_segmentation))
def main():
"""Parse in command line arguments"""
parser = argparse.ArgumentParser(
description='Demonstrate mask-rcnn results')
parser.add_argument(
'--config-file',
required=True,
help='optional config file')
parser.add_argument(
'--model-path',
type=Path,
help=('Path to model pickle file. If not specified, the latest '
'checkpoint, if it exists, or cfg.MODEL.WEIGHT is loaded.'))
parser.add_argument('--gpus', default=[0], nargs='*', type=int)
parser.add_argument('--score-thresh', default=0.7, type=float)
parser.add_argument(
'--image-dir',
type=Path,
help='directory to load images for demo')
parser.add_argument(
'--images', nargs='+',
type=Path,
help=('images to infer. Must not use with --image_dirs. If the model '
'requires multiple input datasets, use --image_dirs instead.'))
parser.add_argument(
'--output-dir',
type=Path,
help='directory to save demo results',
default="infer_outputs")
parser.add_argument(
'--recursive',
help='Whether to search recursively in --image-dir for images.',
action='store_true')
parser.add_argument('--visualize', action='store_true')
parser.add_argument(
'--threads-per-worker',
type=int,
help='Defaults to 1 if using multiple gpus, otherwise unrestricted.')
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER)
args = parser.parse_args()
args.output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output_dir, args)
if len(args.gpus) > 1 and args.threads_per_worker is None:
args.threads_per_worker = 1
if args.threads_per_worker is not None:
torch.set_num_threads(args.threads_per_worker)
os.environ['OMP_NUM_THREADS'] = str(args.threads_per_worker)
# update the config options with the config file
cfg.merge_from_file(args.config_file)
if len(args.gpus) > 1 and 'DATALOADER.NUM_WORKERS' not in args.opts:
args.opts += ['DATALOADER.NUM_WORKERS', '1']
cfg.merge_from_list(args.opts)
cfg.freeze()
if args.image_dir:
logging.info('Collecting images')
images = [x for x in args.image_dir.iterdir() if is_image(x)]
if args.recursive:
# We could just use args.images_dir.rglob('*'), but this ignores
# symlinks. For convenience, we handle immediate child directories
# being symlinks.
for x in args.image_dir.iterdir():
if x.is_dir():
images.extend([y for y in x.rglob('*') if is_image(y)])
outputs = [(args.output_dir / x.relative_to(
args.image_dir)).with_suffix('.pickle') for x in images]
else:
images = args.images
outputs = [(args.output_dir / (x.stem + '.pickle')) for x in images]
if not images:
raise ValueError('No images found!')
logging.info('Inferring on %s images', len(images))
init_args = {
'model_path': str(args.model_path.resolve()),
'config': cfg,
'score_threshold': args.score_thresh
}
infer_tasks = [{
'image_path': image,
'output_path': output,
'visualize': args.visualize
} for image, output in zip(images, outputs)]
if len(args.gpus) == 1:
context = {'gpu': args.gpus[0]}
init_model(init_args, context)
for task in tqdm(infer_tasks):
infer(task, context)
else:
pool = FixedGpuPool(
args.gpus, initializer=init_model, initargs=init_args)
list(
tqdm(
pool.imap_unordered(infer, infer_tasks),
total=len(infer_tasks)))
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config', default=Path('./release/config.yaml'))
parser.add_argument('--frames-dir',
type=Path,
help='Required unless --model set to "motion"')
parser.add_argument('--flow-dir',
type=Path,
help='Required unless --model set to "appearance".')
parser.add_argument('--output-dir', type=Path, required=True)
parser.add_argument('--visualize', action='store_true')
parser.add_argument('--model',
choices=['joint', 'appearance', 'motion'],
default='joint')
args = parser.parse_args()
if args.model != 'appearance':
assert args.flow_dir and args.flow_dir.exists(), (
f'--flow-dir must be specified for --model {args.model}')
if args.model != 'motion':
assert args.frames_dir and args.frames_dir.exists(), (
f'--frames-dir must be specified for --model {args.model}')
with open(args.config, 'r') as f:
config = yaml.load(f)
args = parser.parse_args()
args.output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output_dir)
logging.debug('Args:\n%s', pprint.pformat(vars(args)))
detectron_dir = (Path(__file__).resolve().parent.parent.parent /
'detectron_pytorch')
model_config, model = get_config_ckpt(config['model'][args.model]['dir'],
config['model'][args.model]['step'])
maybe_combine_appearance = []
if args.model == 'joint':
appearance_config, appearance_model = get_config_ckpt(
config['model']['appearance']['dir'],
config['model']['appearance']['step'])
maybe_combine_appearance = [
'--load_appearance_ckpt', appearance_model, '--set',
'MODEL.MERGE_WITH_APPEARANCE.ENABLED', 'True'
]
image_dirs = ['--image_dirs', args.frames_dir, args.flow_dir]
input_type = ['--input_type', 'rgb', 'flow']
elif args.model == 'appearance':
maybe_combine_appearance = []
image_dirs = ['--image_dirs', args.frames_dir]
input_type = ['--input_type', 'rgb']
elif args.model == 'motion':
maybe_combine_appearance = []
image_dirs = ['--image_dirs', args.flow_dir]
input_type = ['--input_type', 'flow']
msg(f'Running {args.model} model on {image_dirs[1].resolve()}')
cmd = ['python', 'tools/infer_simple.py']
args = (['--cfg', model_config, '--num_classes', 2, '--load_ckpt', model] +
maybe_combine_appearance + image_dirs + input_type + [
'--save_images', args.visualize, '--output_dir',
args.output_dir, '--quiet', '--recursive'
])
subprocess_call(cmd + args, cwd=str(detectron_dir))
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config', default=Path('./release/config.yaml'))
parser.add_argument(
'--without-ytvos-train',
action='store_true',
help=("By default, we infer with model that was trained on YTVOS. "
"For fair evaluation of generalization, we use a model without "
"YTVOS training in our manuscript. Set this to True to use "
"the model without YTVOS training."))
args = parser.parse_args()
with open(args.config, 'r') as f:
config = yaml.load(f)
args = parser.parse_args()
output_dir = Path(config['ytvos']['output_dir']) / 'detections'
if args.without_ytvos_train:
output_dir = output_dir / 'without_ytvos'
else:
output_dir = output_dir / 'with_ytvos'
output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, output_dir, args)
detectron_dir = (
Path(__file__).resolve().parent.parent.parent / 'detectron_pytorch')
split_dirs = link_splits(config)
for split in config['ytvos']['splits']:
image_dir = split_dirs[split][0]
flow_dir = Path(config['ytvos']['output_dir']) / 'flow' / split
if args.with_ytvos_train:
joint_config, joint_model = get_config_ckpt(
config['model']['joint']['dir'],
config['model']['joint']['step'])
else:
joint_config, joint_model = get_config_ckpt(
config['model']['joint_no_ytvos']['dir'],
config['model']['joint_no_ytvos']['step'])
appearance_config, appearance_model = get_config_ckpt(
config['model']['appearance']['dir'],
config['model']['appearance']['step'])
cmd = ['python', 'tools/infer_simple.py']
args = [
'--cfg', joint_config,
'--num_classes', 2,
'--load_ckpt', joint_model,
'--load_appearance_ckpt', appearance_model,
'--set', 'MODEL.MERGE_WITH_APPEARANCE.ENABLED', 'True',
'--image_dirs', image_dir, flow_dir,
'--input_type', 'rgb', 'flow',
'--save_images', False,
'--output_dir', output_dir / split,
'--quiet',
'--recursive'
]
subprocess_call(cmd + args, cwd=str(detectron_dir))
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('root', type=Path)
parser.add_argument('--split',
required=True,
choices=['train', 'val', 'test'])
parser.add_argument('--sources',
default=['BDD', 'HACS', 'Charades', 'YFCC100M'],
choices=['BDD', 'HACS', 'Charades', 'YFCC100M'])
parser.add_argument('--workers', default=8, type=int)
args = parser.parse_args()
log_dir = args.root / 'logs'
log_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, log_dir, args)
ann_path = args.root / f'annotations/{args.split}.json'
with open(ann_path, 'r') as f:
tao = json.load(f)
checksums_path = (args.root /
f'annotations/checksums/{args.split}_checksums.json')
with open(checksums_path, 'r') as f:
checksums = json.load(f)
videos_by_dataset = defaultdict(list)
for video in tao['videos']:
videos_by_dataset[video['metadata']['dataset']].append(video)
videos_dir = args.root / 'videos'
frames_dir = args.root / 'frames'
for dataset in args.sources:
# Collect list of videos
ext = '.mov' if dataset == 'BDD' else '.mp4'
videos = videos_by_dataset[dataset]
video_paths = [
videos_dir / f"{video['name']}{ext}" for video in videos
]
output_frame_dirs = [frames_dir / video['name'] for video in videos]
# List of (video, video path, frame directory) tuples
to_dump = []
for video, video_path, frame_dir in zip(videos, video_paths,
output_frame_dirs):
if not video_path.exists():
raise ValueError(f'Could not find video at {video_path}')
video_checksums = checksums[video['name']]
if frame_dir.exists() and are_tao_frames_dumped(
frame_dir, video_checksums, warn=False):
continue
to_dump.append((video, video_path, frame_dir))
# Dump frames from each video
logging.info(f'{dataset}: Extracting frames')
dump_tao_frames([x[1] for x in to_dump], [x[2] for x in to_dump],
workers=args.workers)
to_dump = []
for video, video_path, frame_dir in zip(videos, video_paths,
output_frame_dirs):
video_checksums = checksums[video['name']]
# Remove frames not used for TAO.
remove_non_tao_frames(frame_dir, set(video_checksums.keys()))
# Compare checksums for frames
assert are_tao_frames_dumped(frame_dir, video_checksums), (
f'Not all TAO frames for {video["name"]} were extracted.')
logging.info(
f'{dataset}: Removing non-TAO frames, verifying extraction')
logging.info(f'{dataset}: Successfully extracted!')
def main():
# Use first line of file docstring as description if it exists.
parser = argparse.ArgumentParser(
description=__doc__.split('\n')[0] if __doc__ else '',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--detections-dir',
type=Path,
required=True,
help='Results directory with pickle or mat files')
parser.add_argument('--annotations',
type=Path,
required=True,
help='Annotations json')
parser.add_argument(
'--output-dir',
type=Path,
required=True,
help=('Output directory, where a results.json will be output, as well '
'as a .npz file for each video, containing a boxes array of '
'size (num_boxes, 6), of the format [x0, y0, x1, y1, class, '
'score, box_index, track_id], where box_index maps into the '
'pickle files'))
parser.add_argument('--max-age', default=100, type=int)
parser.add_argument('--min-hits', default=1, type=float)
parser.add_argument('--min-iou', default=0.1, type=float)
parser.add_argument('--score-threshold',
default=0.0005,
help='Float or "none".')
parser.add_argument('--nms-thresh', type=float, default=-1)
parser.add_argument('--workers', default=8, type=int)
args = parser.parse_args()
args.score_threshold = (-float('inf') if args.score_threshold == 'none'
else float(args.score_threshold))
args.output_dir.mkdir(exist_ok=True, parents=True)
common_setup(__file__, args.output_dir, args)
npz_dir = dir_path(args.output_dir / 'npz_files')
def get_output_path(video):
return npz_dir / (video + '.npz')
with open(args.annotations, 'r') as f:
groundtruth = json.load(f)
videos = [x['name'] for x in groundtruth['videos']]
video_paths = {}
for video in tqdm(videos, desc='Collecting paths'):
output = get_output_path(video)
if output.exists():
logging.debug(f'{output} already exists, skipping...')
continue
vid_detections = args.detections_dir / video
assert vid_detections.exists(), (
f'No detections dir at {vid_detections}!')
detection_paths = natsorted(
(args.detections_dir / video).rglob(f'*.pkl'))
assert detection_paths, (f'No detections pickles at {vid_detections}!')
video_paths[video] = detection_paths
if not video_paths:
logging.info(f'Nothing to do! Exiting.')
return
logging.info(f'Found {len(video_paths)} videos to track.')
tasks = []
for video, paths in tqdm(video_paths.items()):
output = get_output_path(video)
tasks.append((paths, output, args.score_threshold, args.nms_thresh, {
'iou_threshold': args.min_iou,
'min_hits': args.min_hits,
'max_age': args.max_age
}))
if args.workers > 0:
pool = Pool(args.workers)
list(
tqdm(pool.imap_unordered(track_and_save_star, tasks),
total=len(tasks),
desc='Tracking'))
else:
for task in tqdm(tasks):
track_and_save(*task)
logging.info(f'Finished')
create_json(npz_dir, groundtruth, args.output_dir)
