| #!/usr/bin/env python3 | |
| """Overlay Pexels segmentation masks onto their raw videos. | |
| Reads pexels_videovo_{train,val,test}_dataset.csv rows (path, start_frame, | |
| end_frame, fps, mask_id, caption), matches each row to a locally downloaded | |
| raw video (videos/<videoId>.mp4) and its mask | |
| (pexels_masks/<index>_<videoId>/all_masks.npz, either already extracted or | |
| still inside a pexels-*.zip archive), and for the frame range | |
| [start_frame, end_frame] (inclusive) writes: | |
| --overlay-dir alpha-blended visualization (original frame + colored | |
| mask_id region), for eyeballing what got selected. | |
| --inpaint-mask-dir the original frame with the mask_id region's axis- | |
| aligned bounding box painted solid black -- a hard | |
| replace, not an alpha blend, and a bbox rather than | |
| the precise mask silhouette so the object's exact | |
| shape isn't leaked into the "inpaint this" signal. | |
| At least one of the two must be given; both may be given to produce both | |
| outputs from a single decode pass. | |
| """ | |
| import argparse | |
| import colorsys | |
| import csv | |
| import io | |
| import os | |
| import re | |
| import sys | |
| import zipfile | |
| from concurrent.futures import ProcessPoolExecutor, as_completed | |
| from pathlib import Path | |
| import av | |
| import cv2 | |
| import numpy as np | |
| from tqdm import tqdm | |
| HERE = Path(__file__).resolve().parent | |
| ROOT = HERE.parent | |
| DEFAULT_CSVS = [ | |
| ROOT / "pexels_videovo_train_dataset.csv", | |
| ] | |
| DEFAULT_VIDEOS_DIR = ROOT / "videos" | |
| DEFAULT_MASKS_DIR = ROOT / "pexels_masks" | |
| # CSV "path" values look like "000000010558_3525403.mp4": a 12-digit row | |
| # index followed by the pexels videoId, which is also how mask folders | |
| # (both on disk and inside the zips) are named. | |
| PATH_RE = re.compile(r"^(\d{12})_(\d+)\.mp4$") | |
| _MASK_INDEX: dict | None = None | |
| _MASKS_DIR: Path | None = None | |
| def build_mask_index(masks_dir: Path) -> dict[str, Path | None]: | |
| """Map '<index>_<videoId>' -> the zip file that holds it (None if already extracted).""" | |
| index: dict[str, Path | None] = {} | |
| for d in masks_dir.iterdir(): | |
| if d.is_dir() and (d / "all_masks.npz").exists(): | |
| index[d.name] = None | |
| for zip_path in sorted(masks_dir.glob("*.zip")): | |
| with zipfile.ZipFile(zip_path) as zf: | |
| for name in zf.namelist(): | |
| if name.endswith("/all_masks.npz"): | |
| index[name.split("/")[0]] = zip_path | |
| return index | |
| def load_mask_array(key: str, zip_path: Path | None, masks_dir: Path) -> np.ndarray: | |
| if zip_path is None: | |
| with open(masks_dir / key / "all_masks.npz", "rb") as f: | |
| return np.load(io.BytesIO(f.read()))["arr_0"] | |
| with zipfile.ZipFile(zip_path) as zf, zf.open(f"{key}/all_masks.npz") as f: | |
| return np.load(io.BytesIO(f.read()))["arr_0"] | |
| def make_palette(n: int = 256) -> np.ndarray: | |
| """Deterministic, visually distinct BGR palette; index 0 (background) is left black.""" | |
| palette = np.zeros((n, 3), dtype=np.uint8) | |
| for i in range(1, n): | |
| hue = (i * 0.6180339887) % 1.0 | |
| r, g, b = colorsys.hsv_to_rgb(hue, 0.85, 0.95) | |
| palette[i] = (int(b * 255), int(g * 255), int(r * 255)) | |
| return palette | |
| PALETTE = make_palette() | |
| def overlay_frame(frame: np.ndarray, mask: np.ndarray, alpha: float) -> np.ndarray: | |
| if mask.shape != frame.shape[:2]: | |
| mask = cv2.resize(mask, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_NEAREST) | |
| color = PALETTE[mask] | |
| blended = cv2.addWeighted(frame, 1 - alpha, color, alpha, 0) | |
| return np.where(mask[..., None] > 0, blended, frame) | |
| def black_out_bbox(frame: np.ndarray, mask: np.ndarray) -> np.ndarray: | |
| """Original frame with the nonzero mask region's bounding box painted | |
| solid black -- a hard replace (no alpha blending), and a bbox rather | |
| than the precise mask contour so the object's exact silhouette isn't | |
| leaked into the "inpaint this" signal.""" | |
| out = frame.copy() | |
| ys, xs = np.nonzero(mask) | |
| if ys.size: | |
| out[ys.min() : ys.max() + 1, xs.min() : xs.max() + 1] = 0 | |
| return out | |
| def open_writer( | |
| path: Path, fps: float, width: int, height: int, lossless: bool = False, | |
| ) -> tuple[av.container.OutputContainer, av.video.stream.VideoStream]: | |
| container = av.open(str(path), mode="w") | |
| stream = container.add_stream("libx264", rate=round(fps)) | |
| stream.width = width | |
| stream.height = height | |
| stream.pix_fmt = "yuv420p" | |
| options = {"preset": "veryfast"} | |
| if lossless: | |
| # crf=0 -> exact 0/255 values, no compression artifacts smearing the | |
| # bbox edges. Needed since the inpainting mask is read back as a | |
| # strict binary mask, not just viewed. | |
| options["crf"] = "0" | |
| stream.options = options | |
| stream.codec_context.thread_count = 2 | |
| return container, stream | |
| def write_frame(container: av.container.OutputContainer, stream: av.video.stream.VideoStream, rgb: np.ndarray) -> None: | |
| av_frame = av.VideoFrame.from_ndarray(rgb, format="rgb24").reformat(format="yuv420p") | |
| for packet in stream.encode(av_frame): | |
| container.mux(packet) | |
| def close_writer(container: av.container.OutputContainer, stream: av.video.stream.VideoStream) -> None: | |
| for packet in stream.encode(): | |
| container.mux(packet) | |
| container.close() | |
| def load_rows(csv_paths: list[Path]) -> list[dict]: | |
| rows = [] | |
| for csv_path in csv_paths: | |
| if not csv_path.exists(): | |
| print(f"warning: csv not found, skipping: {csv_path}", file=sys.stderr) | |
| continue | |
| with open(csv_path, newline="") as f: | |
| for row in csv.DictReader(f): | |
| m = PATH_RE.match(row["path"]) | |
| if m: | |
| row["_index"], row["_video_id"] = m.groups() | |
| rows.append(row) | |
| return rows | |
| def _init_worker(mask_index: dict, masks_dir: Path) -> None: | |
| global _MASK_INDEX, _MASKS_DIR | |
| _MASK_INDEX = mask_index | |
| _MASKS_DIR = masks_dir | |
| def process_row( | |
| row: dict, | |
| videos_dir: Path, | |
| overlay_dir: Path | None, | |
| inpaint_dir: Path | None, | |
| alpha: float, | |
| ) -> tuple[str, bool, str]: | |
| key = f"{row['_index']}_{row['_video_id']}" | |
| if key not in _MASK_INDEX: | |
| return key, False, "mask not found" | |
| video_path = videos_dir / f"{row['_video_id']}.mp4" | |
| if not video_path.exists(): | |
| return key, False, "raw video not found" | |
| start_frame = int(row["start_frame"]) | |
| end_frame = int(row["end_frame"]) # inclusive | |
| mask_id = int(row["mask_id"]) | |
| masks = load_mask_array(key, _MASK_INDEX[key], _MASKS_DIR) | |
| end_frame_excl = min(end_frame + 1, masks.shape[0]) | |
| if start_frame >= end_frame_excl: | |
| return key, False, "empty frame range" | |
| cap = cv2.VideoCapture(str(video_path)) | |
| fps = cap.get(cv2.CAP_PROP_FPS) or float(row.get("fps") or 30) | |
| width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame) | |
| clip_name = f"{key}_{start_frame}-{end_frame}.mp4" | |
| overlay_path = overlay_dir / clip_name if overlay_dir else None | |
| inpaint_path = inpaint_dir / clip_name if inpaint_dir else None | |
| # PyAV + libx264 (H.264, yuv420p) instead of cv2.VideoWriter: this | |
| # FFmpeg build's VideoWriter backend can only open the "mp4v" fourcc | |
| # (MPEG-4 Part 2), which most players/browsers can't play. Matches | |
| # write_video() in VIDGEN-1M/processing/select_and_trim_videos.py. | |
| overlay_writer = open_writer(overlay_path, fps, width, height) if overlay_path else None | |
| inpaint_writer = open_writer(inpaint_path, fps, width, height, lossless=True) if inpaint_path else None | |
| written = 0 | |
| try: | |
| for frame_idx in range(start_frame, end_frame_excl): | |
| ok, frame = cap.read() | |
| if not ok: | |
| break | |
| mask = np.where(masks[frame_idx] == mask_id, masks[frame_idx], 0) | |
| if overlay_writer: | |
| overlaid = overlay_frame(frame, mask, alpha) | |
| write_frame(*overlay_writer, cv2.cvtColor(overlaid, cv2.COLOR_BGR2RGB)) | |
| if inpaint_writer: | |
| blacked = black_out_bbox(frame, mask) | |
| write_frame(*inpaint_writer, cv2.cvtColor(blacked, cv2.COLOR_BGR2RGB)) | |
| written += 1 | |
| if overlay_writer: | |
| close_writer(*overlay_writer) | |
| overlay_writer = None | |
| if inpaint_writer: | |
| close_writer(*inpaint_writer) | |
| inpaint_writer = None | |
| finally: | |
| cap.release() | |
| # Only reached if an exception skipped the close_writer() calls above. | |
| if overlay_writer: | |
| overlay_writer[0].close() | |
| if inpaint_writer: | |
| inpaint_writer[0].close() | |
| if written == 0: | |
| if overlay_path: | |
| overlay_path.unlink(missing_ok=True) | |
| if inpaint_path: | |
| inpaint_path.unlink(missing_ok=True) | |
| return key, False, "no frames read" | |
| return key, True, f"wrote {written} frames -> {clip_name}" | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument( | |
| "--csv", type=Path, nargs="+", default=DEFAULT_CSVS, | |
| help="Dataset CSV(s) with path,start_frame,end_frame,fps,mask_id,caption columns", | |
| ) | |
| parser.add_argument("--videos-dir", type=Path, default=DEFAULT_VIDEOS_DIR, help="Directory of downloaded raw videos, named <videoId>.mp4") | |
| parser.add_argument("--masks-dir", type=Path, default=DEFAULT_MASKS_DIR, help="pexels_masks directory (zip archives and/or extracted folders)") | |
| parser.add_argument("--overlay-dir", type=Path, help="Directory to save alpha-blended overlay clips (original video + colored mask)") | |
| parser.add_argument("--inpaint-mask-dir", type=Path, help="Directory to save inpainting clips: original video with the mask_id region's bbox painted solid black (hard replace, no alpha)") | |
| parser.add_argument("--video-ids", nargs="+", help="Restrict to specific pexels videoIds") | |
| parser.add_argument("--limit", type=int, help="Max number of clips to process") | |
| parser.add_argument("--alpha", type=float, default=0.5, help="Overlay opacity for --overlay-dir clips (0-1)") | |
| parser.add_argument( | |
| "--workers", type=int, default=os.cpu_count() or 4, | |
| help="Parallel worker processes (default: all logical CPUs). Each clip's decode+encode is " | |
| "independent and CPU-bound, so this scales close to linearly with cores -- lower it if the " | |
| "machine is shared with other jobs.", | |
| ) | |
| args = parser.parse_args() | |
| if not args.overlay_dir and not args.inpaint_mask_dir: | |
| sys.exit("Specify at least one of --overlay-dir or --inpaint-mask-dir.") | |
| if args.overlay_dir: | |
| args.overlay_dir.mkdir(parents=True, exist_ok=True) | |
| if args.inpaint_mask_dir: | |
| args.inpaint_mask_dir.mkdir(parents=True, exist_ok=True) | |
| available = {p.stem for p in args.videos_dir.glob("*.mp4")} | |
| rows = load_rows(args.csv) | |
| rows = [r for r in rows if r["_video_id"] in available] | |
| if args.video_ids: | |
| wanted = set(args.video_ids) | |
| rows = [r for r in rows if r["_video_id"] in wanted] | |
| # A raw video can have multiple CSV rows (different mask_id/frame ranges). | |
| # Keep only the first row per video_id so each video is overlaid once. | |
| seen_videos: set[str] = set() | |
| deduped = [] | |
| for r in rows: | |
| if r["_video_id"] in seen_videos: | |
| continue | |
| seen_videos.add(r["_video_id"]) | |
| deduped.append(r) | |
| if len(deduped) < len(rows): | |
| print(f"Deduped {len(rows)} rows -> {len(deduped)} distinct videos.") | |
| rows = deduped | |
| if not rows: | |
| sys.exit("No CSV rows match the videos available in --videos-dir.") | |
| if args.limit: | |
| rows = rows[: args.limit] | |
| print(f"Indexing masks in {args.masks_dir} ...") | |
| mask_index = build_mask_index(args.masks_dir) | |
| print(f"Indexed {len(mask_index)} mask entries. Processing {len(rows)} clips with {args.workers} worker(s).") | |
| ok = 0 | |
| fail = 0 | |
| with ProcessPoolExecutor(max_workers=args.workers, initializer=_init_worker, initargs=(mask_index, args.masks_dir)) as pool: | |
| futures = { | |
| pool.submit(process_row, row, args.videos_dir, args.overlay_dir, args.inpaint_mask_dir, args.alpha): row | |
| for row in rows | |
| } | |
| with tqdm(total=len(futures), unit="clip", dynamic_ncols=True) as bar: | |
| for future in as_completed(futures): | |
| key, success, msg = future.result() | |
| if success: | |
| ok += 1 | |
| else: | |
| fail += 1 | |
| tqdm.write(f"{key}: FAIL ({msg})") | |
| bar.set_postfix(ok=ok, fail=fail) | |
| bar.update(1) | |
| dests = ", ".join(str(d) for d in (args.overlay_dir, args.inpaint_mask_dir) if d) | |
| print(f"\nDone: {ok}/{len(rows)} clips processed -> {dests}") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 13 kB
- Xet hash:
- d48bc9a47622636f4a7667541fa6bcb09ac627d5079d2ba8ef22dfc1daacde9f
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.