khwstolle/csvps

1#!/usr/bin/env python2r"""3Builds a WebDataset from the Cityscapes Video dataset.4 5Adapted from the `WebDataset documentation<https://github.com/webdataset/webdataset/>`_.6"""7 8import itertools9import collections10import typing as T11from pprint import pformat12import argparse13import multiprocessing as mp14import tarfile15import pandas as pd16from io import BytesIO17import json18 19from pathlib import Path20from tqdm import tqdm21 22 23def parse_args():24    ap = argparse.ArgumentParser(25        description="Build a WebDataset from the Cityscapes Video dataset."26    )27 28    # Flags and optional29    ap.add_argument(30        "--shard-size",31        "-s",32        type=int,33        default=10,34        help=("Number of sequences per shard."),35    )36    ap.add_argument(37        "--name",38        "-n",39        type=str,40        default="csvps",41        help=(42            "Name of the dataset. This will be used as the prefix for the tar files."43        ),44    )45    ap.add_argument(46        "--variant",47        type=str,48        default="",49        help=(50            "When passing different manifest variants, this will be used to postfix "51            "each split such that the resulting dataset name is unique."52        ),53    )54    ap.add_argument(55        "--force", "-f", action="store_true", help="Overwrite existing data."56    )57    ap.add_argument(58        "--splits", nargs="+", default=["train", "val", "test"], help="Splits to build."59    )60    ap.add_argument("--compression", "-c", default="", help="Compression to use")61 62    # Positional63    ap.add_argument("manifest", type=Path, help="Path to the manifest CSV file.")64    ap.add_argument("data", type=Path, help="Path to the Cityscapes Video dataset.")65    ap.add_argument("output", type=Path, help="Path to the output directory.")66 67    rt = ap.parse_args()68 69    # Validation70    if rt.shard_size < 1:71        ap.error("Shard size must be a positive integer.")72    if rt.name == "":73        ap.error("Name must be a non-empty string.")74    if not rt.name.isalnum() and not rt.name.islower():75        ap.error("Name must be a lowercase alpha-numeric string.")76    if rt.variant != "" and not rt.variant.isalnum() and not rt.variant.islower():77        ap.error("Variant must be a lowercase alpha-numeric string.")78    if not rt.manifest.exists():79        ap.error(f"Manifest file not found: {rt.manifest}")80    if not rt.data.exists():81        ap.error(f"Data directory not found: {rt.data}")82    if not rt.output.exists():83        rt.output.mkdir(parents=True)84        print(f"Created output directory: {rt.output}")85 86    return rt87 88 89PAD_TO: T.Final[int] = 6  # 06-padding is given by the dataset and should not be changed90 91 92def pad_number(n: int) -> str:93    r"""94    For sorting, numbers are padded with zeros to a fixed width.95    """96    if not isinstance(n, int):97        msg = f"Expected an integer, got {n} of type {type(n)}"98        raise TypeError(msg)99    return f"{n:0{PAD_TO}d}"100 101 102def read_timestamp(path: Path) -> int:103    with path.open("r") as f:104        ts = f.read().strip()105    if not ts.isdigit():106        msg = f"Expected a timestamp, got {ts} from {path}"107        raise ValueError(msg)108    return int(ts)109 110 111def write_bytes(tar: tarfile.TarFile, bt: bytes, arc: str):112    r""" "113    Simple utility to write the bytes (e.g. metadata json) directly from memory to114    the tarfile, since these do not exist as a file.115    """116    with BytesIO() as buf:117        buf.write(bt)118 119        # The TarInfo object must be created manually since the meta-data120        # JSON is written to a buffer (BytesIO) and not a file.121        tar_info = tarfile.TarInfo(arc)122        tar_info.size = buf.tell()  # number of bytes written123 124        # Reset the buffer to the beginning before adding it to the tarfile125        buf.seek(0)126 127        tar.addfile(tar_info, buf)128 129 130def find_sequence_files(131    seq: int,132    group: pd.DataFrame,133    *,134    data_dir: Path,135    dataset_name: str,136    compression: str,137    missing_ok: bool = False,138    frame_inputs: T.Sequence[str] = ("image.png", "vehicle.json"),139    frame_annotations: T.Sequence[str] = ("panoptic.png", "depth.tiff"),140    sequence_data: T.Sequence[str] = ("camera.json",),141    separator: str = "/",142) -> T.Iterator[tuple[Path | bytes, str]]:143    seq_pad = pad_number(seq)144    seq_dir = data_dir / seq_pad145 146    group = group.sort_values("frame")147 148    # Add frame-wise data149    primary_keys = group.index.tolist()150    frame_numbers = list(map(pad_number, group["frame"].tolist()))151 152    for i, meta in enumerate(153        group.drop(columns=["sequence", "frame", "split"]).to_dict(154            orient="records", index=True155        )156    ):157        frame_06 = frame_numbers[i]158        is_ann = meta["is_annotated"]159 160        # Write primary key161        meta["primary_key"] = primary_keys[i]162 163        # Add files to the tarfile164        for var in frame_inputs + frame_annotations:165            path_file = seq_dir / f"{frame_06}.{var}"166            if not path_file.exists():167                if missing_ok or (var in frame_annotations and not is_ann):168                    continue  # missing annotation OK169                msg = f"File not found: {path_file}"170                raise FileNotFoundError(msg)171 172            yield (173                path_file,174                separator.join(175                    (176                        dataset_name,177                        # {seq}.{frame}.{var}.{ext}178                        path_file.relative_to(data_dir).as_posix().replace("/", "."),179                    )180                ),181            )182 183        # Add the timestamp to the meta-data if it exists184        path_ts = seq_dir / f"{frame_06}.timestamp.txt"185        if not path_ts.exists():186            if not missing_ok:187                msg = f"Timestamp file not found: {path_ts}"188                raise FileNotFoundError(msg)189            meta["timestamp"] = None190        else:191            meta["timestamp"] = read_timestamp(path_ts)192 193        # Write frame metadata194        yield (195            json.dumps(meta).encode("utf-8"),196            f"{dataset_name}/{seq_pad}.{frame_06}.metadata.json",197        )198 199    # Add sequence-wise files {seq}.{var}.{ext}, e.g. 000000.camera.json200    for var in sequence_data:201        path_file = seq_dir.with_suffix("." + var)202        if not path_file.exists():203            if missing_ok:204                continue205            msg = f"File not found: {path_file}"206            raise FileNotFoundError(msg)207 208        yield (209            path_file,210            separator.join(211                (212                    dataset_name,213                    # {seq}.{var}.{ext}214                    path_file.relative_to(data_dir).as_posix(),215                )216            ),217        )218 219    # Write frames array220    yield (221        json.dumps(frame_numbers).encode("utf-8"),222        f"{dataset_name}/{seq_pad}.frames.json",223    )224 225 226def run_collector(227    seq: int, group: pd.DataFrame, kwargs: dict228) -> tuple[int, list[tuple[Path | bytes, str]]]:229    r"""230    Worker that collects the files for a single sequence.231    """232    return (seq, list(find_sequence_files(seq, group, **kwargs)))233 234 235def run_writer(236    tar_path: Path, items: list[list[tuple[Path | bytes, str]]], compression: str = "gz"237) -> None:238    r"""239    Worker that writes the files to a tar archive.240    """241    if compression != "":242        tar_path = tar_path.with_suffix(f".tar.{compression}")243        write_mode = f"w:{compression}"244    else:245        tar_path = tar_path.with_suffix(".tar")246        write_mode = "w"247 248    with tarfile.open(tar_path, write_mode) as tar:249        for item in itertools.chain.from_iterable(items):250            try:251                path, arc = item252            except ValueError:253                msg = f"Expected a tuple of length 2, got {item}"254                raise ValueError(msg)255 256            if isinstance(path, Path):257                tar.add(path, arcname=arc)258            else:259                write_bytes(tar, path, arc)260 261 262def build_shard(263    mfst: pd.DataFrame,264    *,265    tar_dir: Path,266    shard_size: int,267    **kwargs,268):269    # Make dirs270    tar_dir.mkdir(exist_ok=True, parents=True)271 272    write_log = collections.defaultdict(list)273 274    # Create a list of all sequences275    # groups = [(seq, group) for seq, group in mfst.groupby("sequence")]276    # shards = [groups[i : i + shard_size] for i in range(0, len(groups), shard_size)277    n_groups = len(mfst["sequence"].unique())278    n_shards = n_groups // shard_size279 280    targets = [None] * n_groups281 282    # Start a multiprocessing pool283    n_proc = min(mp.cpu_count(), 16)284    with mp.Pool(n_proc) as pool:285        write_jobs: list[mp.AsyncResult] = []286 287        # Data collection288        with tqdm(total=n_groups, desc="Collecting data") as pbar_group:289            for seq, files in pool.starmap(290                run_collector,291                [(seq, group, kwargs) for seq, group in mfst.groupby("sequence")],292                chunksize=min(8, shard_size),293            ):294                assert targets[seq] is None, f"Duplicate sequence: {seq}"295 296                pbar_group.update()297 298                # Write to the file specs list299                targets[seq] = files300 301                # Get a view of only the current shards's files302                shard_index = seq // shard_size303                shard_offset = shard_index * shard_size304                shard_specs = targets[shard_offset : shard_offset + shard_size]305 306                # Pad the shard index307                shard_06 = pad_number(shard_index)308 309                write_log[shard_06].append(pad_number(seq))310 311                # If the shard is fully populated, write it to a tar file in another process312                if all(s is not None for s in shard_specs):313                    tar_path = tar_dir / shard_06314 315                    write_jobs.append(316                        pool.apply_async(317                            run_writer,318                            (tar_path, shard_specs, ""),319                        )320                    )321 322        # Wait for write-workers to finish generating the TAR files323        with tqdm(total=n_shards, desc="Writing shards") as pbar_shard:324            for j in write_jobs:325                j.get()326                pbar_shard.update()327 328        pool.close()329        pool.join()330 331    print("Created shard files:\n" + pformat(dict(write_log)))332 333 334def main():335    args = parse_args()336    manifest = pd.read_csv(args.manifest, index_col="primary_key")337 338    # For each split, build a tar archive containing the sorted files339    for split in args.splits:340        split_out = "-".join([s for s in (split, args.variant) if len(s) > 0])341        tar_dir = args.output / split_out342 343        if tar_dir.exists():344            if args.force:345                print(f"Removing existing dataset: {tar_dir}")346                for f in tar_dir.glob("*.tar"):347                    f.unlink()348            else:349                msg = f"Dataset already exists: {tar_dir}"350                raise FileExistsError(msg)351 352        print(f"Generating {split_out} split...")353 354        build_shard(355            manifest[manifest["split"] == split],356            tar_dir=tar_dir,357            data_dir=args.data / split,358            shard_size=args.shard_size,359            dataset_name=f"{args.name}-{split_out}",360            missing_ok=True,361            compression=args.compression362        )363 364 365if __name__ == "__main__":366    main()367