hugging-apps/echo-memory

1#!/usr/bin/env python32"""3Preprocess Context-as-Memory dataset folders into Echo-Memory metadata CSV.4 5Expected dataset layout:6- frames/: frame images organized by video7- jsons/: camera pose information for each video8- overlap_labels/: FOV overlap information for memory retrieval9- captions.txt: video segment captions10"""11 12import argparse13import csv14import json15import os16from typing import Dict, List, Tuple17 18 19def parse_caption_line(line: str) -> Tuple[str, str]:20    """21    Parse a line from captions.txt.22 23    Format: "video_name/start_end.mp4\tcaption text..."24    Returns: (video_path, caption)25    """26    parts = line.strip().split("\t", 1)27    if len(parts) != 2:28        return None, None29    video_path = parts[0]30    caption = parts[1]31    return video_path, caption32 33 34def load_captions(captions_file: str) -> Dict[str, str]:35    """Load captions.txt as video_name -> caption."""36    captions = {}37    if not os.path.exists(captions_file):38        print(f"Warning: Captions file not found: {captions_file}")39        return captions40 41    with open(captions_file, "r", encoding="utf-8") as f:42        for line in f:43            video_path, caption = parse_caption_line(line)44            if video_path and caption:45                video_name = video_path.split("/")[0]46                if video_name not in captions:47                    captions[video_name] = []48                captions[video_name].append(caption)49 50    for video_name in captions:51        captions[video_name] = captions[video_name][0] if captions[video_name] else ""52 53    return captions54 55 56def get_frame_files(frames_dir: str, video_name: str) -> List[str]:57    """Get sorted frame paths for one video, relative to frames_dir."""58    video_frames_dir = os.path.join(frames_dir, video_name)59    if not os.path.exists(video_frames_dir):60        return []61 62    frame_files = []63    for frame_file in sorted(os.listdir(video_frames_dir)):64        if frame_file.endswith(".png"):65            frame_files.append(os.path.join(video_name, frame_file))66 67    return frame_files68 69 70def load_camera_poses(json_file: str) -> Dict:71    """Load camera poses from a JSON file."""72    if not os.path.exists(json_file):73        return {}74 75    with open(json_file, "r", encoding="utf-8") as f:76        data = json.load(f)77 78    if "CineCameraActor" in data:79        return data["CineCameraActor"]80    if isinstance(data, dict):81        return data82    return {}83 84 85def load_overlap_labels(overlap_dir: str, video_name: str, frame_idx: int) -> List[int]:86    """Load overlapping frame indices for a given frame."""87    overlap_file = os.path.join(overlap_dir, video_name, f"{frame_idx}.json")88    if not os.path.exists(overlap_file):89        return []90 91    try:92        with open(overlap_file, "r", encoding="utf-8") as f:93            data = json.load(f)94            overlapping_frames = data.get("overlapping_frames", [])95            return [int(frame) for frame in overlapping_frames if str(frame).isdigit()]96    except Exception:97        return []98 99 100def create_metadata_csv(101    dataset_base_path: str,102    output_csv: str,103    segment_length: int = 81,104    context_frames: int = 5,105):106    """107    Create metadata CSV for the Context-as-Memory dataset.108 109    Args:110        dataset_base_path: root of the dataset.111        output_csv: output CSV path.112        segment_length: frames per training segment.113        context_frames: context frames reserved by downstream workflows.114    """115    frames_dir = os.path.join(dataset_base_path, "frames")116    captions_file = os.path.join(dataset_base_path, "captions.txt")117 118    captions = load_captions(captions_file)119 120    if not os.path.exists(frames_dir):121        print(f"Error: Frames directory not found: {frames_dir}")122        return123 124    video_names = [125        d for d in os.listdir(frames_dir)126        if os.path.isdir(os.path.join(frames_dir, d))127    ]128 129    print(f"Found {len(video_names)} videos")130    print(f"Context frames: {context_frames}")131 132    output_dir = os.path.dirname(output_csv)133    if output_dir:134        os.makedirs(output_dir, exist_ok=True)135 136    with open(output_csv, "w", newline="", encoding="utf-8") as csvfile:137        fieldnames = [138            "video",139            "prompt",140            "video_name",141            "start_frame",142            "end_frame",143        ]144        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)145        writer.writeheader()146 147        total_segments = 0148 149        for video_name in sorted(video_names):150            print(f"Processing video: {video_name}")151 152            frame_files = get_frame_files(frames_dir, video_name)153            if len(frame_files) < segment_length:154                print(155                    f"  Skipping {video_name}: only {len(frame_files)} frames "156                    f"(need at least {segment_length})"157                )158                continue159 160            prompt = captions.get(video_name, f"A scene from {video_name}")161            step = max(1, segment_length // 2)162            video_segments = 0163 164            for start_idx in range(0, len(frame_files) - segment_length + 1, step):165                end_idx = start_idx + segment_length - 1166                segment_frames = frame_files[start_idx:end_idx + 1]167 168                if len(segment_frames) < segment_length:169                    continue170 171                frame_paths = "|".join(segment_frames)172                video_path = os.path.join("frames", frame_paths)173 174                writer.writerow({175                    "video": video_path,176                    "prompt": prompt,177                    "video_name": video_name,178                    "start_frame": start_idx,179                    "end_frame": end_idx,180                })181 182                total_segments += 1183                video_segments += 1184 185            print(f"  Created {video_segments} segments for {video_name}")186 187        print(f"\nTotal segments created: {total_segments}")188        print(f"Metadata CSV saved to: {output_csv}")189 190 191def main():192    parser = argparse.ArgumentParser(description="Preprocess Context-as-Memory Dataset")193    parser.add_argument(194        "--dataset_base_path",195        type=str,196        required=True,197        help="Base path to Context-as-Memory dataset",198    )199    parser.add_argument(200        "--output_csv",201        type=str,202        default="metadata.csv",203        help="Output CSV file path (default: metadata.csv)",204    )205    parser.add_argument(206        "--segment_length",207        type=int,208        default=81,209        help="Length of video segments (default: 81 frames)",210    )211    parser.add_argument(212        "--context_frames",213        type=int,214        default=5,215        help="Number of context frames (default: 5)",216    )217 218    args = parser.parse_args()219 220    if not os.path.isabs(args.output_csv):221        args.output_csv = os.path.join(args.dataset_base_path, args.output_csv)222 223    create_metadata_csv(224        dataset_base_path=args.dataset_base_path,225        output_csv=args.output_csv,226        segment_length=args.segment_length,227        context_frames=args.context_frames,228    )229 230 231if __name__ == "__main__":232    main()233