You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
95 lines
3.2 KiB
95 lines
3.2 KiB
import os
|
|
import sys
|
|
import numpy as np
|
|
import tqdm
|
|
from pathlib import Path
|
|
|
|
# Add project root and ISOLATE paths
|
|
project_root = Path(__file__).parent.parent
|
|
sys.path.append(str(project_root))
|
|
sys.path.append(str(project_root / 'scripts' / 'ISOLATE'))
|
|
|
|
# Import the model wrapper
|
|
try:
|
|
from scripts.ISOLATE.model_wrapper import ShenronRadarModel
|
|
except ImportError as e:
|
|
print(f"Error: Failed to import ShenronRadarModel. Ensure scripts/ISOLATE/model_wrapper.py exists. ({e})")
|
|
sys.exit(1)
|
|
|
|
def process_session(session_path):
|
|
print(f"\n>>> Processing session: {session_path.name}")
|
|
|
|
lidar_dir = session_path / "lidar"
|
|
if not lidar_dir.exists():
|
|
print(f" [SKIP] No 'lidar' folder found.")
|
|
return
|
|
|
|
# Find all .npy files in lidar/
|
|
lidar_files = sorted(list(lidar_dir.glob("*.npy")))
|
|
if not lidar_files:
|
|
print(f" [SKIP] No .npy files in 'lidar' folder.")
|
|
return
|
|
|
|
output_dir = session_path / "shenron_radar"
|
|
output_dir.mkdir(exist_ok=True)
|
|
|
|
# Initialize the model once per session
|
|
print(f" Initializing ShenronRadarModel...")
|
|
model = ShenronRadarModel(radar_type='radarbook')
|
|
|
|
print(f" Generating Shenron Radar data for {len(lidar_files)} frames...")
|
|
|
|
for lidar_file in tqdm.tqdm(lidar_files, desc=" Simulating Radar", unit="frame"):
|
|
try:
|
|
# 1. Load Semantic LiDAR data
|
|
# Expected raw: [x, y, z, cos, obj, tag] (6 cols)
|
|
# Expected Shenron input: [x, y, z, intensity, cos, obj, tag] (7 cols)
|
|
data = np.load(lidar_file)
|
|
|
|
if data.shape[1] == 6:
|
|
# Pad with a dummy intensity column at index 3
|
|
# This aligns 'tag' to index 6 as expected by our lidar.py mapping
|
|
padded_data = np.zeros((data.shape[0], 7), dtype=np.float32)
|
|
padded_data[:, 0:3] = data[:, 0:3] # x, y, z
|
|
padded_data[:, 4:7] = data[:, 3:6] # cos, obj, tag
|
|
data = padded_data
|
|
|
|
# 2. Process through the physics-based model
|
|
# returns rich PCD: [M, 5] (x, y, z, velocity, magnitude)
|
|
rich_pcd = model.process(data)
|
|
|
|
# 3. Save to disk
|
|
output_file = output_dir / lidar_file.name
|
|
np.save(output_file, rich_pcd)
|
|
|
|
except Exception as e:
|
|
print(f"\n [ERROR] Failed to process {lidar_file.name}: {e}")
|
|
|
|
def main():
|
|
data_root = project_root / "data"
|
|
if not data_root.exists():
|
|
print(f"Error: {data_root} not found.")
|
|
return
|
|
|
|
# Get all session folders
|
|
sessions = sorted([d for d in data_root.iterdir() if d.is_dir()])
|
|
|
|
if not sessions:
|
|
print("No simulation sessions found in data/.")
|
|
return
|
|
|
|
print(f"Found {len(sessions)} sessions.")
|
|
|
|
for session in sessions:
|
|
# Check if the session has frames.jsonl to confirm it's a valid data folder
|
|
if (session / "frames.jsonl").exists():
|
|
process_session(session)
|
|
else:
|
|
print(f"Skipping {session.name} (no frames.jsonl found).")
|
|
|
|
print("\n" + "="*50)
|
|
print("SHENRON BATCH PROCESSING COMPLETE!")
|
|
print("="*50)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|