Add start and end params for video creation

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)
 

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)

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

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: