| #!/usr/bin/env python3 | |
| """Sample rows from pexels.csv and download the videos, named by videoId.""" | |
| import argparse | |
| import csv | |
| import random | |
| import re | |
| import sys | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| from pathlib import Path | |
| import requests | |
| HERE = Path(__file__).resolve().parent | |
| DEFAULT_CSV = "pexels.csv" | |
| DEFAULT_OUT = "videos" | |
| RESOLUTION_RE = re.compile(r"_(\d+)_(\d+)_\d+fps") | |
| MASK_INFO_ID_RE = re.compile(r"_(\d+)\.mp4$") | |
| def load_mask_info_ids(csv_path: Path) -> set[str]: | |
| with open(csv_path, newline="") as f: | |
| rows = csv.DictReader(f) | |
| return { | |
| m.group(1) | |
| for row in rows | |
| if (m := MASK_INFO_ID_RE.search(row["path"])) | |
| } | |
| def load_rows(csv_path: Path, resolution: str | None, mask_info_csv: Path | None) -> list[dict]: | |
| with open(csv_path, newline="") as f: | |
| rows = list(csv.DictReader(f)) | |
| if resolution is not None: | |
| width, height = resolution.lower().split("x") | |
| rows = [ | |
| row | |
| for row in rows | |
| if (m := RESOLUTION_RE.search(row["link"])) | |
| and m.group(1) == width | |
| and m.group(2) == height | |
| ] | |
| if mask_info_csv is not None: | |
| mask_info_ids = load_mask_info_ids(mask_info_csv) | |
| rows = [row for row in rows if row["videoId"] in mask_info_ids] | |
| return rows | |
| def download_one(row: dict, out_dir: Path, timeout: int) -> tuple[str, bool, str]: | |
| video_id = row["videoId"] | |
| url = row["link"] | |
| ext = Path(url).suffix or ".mp4" | |
| dest = out_dir / f"{video_id}{ext}" | |
| if dest.exists() and dest.stat().st_size > 0: | |
| return video_id, True, "already exists" | |
| tmp = dest.with_suffix(dest.suffix + ".part") | |
| try: | |
| with requests.get(url, stream=True, timeout=timeout) as resp: | |
| resp.raise_for_status() | |
| with open(tmp, "wb") as f: | |
| for chunk in resp.iter_content(chunk_size=1 << 20): | |
| if chunk: | |
| f.write(chunk) | |
| tmp.rename(dest) | |
| return video_id, True, "downloaded" | |
| except Exception as e: | |
| tmp.unlink(missing_ok=True) | |
| return video_id, False, str(e) | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--csv", type=Path, default=DEFAULT_CSV, help="Path to pexels.csv") | |
| parser.add_argument( | |
| "--mask-info-csv", | |
| type=Path, | |
| default=None, | |
| help="Path to pexels_videovo_train_dataset.csv; if given, only videoIds present in both files are kept", | |
| ) | |
| parser.add_argument("--out", type=Path, default=DEFAULT_OUT, help="Output directory for videos") | |
| parser.add_argument("-n", "--num-samples", type=int, default=2000, help="Number of videos to sample") | |
| parser.add_argument("--seed", type=int, default=42, help="Random seed for sampling") | |
| parser.add_argument("--workers", type=int, default=8, help="Concurrent downloads") | |
| parser.add_argument("--timeout", type=int, default=60, help="Per-request timeout in seconds") | |
| parser.add_argument( | |
| "--resolution", | |
| default="1920x1080", | |
| help="Only keep videos matching this WxH (e.g. 1920x1080). Use 'any' to disable filtering.", | |
| ) | |
| args = parser.parse_args() | |
| resolution = None if args.resolution.lower() == "any" else args.resolution | |
| rows = load_rows(args.csv, resolution, args.mask_info_csv) | |
| if args.num_samples > len(rows): | |
| sys.exit(f"Requested {args.num_samples} samples but only {len(rows)} matching rows in {args.csv}") | |
| random.seed(args.seed) | |
| random.shuffle(rows) | |
| pool_iter = iter(rows) | |
| args.out.mkdir(parents=True, exist_ok=True) | |
| ok = 0 | |
| failed = [] | |
| tried = 0 | |
| with ThreadPoolExecutor(max_workers=args.workers) as pool: | |
| while ok < args.num_samples: | |
| need = args.num_samples - ok | |
| batch = [] | |
| for _ in range(need): | |
| try: | |
| batch.append(next(pool_iter)) | |
| except StopIteration: | |
| break | |
| if not batch: | |
| sys.exit( | |
| f"Ran out of rows after {tried} attempts; only got {ok}/{args.num_samples}. " | |
| f"Failed videoIds: {failed}" | |
| ) | |
| futures = { | |
| pool.submit(download_one, row, args.out, args.timeout): row["videoId"] | |
| for row in batch | |
| } | |
| for future in as_completed(futures): | |
| video_id, success, msg = future.result() | |
| tried += 1 | |
| status = "OK" if success else "FAIL" | |
| print(f"[{tried}] {video_id}: {status} ({msg})") | |
| if success: | |
| ok += 1 | |
| else: | |
| failed.append(video_id) | |
| print(f"\nDone: {ok}/{args.num_samples} succeeded ({tried} attempted).") | |
| if failed: | |
| print(f"Failed videoIds (replaced with new samples): {failed}") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 5.01 kB
- Xet hash:
- ba7a714b6ba2037d1256789ed3d12e6f56af36e745d7c5c07caed3b86afd3dff
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.