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