File size: 3,899 Bytes
01857e4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | #!/usr/bin/env python3
"""
Rewrite absolute paths embedded in processed datasets (e.g. after copying project p32591 -> p32685).
Updates:
- backup_entries/*.json (image, label_mask string fields; template JSON if present)
- template/template_summary.json
- hf_dataset/ (HuggingFace DatasetDict on disk) columns image, label_mask
"""
from __future__ import annotations
import argparse
import json
import os
import shutil
import sys
import tempfile
def replace_in_obj(obj, old: str, new: str):
if isinstance(obj, str):
return obj.replace(old, new)
if isinstance(obj, list):
return [replace_in_obj(x, old, new) for x in obj]
if isinstance(obj, dict):
return {k: replace_in_obj(v, old, new) for k, v in obj.items()}
return obj
def fix_backup_entries(processed_dir: str, old: str, new: str) -> int:
n = 0
bdir = os.path.join(processed_dir, "backup_entries")
for name in ("train_entries.json", "val_entries.json"):
path = os.path.join(bdir, name)
if not os.path.isfile(path):
continue
with open(path) as f:
data = json.load(f)
data = replace_in_obj(data, old, new)
with open(path, "w") as f:
json.dump(data, f, indent=2)
n += 1
return n
def fix_template_summary(processed_dir: str, old: str, new: str) -> bool:
path = os.path.join(processed_dir, "template", "template_summary.json")
if not os.path.isfile(path):
return False
with open(path) as f:
data = json.load(f)
data = replace_in_obj(data, old, new)
with open(path, "w") as f:
json.dump(data, f, indent=2)
return True
def fix_hf_dataset(processed_dir: str, old: str, new: str) -> bool:
hf_path = os.path.join(processed_dir, "hf_dataset")
if not os.path.isdir(hf_path):
return False
from datasets import DatasetDict, load_from_disk
ds = load_from_disk(hf_path)
def fix_row(ex):
out = dict(ex)
for key in ("image", "label_mask"):
if key in out and isinstance(out[key], str):
out[key] = out[key].replace(old, new)
return out
new_splits = {}
for name, split in ds.items():
new_splits[name] = split.map(fix_row)
parent = os.path.dirname(hf_path.rstrip(os.sep))
tmp = tempfile.mkdtemp(prefix=".hf_dataset_rewrite_", dir=parent)
try:
DatasetDict(new_splits).save_to_disk(tmp)
shutil.rmtree(hf_path)
shutil.move(tmp, hf_path)
except Exception:
shutil.rmtree(tmp, ignore_errors=True)
raise
return True
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument(
"--processed_glob",
default="/projects/p32685/RL_dataset/lungx_*percent_processed",
help="Glob of processed dataset dirs",
)
ap.add_argument(
"--old_prefix",
default="/projects/p32591/RL_dataset",
help="Path prefix to replace",
)
ap.add_argument(
"--new_prefix",
default="/projects/p32685/RL_dataset",
help="Replacement prefix",
)
args = ap.parse_args()
import glob
dirs = sorted(glob.glob(args.processed_glob))
if not dirs:
print(f"No dirs matched: {args.processed_glob}", file=sys.stderr)
return 1
old, new = args.old_prefix, args.new_prefix
for d in dirs:
print(d)
fix_backup_entries(d, old, new)
print(" backup_entries: ok")
if fix_template_summary(d, old, new):
print(" template_summary.json: ok")
else:
print(" template_summary.json: skip")
if fix_hf_dataset(d, old, new):
print(" hf_dataset: ok")
else:
print(" hf_dataset: skip")
print("Done. Re-run Step 3 if needed.")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|