349 lines
14 KiB
Python
349 lines
14 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import time
|
|
import tarfile
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
from tqdm import tqdm
|
|
import yaml
|
|
|
|
from src.dataset import build_manifest, load_manifest, SampleConfig, VideoTripletDataset
|
|
|
|
|
|
def load_config(path: str) -> Dict:
|
|
with open(path, "r", encoding="utf-8") as fp:
|
|
return yaml.safe_load(fp)
|
|
|
|
|
|
def encode_png(frame_rgb: np.ndarray) -> bytes:
|
|
bgr = frame_rgb[:, :, ::-1]
|
|
ok, buf = cv2.imencode(".png", bgr)
|
|
if not ok:
|
|
raise RuntimeError("Failed to encode PNG")
|
|
return buf.tobytes()
|
|
|
|
|
|
def encode_jpg(frame_rgb: np.ndarray, quality: int) -> bytes:
|
|
bgr = frame_rgb[:, :, ::-1]
|
|
ok, buf = cv2.imencode(".jpg", bgr, [int(cv2.IMWRITE_JPEG_QUALITY), int(quality)])
|
|
if not ok:
|
|
raise RuntimeError("Failed to encode JPG")
|
|
return buf.tobytes()
|
|
|
|
|
|
def encode_image(frame_rgb: np.ndarray, fmt: str, jpeg_quality: int) -> bytes:
|
|
fmt = (fmt or "png").lower()
|
|
if fmt in ("jpg", "jpeg"):
|
|
return encode_jpg(frame_rgb, jpeg_quality)
|
|
return encode_png(frame_rgb)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", type=str, default="config.yaml")
|
|
parser.add_argument("--out", type=str, required=True)
|
|
parser.add_argument("--count", type=int, required=True)
|
|
parser.add_argument("--batch-size", type=int, default=8)
|
|
parser.add_argument("--num-workers", type=int, default=4)
|
|
parser.add_argument("--backend", type=str, default=None)
|
|
parser.add_argument("--shard-size", type=int, default=1000)
|
|
parser.add_argument("--image-format", type=str, default="png", choices=["png", "jpg", "jpeg"])
|
|
parser.add_argument("--jpeg-quality", type=int, default=95)
|
|
parser.add_argument("--log-path", type=str, default=None)
|
|
parser.add_argument("--log-interval", type=int, default=1000)
|
|
parser.add_argument("--loader-timeout", type=int, default=0)
|
|
parser.add_argument("--persistent-workers", action="store_true")
|
|
parser.add_argument("--resume", action="store_true")
|
|
parser.add_argument("--state-path", type=str, default=None)
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
import webdataset as wds
|
|
except Exception as exc:
|
|
raise RuntimeError("webdataset is not installed") from exc
|
|
|
|
cfg = load_config(args.config)
|
|
os.makedirs(args.out, exist_ok=True)
|
|
if not args.log_path:
|
|
args.log_path = os.path.join(args.out, "cache_wds.log.jsonl")
|
|
if not args.state_path:
|
|
args.state_path = os.path.join(args.out, "cache_wds.state.json")
|
|
if args.state_path:
|
|
state_dir = os.path.dirname(args.state_path)
|
|
if state_dir:
|
|
os.makedirs(state_dir, exist_ok=True)
|
|
|
|
data_cfg = cfg["data"]
|
|
manifest = data_cfg.get("manifest")
|
|
if not manifest:
|
|
if not data_cfg.get("roots"):
|
|
raise ValueError("data.roots must be set when data.manifest is empty")
|
|
manifest = os.path.join(args.out, "dataset_manifest.jsonl")
|
|
build_manifest(
|
|
data_cfg["roots"],
|
|
manifest,
|
|
data_cfg.get("extensions"),
|
|
backend=str(args.backend or data_cfg.get("backend", "opencv")),
|
|
errors_path=data_cfg.get("manifest_errors_log"),
|
|
force_count=bool(data_cfg.get("manifest_force_count", False)),
|
|
)
|
|
|
|
records = load_manifest(manifest)
|
|
sample_cfg = SampleConfig(
|
|
crop_size=tuple(data_cfg.get("train_crop", [256, 256])),
|
|
min_size=tuple(data_cfg.get("min_size", [256, 256])),
|
|
resize_short_edge=int(data_cfg.get("resize_short_edge", 0)),
|
|
random_flip=bool(data_cfg.get("random_horizontal_flip", True)),
|
|
stride_min=int(data_cfg.get("frame_stride_min", 1)),
|
|
stride_max=int(data_cfg.get("frame_stride_max", 1)),
|
|
)
|
|
|
|
dataset = VideoTripletDataset(
|
|
records=records,
|
|
sample_cfg=sample_cfg,
|
|
reader_cache=int(data_cfg.get("reader_cache", 4)),
|
|
backend=str(args.backend or data_cfg.get("backend", "opencv")),
|
|
max_tries=int(data_cfg.get("max_tries", 10)),
|
|
skip_failed=bool(data_cfg.get("skip_failed", True)),
|
|
max_total_tries=data_cfg.get("max_total_tries"),
|
|
return_path=True,
|
|
)
|
|
|
|
def build_loader() -> DataLoader:
|
|
timeout = int(args.loader_timeout) if args.num_workers > 0 else 0
|
|
return DataLoader(
|
|
dataset,
|
|
batch_size=args.batch_size,
|
|
shuffle=True,
|
|
num_workers=args.num_workers,
|
|
pin_memory=False,
|
|
drop_last=True,
|
|
timeout=timeout,
|
|
persistent_workers=bool(args.persistent_workers) if args.num_workers > 0 else False,
|
|
)
|
|
|
|
shard_id = -1
|
|
shard_count = 0
|
|
writer = None
|
|
log_fp: Optional[object] = None
|
|
if args.log_path:
|
|
os.makedirs(os.path.dirname(args.log_path), exist_ok=True)
|
|
log_fp = open(args.log_path, "a", encoding="utf-8")
|
|
|
|
def log_event(payload: Dict) -> None:
|
|
if log_fp is None:
|
|
return
|
|
payload["ts"] = time.strftime("%Y-%m-%d %H:%M:%S")
|
|
log_fp.write(json.dumps(payload, ensure_ascii=False) + "\n")
|
|
log_fp.flush()
|
|
|
|
def save_state(total_count: int, shard_id_val: int) -> None:
|
|
tmp_path = f"{args.state_path}.tmp"
|
|
payload = {
|
|
"total": int(total_count),
|
|
"shard": int(shard_id_val),
|
|
"shard_size": int(args.shard_size),
|
|
"image_format": (args.image_format or "png").lower(),
|
|
"updated_at": time.strftime("%Y-%m-%d %H:%M:%S"),
|
|
}
|
|
with open(tmp_path, "w", encoding="utf-8") as fp:
|
|
json.dump(payload, fp, ensure_ascii=False)
|
|
os.replace(tmp_path, args.state_path)
|
|
|
|
def _as_file_url(path: str) -> str:
|
|
path = os.path.abspath(path).replace("\\", "/")
|
|
return f"file:{path}"
|
|
|
|
def open_writer(sid: int):
|
|
name = os.path.join(args.out, f"shard-{sid:06d}.tar")
|
|
return wds.TarWriter(_as_file_url(name))
|
|
|
|
def find_shards(out_dir: str) -> List[Tuple[int, str]]:
|
|
shards: List[Tuple[int, str]] = []
|
|
for name in os.listdir(out_dir):
|
|
if not (name.startswith("shard-") and name.endswith(".tar")):
|
|
continue
|
|
sid_str = name[len("shard-") : -len(".tar")]
|
|
if not sid_str.isdigit():
|
|
continue
|
|
sid = int(sid_str)
|
|
shards.append((sid, os.path.join(out_dir, name)))
|
|
shards.sort(key=lambda x: x[0])
|
|
return shards
|
|
|
|
def count_samples_in_tar(path: str, ext: str) -> int:
|
|
count = 0
|
|
try:
|
|
with tarfile.open(path, "r") as tf:
|
|
for member in tf.getmembers():
|
|
name = member.name
|
|
if name.endswith(f"frame0.{ext}"):
|
|
count += 1
|
|
except Exception as exc:
|
|
log_event({"event": "tar_read_error", "path": path, "error": str(exc)})
|
|
return 0
|
|
return count
|
|
|
|
total = 0
|
|
fmt = (args.image_format or "png").lower()
|
|
ext = "jpg" if fmt in ("jpg", "jpeg") else "png"
|
|
if args.resume:
|
|
resume_total = 0
|
|
resume_shard = -1
|
|
if args.state_path and os.path.exists(args.state_path):
|
|
try:
|
|
with open(args.state_path, "r", encoding="utf-8") as fp:
|
|
state = json.load(fp)
|
|
resume_total = int(state.get("total", 0))
|
|
resume_shard = int(state.get("shard", resume_total // args.shard_size))
|
|
log_event(
|
|
{
|
|
"event": "resume_state",
|
|
"total": int(resume_total),
|
|
"shard": int(resume_shard),
|
|
"state_path": args.state_path,
|
|
}
|
|
)
|
|
except Exception as exc:
|
|
log_event({"event": "resume_state_error", "error": str(exc), "state_path": args.state_path})
|
|
if resume_total <= 0:
|
|
shards = find_shards(args.out)
|
|
if shards:
|
|
contiguous = []
|
|
expected = 0
|
|
for sid, path in shards:
|
|
if sid != expected:
|
|
log_event(
|
|
{
|
|
"event": "resume_gap",
|
|
"expected_shard": int(expected),
|
|
"found_shard": int(sid),
|
|
"path": path,
|
|
}
|
|
)
|
|
break
|
|
contiguous.append((sid, path))
|
|
expected += 1
|
|
if contiguous:
|
|
last_sid, last_path = contiguous[-1]
|
|
last_count = count_samples_in_tar(last_path, ext)
|
|
if last_count >= args.shard_size:
|
|
resume_total = len(contiguous) * args.shard_size
|
|
else:
|
|
resume_total = max(len(contiguous) - 1, 0) * args.shard_size
|
|
try:
|
|
os.remove(last_path)
|
|
log_event({"event": "resume_drop_partial", "path": last_path, "count": int(last_count)})
|
|
except Exception as exc:
|
|
log_event({"event": "resume_drop_error", "path": last_path, "error": str(exc)})
|
|
log_event(
|
|
{
|
|
"event": "resume_scan",
|
|
"total": int(resume_total),
|
|
"shards": int(len(contiguous)),
|
|
}
|
|
)
|
|
total = max(0, min(int(resume_total), int(args.count)))
|
|
if total % args.shard_size != 0:
|
|
floored = (total // args.shard_size) * args.shard_size
|
|
log_event({"event": "resume_floor", "from": int(total), "to": int(floored)})
|
|
total = floored
|
|
if total % args.shard_size == 0:
|
|
resume_path = os.path.join(args.out, f"shard-{total // args.shard_size:06d}.tar")
|
|
if os.path.exists(resume_path):
|
|
try:
|
|
os.remove(resume_path)
|
|
log_event({"event": "resume_drop_current", "path": resume_path})
|
|
except Exception as exc:
|
|
log_event({"event": "resume_drop_error", "path": resume_path, "error": str(exc)})
|
|
save_state(total, total // args.shard_size)
|
|
pbar = tqdm(total=args.count, initial=total, desc="wds-cache")
|
|
loader = build_loader()
|
|
data_iter = iter(loader)
|
|
try:
|
|
while total < args.count:
|
|
try:
|
|
batch = next(data_iter)
|
|
except StopIteration:
|
|
data_iter = iter(loader)
|
|
continue
|
|
except Exception as exc:
|
|
log_event({"event": "loader_error", "total": int(total), "error": str(exc)})
|
|
loader = build_loader()
|
|
data_iter = iter(loader)
|
|
continue
|
|
|
|
paths = batch.get("path")
|
|
batch_size = batch["frame0"].shape[0]
|
|
for i in range(batch_size):
|
|
sid = total // args.shard_size
|
|
if sid != shard_id:
|
|
if writer is not None:
|
|
writer.close()
|
|
writer = open_writer(sid)
|
|
shard_id = sid
|
|
shard_count = 0
|
|
save_state(total, shard_id)
|
|
log_event(
|
|
{
|
|
"event": "shard_open",
|
|
"shard": int(shard_id),
|
|
"total": int(total),
|
|
"path": os.path.join(args.out, f"shard-{sid:06d}.tar"),
|
|
}
|
|
)
|
|
path = paths[i] if paths is not None else None
|
|
f0 = (batch["frame0"][i].permute(1, 2, 0).numpy() * 255.0).astype("uint8")
|
|
f1 = (batch["frame1"][i].permute(1, 2, 0).numpy() * 255.0).astype("uint8")
|
|
gt = (batch["gt"][i].permute(1, 2, 0).numpy() * 255.0).astype("uint8")
|
|
try:
|
|
sample = {
|
|
"__key__": f"{total:08d}",
|
|
f"frame0.{ext}": encode_image(f0, fmt, args.jpeg_quality),
|
|
f"frame1.{ext}": encode_image(f1, fmt, args.jpeg_quality),
|
|
f"gt.{ext}": encode_image(gt, fmt, args.jpeg_quality),
|
|
}
|
|
writer.write(sample)
|
|
except Exception as exc:
|
|
log_event(
|
|
{
|
|
"event": "write_error",
|
|
"total": int(total),
|
|
"path": path,
|
|
"error": str(exc),
|
|
}
|
|
)
|
|
continue
|
|
shard_count += 1
|
|
total += 1
|
|
pbar.update(1)
|
|
if args.log_interval > 0 and total % int(args.log_interval) == 0:
|
|
log_event(
|
|
{
|
|
"event": "progress",
|
|
"total": int(total),
|
|
"shard": int(shard_id),
|
|
"path": path,
|
|
}
|
|
)
|
|
if total % args.shard_size == 0:
|
|
save_state(total, shard_id)
|
|
if total >= args.count:
|
|
break
|
|
finally:
|
|
pbar.close()
|
|
if writer is not None:
|
|
writer.close()
|
|
if log_fp is not None:
|
|
log_fp.close()
|
|
return
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|