Spaces:
Running
on
Zero
Running
on
Zero
| import pandas as pd | |
| from pathlib import Path | |
| import ijson | |
| from typing import Dict, Any | |
| from .video_base import VideoDataset | |
| class Ego4DVideoDataset(VideoDataset): | |
| def download(self): | |
| from ego4d.cli.cli import main_cfg as download_ego4d | |
| from ego4d.cli.config import Config as Ego4DConfig | |
| raw_dir = self.data_root / "raw" | |
| raw_dir.mkdir(parents=True, exist_ok=True) | |
| aws_credentials_path = Path.home() / ".aws" / "credentials" | |
| if not aws_credentials_path.exists(): | |
| raise FileNotFoundError( | |
| f"AWS credentials file not found at {aws_credentials_path}" | |
| "For Ego4D auto download, you need to request access and use the " | |
| "emailed key to set up AWS credentials first." | |
| "See https://ego4d-data.org/ for more information." | |
| ) | |
| cfg = Ego4DConfig( | |
| output_directory=str(raw_dir), | |
| datasets=["annotations", "clips"], | |
| benchmarks=["FHO"], | |
| metadata=True, | |
| assume_yes=True, | |
| ) | |
| import botocore | |
| try: | |
| download_ego4d(cfg) | |
| except botocore.exceptions.ClientError as e: | |
| print(e) | |
| raise RuntimeError( | |
| "Failed to download Ego4D dataset due to the above error." | |
| "If you see an error occurred (403) when calling the HeadObject operation: Forbidden", | |
| "It's likely due to an expired Ego4D AWS credential. Renew the dataset's online form and update the AWS credentials.", | |
| ) | |
| annotation_file = "v2/annotations/fho_main.json" | |
| print("Creating metadata CSV...") | |
| records = [] | |
| with open(raw_dir / annotation_file, "rb") as file: | |
| # Create a parser for the videos array | |
| videos = ijson.items(file, "videos.item") | |
| total = 0 | |
| for v in videos: | |
| fps = round(v["video_metadata"]["fps"]) | |
| n_frames = v["video_metadata"]["num_frames"] | |
| width = v["video_metadata"]["width"] | |
| height = v["video_metadata"]["height"] | |
| for c in v["annotated_intervals"]: | |
| video_path = "raw/v2/clips/" + c["clip_uid"] + ".mp4" | |
| if not Path(self.data_root / video_path).exists(): | |
| continue | |
| for a in c["narrated_actions"]: | |
| total += 1 | |
| critical_frames = a["clip_critical_frames"] | |
| is_valid_action = a["is_valid_action"] | |
| is_rejected = a["is_rejected"] | |
| is_invalid_annotation = a["is_invalid_annotation"] | |
| is_partial = a["is_partial"] | |
| if ( | |
| not critical_frames | |
| or not is_valid_action | |
| or is_rejected | |
| or is_invalid_annotation | |
| or is_partial | |
| ): | |
| continue | |
| caption = a["narration_text"] | |
| caption = ( | |
| caption.replace("#cC c ", " ") | |
| .replace("#Cc C ", " ") | |
| .replace("#C C ", "") | |
| .replace("#c c ", " ") | |
| .replace("#c- c ", " ") | |
| .replace("#c C ", " ") | |
| .replace("#c c", " ") | |
| .replace("#CC ", " ") | |
| .replace("#C C ", " ") | |
| .replace("#C c ", " ") | |
| .replace("#cc ", " ") | |
| .replace("#C- C ", " ") | |
| .replace("#c C ", " ") | |
| .replace("#C ", " ") | |
| .replace("#c ", " ") | |
| .replace("#", " ") | |
| ) | |
| pre_frame = critical_frames["pre_frame"] | |
| post_frame = critical_frames["post_frame"] | |
| pnr_frame = critical_frames["pnr_frame"] | |
| contact_frame = critical_frames["contact_frame"] | |
| # some manual heuristics to trim the video | |
| target_len = self._n_frames_in_src(fps) | |
| trim_start = pre_frame | |
| psudo_min_end = int((post_frame - pnr_frame) * 0.1) + pnr_frame | |
| if psudo_min_end - pre_frame >= target_len: | |
| trim_end = psudo_min_end | |
| elif post_frame - pnr_frame < target_len: | |
| trim_end = post_frame | |
| trim_start = max(trim_end - target_len, pre_frame - 15) | |
| else: | |
| trim_end = target_len + pre_frame | |
| trim_start = max(0, trim_start) | |
| trim_end = min(n_frames, trim_end) | |
| records.append( | |
| { | |
| "video_path": video_path, | |
| "height": height, | |
| "width": width, | |
| "n_frames": n_frames, | |
| "fps": fps, | |
| "original_caption": caption, | |
| "trim_start": trim_start, | |
| "trim_end": trim_end, | |
| "pre_frame": pre_frame, | |
| "pnr_frame": pnr_frame, | |
| "post_frame": post_frame, | |
| "contact_frame": contact_frame, | |
| } | |
| ) | |
| metadata_path = self.data_root / self.metadata_path | |
| metadata_path.parent.mkdir(parents=True, exist_ok=True) | |
| df = pd.DataFrame.from_records(records) | |
| df.to_csv(metadata_path, index=False) | |
| print(f"Created metadata CSV with {len(records)} records") | |