PlantShade-code/2_build_split.py at main · DaRL-GenAI/PlantShade-code · GitHub

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
"""Assemble (source, target, prompt) triplets into train/test JSON lists.

  source : plant structure frame  <source_root>/<scene>/<scene>/<stage>/seg_<k>.png
  target : shadow map from extract_shadow.py  <target_root>/<scene>/<stage>/seg_<k>.png
  prompt : supplementary light position "x y z" for light index k

The light travels a fixed circle (radius 1.25 m -> 125, height z=280) sampled at
100 positions. Position k is derived from the trajectory used during capture:
    theta_k = 180 - (k-1) * 3.6  degrees ;  x = 125 cos(theta), y = 125 sin(theta)
(verified against the released filenames, e.g. seg_1 -> "-125.000 0.000 280.000").

Usage (from repo root):
    python build_split.py \
        --source_root data/PlantShadeUnzip \
        --target_root data/shadow/target \
        --out_dir     data \
        --test_ratio  0.2
"""
import os, json, glob, math, random, argparse

RADIUS, Z, N_POS, STEP_DEG, START_DEG = 125.0, 280.0, 100, 3.6, 180.0


def prompt_for(k):
    theta = math.radians(START_DEG - (k - 1) * STEP_DEG)
    return f"{RADIUS*math.cos(theta):.3f} {RADIUS*math.sin(theta):.3f} {Z:.3f}"


def scene_source_root(source_root, scene):
    nested = os.path.join(source_root, scene, scene)
    return nested if os.path.isdir(nested) else os.path.join(source_root, scene)


def collect(source_root, target_root):
    triplets = []
    scenes = sorted(d for d in os.listdir(target_root)
                    if os.path.isdir(os.path.join(target_root, d)))
    for scene in scenes:
        src_root = scene_source_root(source_root, scene)
        tgt_scene = os.path.join(target_root, scene)
        for stage in sorted(os.listdir(tgt_scene)):
            tgt_stage = os.path.join(tgt_scene, stage)
            if not os.path.isdir(tgt_stage):
                continue
            for tgt in glob.glob(os.path.join(tgt_stage, "seg_*.png")):
                fname = os.path.basename(tgt)
                try:
                    k = int(fname.replace("seg_", "").replace(".png", ""))
                except ValueError:
                    continue
                src = os.path.join(src_root, stage, fname)
                if os.path.exists(src):
                    triplets.append({"source": os.path.abspath(src),
                                     "target": os.path.abspath(tgt),
                                     "prompt": prompt_for(k)})
    return triplets


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--source_root", required=True)
    ap.add_argument("--target_root", required=True)
    ap.add_argument("--out_dir", default="data")
    ap.add_argument("--test_ratio", type=float, default=0.2)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    triplets = collect(args.source_root, args.target_root)
    random.seed(args.seed)
    random.shuffle(triplets)
    n_test = int(len(triplets) * args.test_ratio)
    test, train = triplets[:n_test], triplets[n_test:]

    os.makedirs(args.out_dir, exist_ok=True)
    for name, rows in [("train_split.json", train), ("test_split.json", test)]:
        with open(os.path.join(args.out_dir, name), "w") as f:
            for r in rows:
                f.write(json.dumps(r) + "\n")
    print(f"total={len(triplets)}  train={len(train)}  test={len(test)}  -> {args.out_dir}/")


if __name__ == "__main__":
    main()