120 lines
4.6 KiB
Python
120 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
||
"""Generate 2-level terrain: level 0=flat, level 1=pyramid stairs.
|
||
|
||
1cm = 1px. OpenCV draws concentric filled rectangles (outside→in, 10 steps).
|
||
Higher values overwrite lower = convex pyramid (stairs up toward center).
|
||
Lower values overwrite higher = concave pyramid (stairs down from center).
|
||
|
||
Usage:
|
||
uv run scripts/gen_flat_stairs.py # convex (stairs UP)
|
||
uv run scripts/gen_flat_stairs.py --concave # concave (stairs DOWN)
|
||
uv run scripts/gen_flat_stairs.py --step-h 0.10 --num-steps 5
|
||
"""
|
||
import cv2, numpy as np, os, argparse
|
||
|
||
# ═══ fixed params ═══
|
||
HS = 0.01 # 1cm/px
|
||
VS = 0.005 # height unit = 0.5cm
|
||
CELL_M = 8.0 # 8m cell
|
||
BORDER_M = 5.0 # 5m border
|
||
NUM_ROWS = 2 # flat + stairs
|
||
NUM_COLS = 4 # columns
|
||
|
||
CELL_PX = int(CELL_M / HS) # 800
|
||
BORDER_PX = int(BORDER_M / HS) # 500
|
||
PLATFORM_PX = int(1.0 / HS) # 1m platform = 100px
|
||
|
||
TOT_ROWS = NUM_ROWS * CELL_PX + 2 * BORDER_PX
|
||
TOT_COLS = NUM_COLS * CELL_PX + 2 * BORDER_PX
|
||
|
||
# ═══ stairs params (override via CLI) ═══
|
||
NUM_STEPS = 10 # 10 steps
|
||
STEP_H_CM = 20 # 20cm rise per step
|
||
STEP_D_CM = 20 # 20cm tread per step
|
||
|
||
STEP_H_VS = int(STEP_H_CM / 100.0 / VS) # 0.20 / 0.005 = 40
|
||
STEP_D_PX = int(STEP_D_CM / 100.0 / HS) # 0.20 / 0.01 = 20
|
||
|
||
|
||
def draw_pyramid(canvas, x0, y0, num_steps, step_d_px, step_h_vs, concave=False):
|
||
"""Draw concentric rectangles from outside→in.
|
||
|
||
Convex: edge=0 → platform=max (stairs up toward center)
|
||
Concave: raise whole cell to max, then draw pit: edge=max → platform=0
|
||
"""
|
||
cx, cy = x0 + CELL_PX // 2, y0 + CELL_PX // 2
|
||
p2 = PLATFORM_PX // 2
|
||
h_max = step_h_vs * num_steps
|
||
|
||
# Fill cell to cell boundary with reference-plane height
|
||
half_max = CELL_PX // 2 # extend to cell edge
|
||
cv2.rectangle(canvas, (cx - half_max, cy - half_max),
|
||
(cx + half_max, cy + half_max), int(h_max), -1)
|
||
|
||
if concave:
|
||
# Pit: rings going DOWN from reference plane
|
||
for i in range(num_steps + 1):
|
||
half = p2 + (num_steps - i) * step_d_px
|
||
x1, y1 = cx - half, cy - half
|
||
x2, y2 = cx + half, cy + half
|
||
h = h_max - step_h_vs * i
|
||
cv2.rectangle(canvas, (x1, y1), (x2, y2), int(h), -1)
|
||
else:
|
||
# Mound: rings going UP from reference plane
|
||
for i in range(num_steps + 1):
|
||
half = p2 + (num_steps - i) * step_d_px
|
||
x1, y1 = cx - half, cy - half
|
||
x2, y2 = cx + half, cy + half
|
||
h = h_max + step_h_vs * i
|
||
cv2.rectangle(canvas, (x1, y1), (x2, y2), int(h), -1)
|
||
|
||
|
||
def main():
|
||
p = argparse.ArgumentParser()
|
||
p.add_argument("--concave", action="store_true", help="concave pyramid (stairs down from center)")
|
||
p.add_argument("--step-h", type=float, default=0.20, help="step rise (m)")
|
||
p.add_argument("--step-d", type=float, default=0.20, help="step tread (m)")
|
||
p.add_argument("--num-steps", type=int, default=10, help="number of steps")
|
||
args = p.parse_args()
|
||
|
||
step_h_vs = int(args.step_h / VS)
|
||
step_d_px = int(args.step_d / HS)
|
||
total_h = step_h_vs * args.num_steps * VS
|
||
|
||
print(f"Building {TOT_COLS}×{TOT_ROWS}px ({TOT_COLS*HS:.0f}×{TOT_ROWS*HS:.0f}m)")
|
||
print(f" type={'concave' if args.concave else 'convex'} "
|
||
f"steps={args.num_steps} rise={step_h_vs*VS*100:.0f}cm "
|
||
f"tread={step_d_px*HS*100:.0f}cm total_h={total_h*100:.0f}cm")
|
||
|
||
canvas = np.zeros((TOT_ROWS, TOT_COLS), dtype=np.uint16)
|
||
|
||
for row in range(NUM_ROWS):
|
||
for col in range(NUM_COLS):
|
||
x0 = BORDER_PX + col * CELL_PX
|
||
y0 = BORDER_PX + row * CELL_PX
|
||
if row == 1:
|
||
# alternate convex/concave across cols
|
||
concave_cell = (col % 2 == 1)
|
||
draw_pyramid(canvas, x0, y0, args.num_steps,
|
||
step_d_px, step_h_vs, concave_cell)
|
||
|
||
hf_m = canvas.astype(np.float32) * VS
|
||
z_min, z_max = float(hf_m.min()), float(hf_m.max())
|
||
z_range = max(z_max - z_min, 0.001)
|
||
|
||
png = ((hf_m - z_min) / z_range * 65535).astype(np.uint16)
|
||
|
||
out_dir = os.path.join(os.path.dirname(__file__), "..",
|
||
"motrix_envs", "src", "motrix_envs", "locomotion",
|
||
"go1", "xmls", "assets")
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
out_path = os.path.join(out_dir, "flat_stairs.png")
|
||
cv2.imwrite(out_path, png)
|
||
print(f" saved: {out_path}")
|
||
print(f" XML: size=\"{TOT_COLS*HS/2:.1f} {TOT_ROWS*HS/2:.1f} "
|
||
f"{z_range:.3f} {max(z_min,0.001):.3f}\"")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|