trans_video.py

# Copyright 2026 OPPO. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Merge short egocentric clips into longer streaming videos.

The script scans a directory of short `.mp4` clips whose file names embed a
timestamp (e.g. `DAY7_A2_ALICE_<timestamp>.mp4`), sorts them chronologically,
groups consecutive clips spread across different time segments, and uses
`ffmpeg` to concatenate each group into a single longer video. This is used to
build the long-form egocentric videos referenced by the benchmark.
"""

import os
import re
import argparse
import subprocess


def parse_time_from_filename(filename):
    """Extract the integer timestamp embedded in a clip file name."""
    match = re.search(r"_(\d+)\.mp4$", filename)
    if match:
        return int(match.group(1))
    return 0


def merge_videos(video_files, output_file):
    """Concatenate the given clips into a single video using ffmpeg."""
    if not video_files:
        return False

    list_file = "temp_file_list.txt"
    with open(list_file, "w", encoding="utf-8") as f:
        for video_file in video_files:
            f.write(f"file '{video_file}'\n")

    cmd = [
        "ffmpeg",
        "-f", "concat",
        "-safe", "0",
        "-i", list_file,
        "-c", "copy",
        output_file,
        "-y",
    ]

    try:
        subprocess.run(cmd, check=True, capture_output=True, text=True)
        print(f"[OK] merged: {os.path.basename(output_file)}")
        success = True
    except subprocess.CalledProcessError as e:
        print(f"[FAIL] merge failed: {os.path.basename(output_file)}")
        print(f"stderr: {e.stderr}")
        success = False
    except FileNotFoundError:
        print("[FAIL] ffmpeg not found; please install ffmpeg and add it to PATH")
        success = False

    if os.path.exists(list_file):
        os.remove(list_file)

    return success


def create_distributed_groups(video_files, group_size, num_merged_videos_per_class=10):
    """Build groups of `group_size` clips spread across different time segments.

    Produces up to `num_merged_videos_per_class` groups per call.
    """
    if not video_files:
        return []

    sorted_videos = sorted(video_files, key=parse_time_from_filename)
    total_videos = len(sorted_videos)

    total_groups_needed = num_merged_videos_per_class * group_size

    if total_videos < total_groups_needed:
        print(f"Warning: need at least {total_groups_needed} clips, but only {total_videos} are available")
        num_merged_videos_per_class = total_videos // group_size
        print(f"Adjusting to {num_merged_videos_per_class} merged videos per class")

    num_segments = min(4, num_merged_videos_per_class)
    videos_per_segment = total_videos // num_segments if num_segments else total_videos

    all_groups = []

    for group_idx in range(num_merged_videos_per_class):
        segment_start = (group_idx % num_segments) * videos_per_segment
        segment_end = min(segment_start + videos_per_segment, total_videos)

        start_idx = segment_start + (group_idx // num_segments) * group_size
        end_idx = start_idx + group_size

        if end_idx <= segment_end:
            group_videos = sorted_videos[start_idx:end_idx]
        else:
            group_videos = sorted_videos[start_idx:segment_end]
            remaining = group_size - len(group_videos)
            if remaining > 0:
                group_videos.extend(sorted_videos[:remaining])

        if len(group_videos) == group_size:
            all_groups.append(group_videos)
        else:
            print(f"Warning: could not collect {group_size} clips for group {group_idx}")

    return all_groups


def main():
    parser = argparse.ArgumentParser(
        description="Merge short egocentric clips into longer videos via ffmpeg."
    )
    parser.add_argument("--video_dir", type=str, required=True,
                        help="Directory containing the short .mp4 clips.")
    parser.add_argument("--output_dir", type=str, required=True,
                        help="Directory where merged videos are written.")
    parser.add_argument("--prefix", type=str, required=True,
                        help="Clip file-name prefix to match, e.g. 'DAY7_A2_ALICE_'.")
    parser.add_argument("--group_sizes", type=int, nargs="+", default=[6, 8, 10],
                        help="Number of clips per merged video (one or more values).")
    parser.add_argument("--num_per_class", type=int, default=10,
                        help="Number of merged videos to produce per group size.")
    args = parser.parse_args()

    if not os.path.exists(args.video_dir):
        print(f"Error: directory does not exist: {args.video_dir}")
        return

    video_files = [
        os.path.join(args.video_dir, f)
        for f in os.listdir(args.video_dir)
        if f.endswith(".mp4") and f.startswith(args.prefix)
    ]

    if not video_files:
        print(f"No .mp4 files starting with '{args.prefix}' were found")
        return

    print(f"Found {len(video_files)} clips in {args.video_dir}")

    base_name = args.prefix.rstrip("_")

    for group_size in args.group_sizes:
        print(f"\n{'=' * 60}")
        print(f"Processing groups of {group_size} clips...")

        distributed_groups = create_distributed_groups(
            video_files, group_size, args.num_per_class
        )

        if not distributed_groups:
            print(f"Could not build any group for group size {group_size}")
            continue

        print(f"Will create {len(distributed_groups)} merged videos")

        output_dir = os.path.join(args.output_dir, f"merged_{group_size}_videos")
        os.makedirs(output_dir, exist_ok=True)

        success_count = 0
        for i, group in enumerate(distributed_groups):
            print(f"\nProcessing group {i + 1}/{len(distributed_groups)}:")
            print(f"clips: {[os.path.basename(f) for f in group]}")

            first_time = parse_time_from_filename(os.path.basename(group[0]))
            last_time = parse_time_from_filename(os.path.basename(group[-1]))

            output_filename = f"{base_name}_{first_time:08d}_{last_time:08d}.mp4"
            output_path = os.path.join(output_dir, output_filename)

            if merge_videos(group, output_path):
                success_count += 1

        print(f"\nGroup size {group_size}: created {success_count}/{len(distributed_groups)} merged videos")
        print(f"Output directory: {output_dir}")

    print(f"\n{'=' * 60}")
    print("All done!")


if __name__ == "__main__":
    main()