From fbf52e00b59afb5cec4b4e01b38ab7f323505d64 Mon Sep 17 00:00:00 2001 From: Melokeo Date: Fri, 31 Oct 2025 14:00:16 -0400 Subject: [PATCH] Add start and end params for video creation Add `--start` and `--end` params for `label-3d` and `label-combined`, so that one can make shorter videos for quick quality check. --- anipose/anipose.py | 16 ++++++++-- anipose/common.py | 64 ++++++++++++++++++++++++++++++++++++++ anipose/label_combined.py | 23 +++++++++++--- anipose/label_videos_3d.py | 13 ++++++-- 4 files changed, 106 insertions(+), 10 deletions(-) diff --git a/anipose/anipose.py b/anipose/anipose.py index 1a02ccb..62983c1 100644 --- a/anipose/anipose.py +++ b/anipose/anipose.py @@ -258,9 +258,15 @@ def label_2d_filter(config): label_videos_filtered_all(config) @cli.command() +@click.option('--start', default=None, type=str, + help='Start time: timestamp (00:00:05) or fraction (0.1)') +@click.option('--end', default=None, type=str, + help='End time: timestamp (00:01:30) or fraction (0.9)') @pass_config -def label_3d(config): +def label_3d(config, start, end): from .label_videos_3d import label_videos_3d_all + from .common import parse_range + config['export_range'] = parse_range(start, end) click.echo('Labeling videos in 3D...') label_videos_3d_all(config) @@ -272,9 +278,15 @@ def label_3d_filter(config): label_videos_3d_filtered_all(config) @cli.command() +@click.option('--start', default=None, type=str, + help='Start time: timestamp (00:00:05) or fraction (0.1)') +@click.option('--end', default=None, type=str, + help='End time: timestamp (00:01:30) or fraction (0.9)') @pass_config -def label_combined(config): +def label_combined(config, start, end): from .label_combined import label_combined_all + from .common import parse_range + config['export_range'] = parse_range(start, end) click.echo('Labeling combined videos...') label_combined_all(config) diff --git a/anipose/common.py b/anipose/common.py index 079b65f..2a4c725 100644 --- a/anipose/common.py +++ b/anipose/common.py @@ -226,3 +226,67 @@ def get_calibration_board_image(config): size = numx*200, numy*200 img = board.draw(size) return img + +# range parsing utils for start-end param in video labeling +def parse_range(start: str | None, end: str | None) -> dict | None: + """parse timestamp (HH:MM:SS) or fraction (0.0-1.0)""" + if start is None and end is None: + return None + + def parse_value(val: str | None) -> float | str | None: + if val is None: + return None + # check if is fraction + try: + f = float(val) + if 0 <= f <= 1: + return f + except ValueError: + pass + # assume timestamp format (unsafe) + return val # absolute, convert later with fps + + return { + 'start': parse_value(start), + 'end': parse_value(end), + 'mode': 'relative' if isinstance(parse_value(start or end), float) else 'absolute' + } + +def get_frame_range(config: dict, total_frames: int, fps: float) -> tuple[int, int]: + """convert export_range (absolute or relative) to (start_frame, end_frame)""" + range_cfg = config.get('export_range') + if range_cfg is None: + return 0, total_frames + + start = range_cfg.get('start') + end = range_cfg.get('end') + + # relative -> absolute + if isinstance(start, float): + start_frame = int(start * total_frames) + elif isinstance(start, str): + start_frame = timestamp_to_frame(start, fps) + else: + start_frame = 0 + + if isinstance(end, float): + end_frame = int(end * total_frames) + elif isinstance(end, str): + end_frame = timestamp_to_frame(end, fps) + else: + end_frame = total_frames + + return start_frame, end_frame + +def timestamp_to_frame(timestamp: str, fps: float) -> int: + """convert HH:MM:SS or MM:SS to frame#""" + parts = list(map(float, timestamp.split(':'))) + if len(parts) == 3: + h, m, s = parts + total_sec = h * 3600 + m * 60 + s + elif len(parts) == 2: + m, s = parts + total_sec = m * 60 + s + else: + total_sec = parts[0] + return int(total_sec * fps) \ No newline at end of file diff --git a/anipose/label_combined.py b/anipose/label_combined.py index 7029a17..52e5234 100644 --- a/anipose/label_combined.py +++ b/anipose/label_combined.py @@ -13,10 +13,12 @@ from aniposelib.cameras import CameraGroup -from .common import make_process_fun, get_nframes, \ - get_video_name, get_cam_name, \ - get_video_params, get_video_params_cap, \ - get_data_length, natural_keys, true_basename, find_calibration_folder +from .common import (make_process_fun, get_nframes, + get_video_name, get_cam_name, + get_video_params, get_video_params_cap, + get_data_length, natural_keys, true_basename, find_calibration_folder, + get_frame_range, +) from .triangulate import load_offsets_dict @@ -368,6 +370,7 @@ def visualize_combined(config, pose_fname, cgroup, offsets_dict, pp = get_plotting_params(caps_2d, cap_3d, ang_names) nframes = pp['nframes'] + nframes_raw = int(caps_2d[0].get(cv2.CAP_PROP_FRAME_COUNT)) fps = pp['fps'] start_img = get_start_image(pp, ang_names) @@ -389,7 +392,17 @@ def visualize_combined(config, pose_fname, cgroup, offsets_dict, args=(writer, q)) thread.start() - for framenum in trange(nframes, ncols=70): + start_frame_num, end_frame_num = get_frame_range(config, nframes_raw, fps) + # set frame seeking for 2d vids + for cap in caps_2d: + cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame_num) + # keep 3d vid at frame 0; it should already be trimmed in label-3d + pp['nframes'] = min(end_frame_num - start_frame_num, pp['nframes']) + + if (end_frame_num - start_frame_num) != nframes: + print(f'Exporting with frame range {start_frame_num} to {end_frame_num} out of {nframes_raw}') + + for framenum in trange(start_frame_num, end_frame_num, ncols=70): ret, frames_2d, frame_3d = read_frames(caps_2d, cap_3d) if not ret: break diff --git a/anipose/label_videos_3d.py b/anipose/label_videos_3d.py index 766454a..0aaf18b 100644 --- a/anipose/label_videos_3d.py +++ b/anipose/label_videos_3d.py @@ -15,8 +15,10 @@ from collections import defaultdict from matplotlib.pyplot import get_cmap -from .common import make_process_fun, get_nframes, get_video_name, get_video_params, get_data_length, natural_keys - +from .common import ( + make_process_fun, get_nframes, get_video_name, get_video_params, + get_data_length, natural_keys, get_frame_range, +) def connect(points, bps, bp_dict, color): ixs = [bp_dict[bp] for bp in bps] @@ -129,7 +131,12 @@ def visualize_labels(config, labels_fname, outname, fps=300): mlab.view(focalpoint='auto', distance='auto') - for framenum in trange(data.shape[0], ncols=70): + nframes = data.shape[0] + start_frame_num, end_frame_num = get_frame_range(config, nframes, fps) + if (end_frame_num - start_frame_num) != nframes: + print(f'Exporting with frame range {start_frame_num} to {end_frame_num} out of {nframes}') + + for framenum in trange(start_frame_num, end_frame_num, ncols=70): fig.scene.disable_render = True if framenum in framedict: