diff --git a/src/camera-based-e2e/Point_cloud/create_occ_labels.py b/src/camera-based-e2e/Point_cloud/create_occ_labels.py new file mode 100644 index 0000000..83996b3 --- /dev/null +++ b/src/camera-based-e2e/Point_cloud/create_occ_labels.py @@ -0,0 +1,485 @@ +"""Offline occupancy pseudo-label generation. +For each frame in the Waymo E2E dataset this script: + 1. Decodes all camera images. + 2. Runs SegFormer-B2 to produce per-pixel semantic masks. + 3. Runs Depth-Anything-V2 to get metric depth maps. + 4. Back-projects depth -> 3D points in vehicle frame. + 5. Assigns semantic class to each point from the segmentation mask. + 6. Filters points that appear in fewer than min_views cameras (multiview filter). + 7. Ray-casts from each camera origin to mark free space (numba JIT). + 8. Voxelises into a (100, 100, 16) uint8 occupancy grid with majority voting. + +Semantic class map (Cityscapes → our 6 classes): + 0 free (confirmed empty along camera rays) + 1 vehicle + 2 pedestrian + 3 cyclist + 4 road + 5 static + 255 unknown (unobserved — ignored in loss) +""" + +import argparse +import os +import pickle +from pathlib import Path + +import cv2 +import numba +import numpy as np +import torch +import torch.nn.functional as F +from tqdm import tqdm +from transformers import ( + AutoImageProcessor, + AutoModelForDepthEstimation, + SegformerForSemanticSegmentation, +) + +from protos import e2e_pb2 +from point_cloud_gpu import ( + BEV_RANGE, + backproject_depth_to_points, + points_to_vehicle_frame, + undistort_image, +) + +# voxel grid config +VOX_XY_RANGE = 25.0 +VOX_XY_RES = 0.5 +VOX_Z_MIN = -3.0 +VOX_Z_MAX = 5.0 +VOX_Z_RES = 0.5 +VOX_XY_SIZE = int(2 * VOX_XY_RANGE / VOX_XY_RES) # 100 +VOX_Z_SIZE = int((VOX_Z_MAX - VOX_Z_MIN) / VOX_Z_RES) # 16 +NUM_CLASSES = 6 + +RAY_SUBSAMPLE = 8 + +# cityscapes -> our map +CITYSCAPES_TO_OCC = np.full(256, 5, dtype=np.uint8) +CITYSCAPES_TO_OCC[0] = 4 # road +CITYSCAPES_TO_OCC[1] = 5 # sidewalk +CITYSCAPES_TO_OCC[2] = 5 # building +CITYSCAPES_TO_OCC[3] = 5 # wall +CITYSCAPES_TO_OCC[4] = 5 # fence +CITYSCAPES_TO_OCC[5] = 5 # pole +CITYSCAPES_TO_OCC[6] = 5 # traffic light +CITYSCAPES_TO_OCC[7] = 5 # traffic sign +CITYSCAPES_TO_OCC[8] = 5 # vegetation +CITYSCAPES_TO_OCC[9] = 5 # terrain +CITYSCAPES_TO_OCC[10] = 255 # sky +CITYSCAPES_TO_OCC[11] = 2 # person +CITYSCAPES_TO_OCC[12] = 3 # rider +CITYSCAPES_TO_OCC[13] = 1 # car +CITYSCAPES_TO_OCC[14] = 1 # truck +CITYSCAPES_TO_OCC[15] = 1 # bus +CITYSCAPES_TO_OCC[16] = 1 # train +CITYSCAPES_TO_OCC[17] = 3 # motorcycle +CITYSCAPES_TO_OCC[18] = 3 # bicycle +CITYSCAPES_TO_OCC[255] = 0 # ignore → free + +# model loading + +def load_seg_model(model_id: str, device: torch.device): + print(f"Loading segmentation model: {model_id}") + processor = AutoImageProcessor.from_pretrained(model_id) + model = SegformerForSemanticSegmentation.from_pretrained(model_id).to(device) + model.eval() + if device.type == "cuda": + model = model.half() + return model, processor + + +def load_depth_model(device: torch.device): + depth_id = "depth-anything/Depth-Anything-V2-Small-hf" + print(f"Loading depth model: {depth_id}") + processor = AutoImageProcessor.from_pretrained(depth_id) + model = AutoModelForDepthEstimation.from_pretrained(depth_id).to(device) + model.eval() + if device.type == "cuda": + model = model.half() + return model, processor + +# segmentation + +def run_segmentation_batch(rgb_images, seg_model, seg_processor, device): + inputs = seg_processor(images=rgb_images, return_tensors="pt") + inputs = { + k: v.half().to(device) if v.dtype == torch.float32 else v.to(device) + for k, v in inputs.items() + } + with torch.no_grad(): + logits = seg_model(**inputs).logits + masks = [] + for i, img in enumerate(rgb_images): + orig_h, orig_w = img.shape[:2] + pred = logits[i:i+1].float() + pred_up = F.interpolate(pred, size=(orig_h, orig_w), mode="bilinear", align_corners=False) + class_map = pred_up.argmax(dim=1).squeeze().cpu().numpy().astype(np.uint8) + masks.append(CITYSCAPES_TO_OCC[class_map]) + return masks + +# depth inference + +def run_depth_batch(rgb_images, depth_model, depth_processor, device, max_depth=30.0): + """Handles cameras with different resolutions by padding to a common size.""" + original_sizes = [(img.shape[0], img.shape[1]) for img in rgb_images] + max_h = max(s[0] for s in original_sizes) + max_w = max(s[1] for s in original_sizes) + padded = [] + for img in rgb_images: + h, w = img.shape[:2] + p = cv2.copyMakeBorder(img, 0, max_h - h, 0, max_w - w, cv2.BORDER_REPLICATE) + padded.append(p) + inputs = depth_processor(images=padded, return_tensors="pt").to(device) + with torch.no_grad(): + preds = depth_model(**inputs).predicted_depth + depth_maps = [] + for i, (orig_h, orig_w) in enumerate(original_sizes): + pred_i = preds[i:i+1].unsqueeze(1).float() + pred_full = F.interpolate( + pred_i, size=(max_h, max_w), mode="bicubic", align_corners=False + ).squeeze() + depth_maps.append(pred_full[:orig_h, :orig_w].clamp(0.1, max_depth)) + return depth_maps + +# multiview filter (vectorized) + +def multiview_filter(all_pts, all_labels, min_views: int = 2): + """Keep only points whose voxel is observed by at least min_views cameras. + + Uses fully vectorized numpy ops — no Python loops over points. + Points from dynamic classes (vehicle/ped/cyclist) skip the filter + since they're unlikely to overlap across cameras anyway. + """ + num_vox = VOX_XY_SIZE * VOX_XY_SIZE * VOX_Z_SIZE + + # compute valid voxel flat index for every point from every camera + cam_vox_keys = [] # flat index per point (-1 = out of bounds) + cam_valid_masks = [] # which points are in-bounds + for pts in all_pts: + pts_np = pts.cpu().numpy() + ix = np.floor((VOX_XY_RANGE - pts_np[:, 0]) / VOX_XY_RES).astype(np.int32) + iy = np.floor((VOX_XY_RANGE - pts_np[:, 1]) / VOX_XY_RES).astype(np.int32) + iz = np.floor((pts_np[:, 2] - VOX_Z_MIN) / VOX_Z_RES).astype(np.int32) + in_bounds = ( + (ix >= 0) & (ix < VOX_XY_SIZE) & + (iy >= 0) & (iy < VOX_XY_SIZE) & + (iz >= 0) & (iz < VOX_Z_SIZE) + ) + flat = np.where(in_bounds, + ix * VOX_XY_SIZE * VOX_Z_SIZE + iy * VOX_Z_SIZE + iz, + -1) + cam_vox_keys.append(flat) + cam_valid_masks.append(in_bounds) + + # count how many distinct cameras observe each voxel + view_counts = np.zeros(num_vox, dtype=np.int8) + for flat in cam_vox_keys: + valid_flat = flat[flat >= 0] + unique_vox = np.unique(valid_flat) + view_counts[unique_vox] += 1 + + # Fflter each camera's points — dynamic objects bypass the view count check + filtered_pts, filtered_labels = [], [] + for pts, labels, flat, in_bounds in zip(all_pts, all_labels, cam_vox_keys, cam_valid_masks): + lbl_np = labels.cpu().numpy() + is_dynamic = (lbl_np >= 1) & (lbl_np <= 3) + # out-of-bounds points get view_count=0, so they're dropped for static + seen_enough = np.zeros(len(flat), dtype=bool) + valid_idx = np.where(in_bounds)[0] + seen_enough[valid_idx] = view_counts[flat[valid_idx]] >= min_views + keep = is_dynamic | seen_enough + filtered_pts.append(pts[torch.from_numpy(keep).to(pts.device)]) + filtered_labels.append(labels[torch.from_numpy(keep).to(labels.device)]) + + return filtered_pts, filtered_labels + +# DDA ray casting - marks intermediate voxels as free, modifies voxels marked unknown + +@numba.njit +def _dda_ray(ox: float, oy: float, oz: float, ex: float, ey: float, ez: float, grid: np.ndarray) -> None: + dx = ex - ox + dy = ey - oy + dz = ez - oz + n = int(max(abs(dx), abs(dy), abs(dz))) + if n < 2 or n > 90: # 90 voxels = 45m + return + step_x = dx / n + step_y = dy / n + step_z = dz / n + cx = float(ox) + cy = float(oy) + cz = float(oz) + for _ in range(n - 1): + cx += step_x + cy += step_y + cz += step_z + vx = int(cx) + vy = int(cy) + vz = int(cz) + if 0 <= vx < 100 and 0 <= vy < 100 and 0 <= vz < 16: + if grid[vx, vy, vz] == 255: + grid[vx, vy, vz] = 0 + + +def mark_free_space(grid: np.ndarray, cam_origin: np.ndarray, + pts_veh: np.ndarray) -> np.ndarray: + if len(pts_veh) == 0: + return grid + ox = (VOX_XY_RANGE - cam_origin[0]) / VOX_XY_RES + oy = (VOX_XY_RANGE - cam_origin[1]) / VOX_XY_RES + oz = (cam_origin[2] - VOX_Z_MIN) / VOX_Z_RES + pts_sub = pts_veh[::RAY_SUBSAMPLE] + px = (VOX_XY_RANGE - pts_sub[:, 0]) / VOX_XY_RES + py = (VOX_XY_RANGE - pts_sub[:, 1]) / VOX_XY_RES + pz = (pts_sub[:, 2] - VOX_Z_MIN) / VOX_Z_RES + for i in range(len(px)): + _dda_ray(ox, oy, oz, px[i], py[i], pz[i], grid) + return grid + +# voxelisation with majority voting + +def voxelise_labelled_cloud(pts: torch.Tensor, labels: torch.Tensor) -> np.ndarray: + x = pts[:, 0].cpu().numpy() + y = pts[:, 1].cpu().numpy() + z = pts[:, 2].cpu().numpy() + lbl = labels.cpu().numpy() + valid_lbl = lbl < NUM_CLASSES + x, y, z, lbl = x[valid_lbl], y[valid_lbl], z[valid_lbl], lbl[valid_lbl] + ix = ((VOX_XY_RANGE - x) / VOX_XY_RES).astype(np.int32) + iy = ((VOX_XY_RANGE - y) / VOX_XY_RES).astype(np.int32) + iz = ((z - VOX_Z_MIN) / VOX_Z_RES).astype(np.int32) + valid = ( + (ix >= 0) & (ix < VOX_XY_SIZE) & + (iy >= 0) & (iy < VOX_XY_SIZE) & + (iz >= 0) & (iz < VOX_Z_SIZE) + ) + ix, iy, iz, lbl = ix[valid], iy[valid], iz[valid], lbl[valid] + flat_idx = ix * VOX_XY_SIZE * VOX_Z_SIZE + iy * VOX_Z_SIZE + iz + num_vox = VOX_XY_SIZE * VOX_XY_SIZE * VOX_Z_SIZE + counts = np.zeros((num_vox, NUM_CLASSES), dtype=np.int32) + np.add.at(counts, (flat_idx, lbl.astype(np.int32)), 1) + counts_sum = counts.sum(axis=1) + winner = counts.argmax(axis=1).astype(np.uint8) + occupied = counts_sum > 0 + + is_static_road = np.isin(winner, [4, 5]) + strong = counts_sum >= 1 + keep = occupied & (~is_static_road | strong) + + flat_grid = np.full(num_vox, 255, dtype=np.uint8) + flat_grid[keep] = winner[keep] + return flat_grid.reshape(VOX_XY_SIZE, VOX_XY_SIZE, VOX_Z_SIZE) + +# per-frame processing + +def process_frame(frame, seg_model, seg_processor, depth_model, depth_processor, + device, max_depth=30.0, min_views=1): + """Full pipeline for one frame -> (100, 100, 16) uint8 occupancy grid.""" + + calib_by_name = {} + for calib in frame.frame.context.camera_calibrations: + intr = list(calib.intrinsic) + extr = np.array(list(calib.extrinsic.transform), dtype=np.float64).reshape(4, 4) + calib_by_name[calib.name] = { + "intrinsic": intr, + "extrinsic": extr, + "width": calib.width, + "height": calib.height, + } + + rgb_images, intrinsics_list, extrinsics_list = [], [], [] + for img_proto in frame.frame.images: + cam_name = img_proto.name + if cam_name not in calib_by_name: + continue + cal = calib_by_name[cam_name] + jpg = np.frombuffer(img_proto.image, dtype=np.uint8) + rgb = cv2.imdecode(jpg, cv2.IMREAD_COLOR) + if rgb is None: + continue + rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB) + intr = cal["intrinsic"] + intrinsics_list.append( + [intr[0], intr[1], intr[2], intr[3], cal["width"], cal["height"]] + intr[4:9] + ) + extrinsics_list.append(cal["extrinsic"]) + rgb_images.append(rgb) + + if not rgb_images: + return np.full((VOX_XY_SIZE, VOX_XY_SIZE, VOX_Z_SIZE), 255, dtype=np.uint8) + + seg_maps = run_segmentation_batch(rgb_images, seg_model, seg_processor, device) + depth_maps_gpu = run_depth_batch(rgb_images, depth_model, depth_processor, device, max_depth) + + all_pts, all_labels, cam_origins = [], [], [] + + for rgb, depth_gpu, intr_vals, extr_np, seg_mask in zip( + rgb_images, depth_maps_gpu, intrinsics_list, extrinsics_list, seg_maps + ): + fx, fy, cx, cy = intr_vals[0], intr_vals[1], intr_vals[2], intr_vals[3] + + if len(intr_vals) > 6: + dist = intr_vals[6:11] + rgb = undistort_image(rgb, fx, fy, cx, cy, dist) + depth_np = depth_gpu.cpu().numpy() + depth_np = undistort_image(depth_np, fx, fy, cx, cy, dist) + depth_gpu = torch.from_numpy(depth_np).to(device) + seg_mask = undistort_image(seg_mask, fx, fy, cx, cy, dist) + + pts_cv = backproject_depth_to_points(depth_gpu, fx, fy, cx, cy) + d_flat = depth_gpu.reshape(-1) + valid = (d_flat > 0.1) & (d_flat < max_depth) + pts_cv = pts_cv[valid] + + extr_gpu = torch.from_numpy(extr_np.astype(np.float32)).to(device) + pts_veh = points_to_vehicle_frame(pts_cv, extr_gpu, device) + + x_, y_, z_ = pts_veh[:, 0], pts_veh[:, 1], pts_veh[:, 2] + keep = ( + (x_.abs() < VOX_XY_RANGE) & (y_.abs() < VOX_XY_RANGE) & + (z_ > VOX_Z_MIN - 0.5) & (z_ < VOX_Z_MAX + 0.5) + ) + pts_veh = pts_veh[keep] + + valid_idx = valid.nonzero(as_tuple=True)[0] + keep_idx = keep.nonzero(as_tuple=True)[0] + depth_vals = d_flat[valid_idx][keep_idx] + + H, W = depth_gpu.shape + v_grid, u_grid = torch.meshgrid( + torch.arange(H, device=device, dtype=torch.float32), + torch.arange(W, device=device, dtype=torch.float32), + indexing="ij", + ) + u_flat = u_grid.reshape(-1)[valid_idx][keep_idx].long().clamp(0, W - 1) + v_flat = v_grid.reshape(-1)[valid_idx][keep_idx].long().clamp(0, H - 1) + seg_t = torch.from_numpy(seg_mask.astype(np.int64)).to(device) + point_labels = seg_t[v_flat, u_flat].to(torch.uint8) + + # remove sky points + sky_mask = point_labels != 255 + pts_veh = pts_veh[sky_mask] + point_labels = point_labels[sky_mask] + depth_vals = depth_vals[sky_mask] + + # class-aware depth filter + is_dynamic = (point_labels >= 1) & (point_labels <= 3) + # depth_keep = (is_dynamic | (depth_vals < 15.0)) & (depth_vals < 25.0) + depth_keep = depth_vals < 25.0 + pts_veh = pts_veh[depth_keep] + # print("raw pts:", pts_veh.shape[0]) + point_labels = point_labels[depth_keep] + + all_pts.append(pts_veh) + all_labels.append(point_labels) + cam_origins.append(extr_np[:3, 3]) + + if not all_pts: + return np.full((VOX_XY_SIZE, VOX_XY_SIZE, VOX_Z_SIZE), 255, dtype=np.uint8) + + # multiview filter — remove static noise seen by only one camera + all_pts, all_labels = multiview_filter(all_pts, all_labels, min_views=min_views) + + fused_pts = torch.cat(all_pts, dim=0) + fused_labels = torch.cat(all_labels, dim=0) + + # ray cast on blank gris -> free space + # Only use static + road points for ray casting + grid = np.full((VOX_XY_SIZE, VOX_XY_SIZE, VOX_Z_SIZE), 255, dtype=np.uint8) + for i, cam_origin in enumerate(cam_origins): + lbl_np = all_labels[i].cpu().numpy() + static_mask = lbl_np >= 4 # road (4) and static (5) only + ray_pts = all_pts[i][torch.from_numpy(static_mask).to(all_pts[i].device)] + grid = mark_free_space(grid, cam_origin, ray_pts.cpu().numpy()) + + # overlay semantic labels with majority voting + semantic_grid = voxelise_labelled_cloud(fused_pts, fused_labels) + occupied_mask = semantic_grid != 255 + grid[occupied_mask] = semantic_grid[occupied_mask] + return grid + +def main(): + parser = argparse.ArgumentParser(description="Generate occupancy pseudo-labels") + parser.add_argument("--data_dir", type=str, required=True) + parser.add_argument("--split", type=str, default="train", + choices=["train", "val", "test"]) + parser.add_argument("--index_dir", type=str, default=".") + parser.add_argument("--start_idx", type=int, default=0) + parser.add_argument("--n_items", type=int, default=None) + parser.add_argument("--seg_model", type=str, + default="nvidia/segformer-b2-finetuned-cityscapes-1024-1024") + parser.add_argument("--min_views", type=int, default=2, + help="Min cameras that must observe a static voxel to keep it") + parser.add_argument("--device", type=str, + default="cuda" if torch.cuda.is_available() else "cpu") + parser.add_argument("--output_dir", type=str, default=None) + args = parser.parse_args() + + device = torch.device(args.device) + + occ_base = Path(args.output_dir) if args.output_dir else Path(args.data_dir).parent / "occ" + occ_split_dir = occ_base / args.split + occ_split_dir.mkdir(parents=True, exist_ok=True) + + index_file = os.path.join(args.index_dir, f"index_{args.split}.pkl") + with open(index_file, "rb") as f: + indexes = pickle.load(f) + + end_idx = len(indexes) + if args.n_items is not None: + end_idx = min(args.start_idx + args.n_items, end_idx) + print(f"Processing frames {args.start_idx}–{end_idx - 1}") + + seg_model, seg_processor = load_seg_model(args.seg_model, device) + depth_model, depth_processor = load_depth_model(device) + + _dummy = np.full((100, 100, 16), 255, dtype=np.uint8) + _dda_ray(50.0, 50.0, 6.0, 70.0, 50.0, 6.0, _dummy) + print("Numba JIT warmed up.") + + occ_index = [] + open_file, open_filename = None, "" + + for idx in tqdm(range(args.start_idx, end_idx), desc=f"OccLabels [{args.split}]"): + occ_path = occ_split_dir / f"occ_{idx:07d}.npy" + + if occ_path.exists(): + occ_index.append((idx, str(occ_path))) + continue + + filename, start_byte, byte_length = indexes[idx] + if open_filename != filename: + if open_file is not None: + open_file.close() + open_file = open(os.path.join(args.data_dir, filename), "rb") + open_filename = filename + + open_file.seek(start_byte) + frame = e2e_pb2.E2EDFrame() + frame.ParseFromString(open_file.read(byte_length)) + + occ_grid = process_frame( + frame, seg_model, seg_processor, depth_model, depth_processor, + device, min_views=args.min_views + ) + + np.save(str(occ_path), occ_grid) + occ_index.append((idx, str(occ_path))) + + if open_file is not None: + open_file.close() + + index_path = occ_base / f"occ_index_{args.split}.pkl" + with open(str(index_path), "wb") as f: + pickle.dump(occ_index, f) + + print(f"\nSaved {len(occ_index)} occupancy grids to {occ_split_dir}") + print(f"Index written to {index_path}") + + +if __name__ == "__main__": + main() + diff --git a/src/camera-based-e2e/Point_cloud/create_occ_train.slurm b/src/camera-based-e2e/Point_cloud/create_occ_train.slurm new file mode 100644 index 0000000..cb8b864 --- /dev/null +++ b/src/camera-based-e2e/Point_cloud/create_occ_train.slurm @@ -0,0 +1,36 @@ +#!/bin/bash +#SBATCH --job-name=waymo_occ_test +#SBATCH --output=logs/%x_%A_%a.out +#SBATCH --error=logs/%x_%A_%a.err +#SBATCH --partition=a10 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=40G +#SBATCH --time=72:00:00 +#SBATCH --account=csso +#SBATCH --array=0-4 + +PYTHON=/scratch/gilbreth/kumar753/conda_envs/robo_env_310_new/bin/python +SCRIPT_DIR=/scratch/gilbreth/kumar753/robotvision/robotvision/src/camera-based-e2e/Point_cloud +PARENT_DIR=/scratch/gilbreth/kumar753/robotvision/robotvision/src/camera-based-e2e +DATA_DIR=/scratch/gilbreth/kumar753/robotvision/waymo_end_to_end_camera_v1_0_0/waymo_open_dataset_end_to_end_camera_v_1_0_0 + +export PYTHONPATH="${PARENT_DIR}:${SCRIPT_DIR}:${PYTHONPATH}" +mkdir -p ${SCRIPT_DIR}/logs + +TOTAL=50000 +NUM_TASKS=5 +CHUNK=$(( (TOTAL + NUM_TASKS - 1) / NUM_TASKS )) +START=$(( SLURM_ARRAY_TASK_ID * CHUNK )) + +echo "OCC train_new task ${SLURM_ARRAY_TASK_ID}/${NUM_TASKS}: frames ${START}..${START}+${CHUNK}" + +$PYTHON ${SCRIPT_DIR}/create_occ_labels.py \ + --data_dir "$DATA_DIR" \ + --split train \ + --index_dir "${PARENT_DIR}" \ + --start_idx "$START" \ + --n_items "$CHUNK" \ + --output_dir /scratch/gilbreth/kumar753/waymo_occ_new \ + --device cuda + diff --git a/src/camera-based-e2e/Point_cloud/create_occ_val.slurm b/src/camera-based-e2e/Point_cloud/create_occ_val.slurm new file mode 100644 index 0000000..a0f6f9e --- /dev/null +++ b/src/camera-based-e2e/Point_cloud/create_occ_val.slurm @@ -0,0 +1,36 @@ +#!/bin/bash +#SBATCH --job-name=waymo_occ_val +#SBATCH --output=logs/%x_%A_%a.out +#SBATCH --error=logs/%x_%A_%a.err +#SBATCH --partition=a10 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=40G +#SBATCH --time=24:00:00 +#SBATCH --account=csso +#SBATCH --array=0-1 + +PYTHON=/scratch/gilbreth/kumar753/conda_envs/robo_env_310_new/bin/python +SCRIPT_DIR=/scratch/gilbreth/kumar753/robotvision/robotvision/src/camera-based-e2e/Point_cloud +PARENT_DIR=/scratch/gilbreth/kumar753/robotvision/robotvision/src/camera-based-e2e +DATA_DIR=/scratch/gilbreth/kumar753/robotvision/waymo_end_to_end_camera_v1_0_0/waymo_open_dataset_end_to_end_camera_v_1_0_0 + +export PYTHONPATH="${PARENT_DIR}:${SCRIPT_DIR}:${PYTHONPATH}" +mkdir -p ${SCRIPT_DIR}/logs + +TOTAL=10000 +NUM_TASKS=2 +CHUNK=$(( (TOTAL + NUM_TASKS - 1) / NUM_TASKS )) +START=$(( SLURM_ARRAY_TASK_ID * CHUNK )) + +echo "OCC val_new task ${SLURM_ARRAY_TASK_ID}/${NUM_TASKS}: frames ${START}..${START}+${CHUNK}" + +$PYTHON ${SCRIPT_DIR}/create_occ_labels.py \ + --data_dir "$DATA_DIR" \ + --split val \ + --index_dir "${PARENT_DIR}" \ + --start_idx "$START" \ + --n_items "$CHUNK" \ + --output_dir /scratch/gilbreth/kumar753/waymo_occ_new \ + --device cuda + diff --git a/src/camera-based-e2e/Point_cloud/vim_occ.py b/src/camera-based-e2e/Point_cloud/vim_occ.py new file mode 100644 index 0000000..3731927 --- /dev/null +++ b/src/camera-based-e2e/Point_cloud/vim_occ.py @@ -0,0 +1,158 @@ +"""Quick sanity-check visualizer for occupancy grids. + +Loads a few occ_{idx}.npy files and shows: + 1. Top-down 2D slice (XY at ground level) — colored by class + 2. Side view (XZ slice through center) + 3. Class distribution bar chart + +Does NOT require Open3D. Uses matplotlib only. + +Usage (on login node — no GPU needed): + python viz_occ.py \ + --occ_dir /scratch/gilbreth/kumar753/waymo_occ/train \ + --indices 0 100 500 1000 5000 + +Classes: + 0 = free (white) + 1 = vehicle (red) + 2 = pedestrian (blue) + 3 = cyclist (orange) + 4 = road (gray) + 5 = static (green) + 255 = unknown (black) +""" + +import argparse +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches +import os + +# Class colors (RGB, normalized) +CLASS_COLORS = { + 255: (0.0, 0.0, 0.0), # unknown — black + 0: (1.0, 1.0, 1.0), # free — white + 1: (0.9, 0.1, 0.1), # vehicle — red + 2: (0.1, 0.3, 0.9), # pedestrian — blue + 3: (1.0, 0.6, 0.0), # cyclist — orange + 4: (0.5, 0.5, 0.5), # road — gray + 5: (0.2, 0.7, 0.2), # static — green +} +CLASS_NAMES = {255: "unknown", 0: "free", 1: "vehicle", 2: "pedestrian", + 3: "cyclist", 4: "road", 5: "static"} + +VOX_XY_SIZE = 100 +VOX_Z_SIZE = 16 +VOX_Z_MIN = -3.0 +VOX_Z_RES = 0.5 + + +def grid_to_rgb(slice_2d): + """Convert a 2D (H, W) uint8 class grid to (H, W, 3) RGB image.""" + h, w = slice_2d.shape + rgb = np.zeros((h, w, 3), dtype=np.float32) + for cls, color in CLASS_COLORS.items(): + mask = slice_2d == cls + rgb[mask] = color + return rgb + + +def ground_level_z(): + """Return the voxel Z index closest to Z=0 (ground level).""" + return int((0.0 - VOX_Z_MIN) / VOX_Z_RES) + + +def visualize_occ(occ_path, ax_row): + """Visualize one occupancy grid on a row of 3 axes.""" + grid = np.load(occ_path) # (100, 100, 16) + idx = os.path.basename(occ_path).replace("occ_", "").replace(".npy", "") + + # ── Slice 1: top-down XY at ground level ── + z_ground = ground_level_z() + xy_slice = grid[:, :, z_ground] # (100, 100) + ax_row[0].imshow(grid_to_rgb(xy_slice), origin="upper") + ax_row[0].set_title(f"idx={idx} XY@Z=0m", fontsize=9) + ax_row[0].axis("off") + + # ── Slice 2: top-down XY — max occupied class along Z ── + # For each XY cell take the most common non-unknown, non-free class + best = np.full((VOX_XY_SIZE, VOX_XY_SIZE), 255, dtype=np.uint8) + for z in range(VOX_Z_SIZE): + layer = grid[:, :, z] + occupied = (layer >= 1) & (layer <= 5) + best[occupied] = layer[occupied] + ax_row[1].imshow(grid_to_rgb(best), origin="upper") + ax_row[1].set_title(f"idx={idx} XY max-Z", fontsize=9) + ax_row[1].axis("off") + + # ── Slice 3: class distribution ── + classes = [0, 1, 2, 3, 4, 5, 255] + counts = [int((grid == c).sum()) for c in classes] + colors = [CLASS_COLORS[c] for c in classes] + labels = [CLASS_NAMES[c] for c in classes] + bars = ax_row[2].bar(labels, counts, color=colors, edgecolor="black", linewidth=0.5) + ax_row[2].set_title(f"idx={idx} voxel counts", fontsize=9) + ax_row[2].tick_params(axis="x", labelsize=7, rotation=30) + ax_row[2].tick_params(axis="y", labelsize=7) + + # Print summary + total = grid.size + known = (grid != 255).sum() + free = (grid == 0).sum() + occupied_sum = sum((grid == c).sum() for c in range(1, 6)) + print(f" idx={idx}: known={known/total:.1%} free={free/total:.1%} " + f"occupied={occupied_sum/total:.1%} unknown={((grid==255).sum())/total:.1%}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--occ_dir", type=str, required=True, + help="Directory containing occ_{idx}.npy files") + parser.add_argument("--indices", type=int, nargs="+", default=[0, 100, 500, 1000, 5000], + help="Frame indices to visualize") + parser.add_argument("--out", type=str, default="occ_viz.png", + help="Output image path") + args = parser.parse_args() + + # Filter to indices that actually exist + paths = [] + for idx in args.indices: + p = os.path.join(args.occ_dir, f"occ_{idx:07d}.npy") + if os.path.exists(p): + paths.append(p) + else: + print(f" Warning: {p} not found, skipping") + + if not paths: + print("No valid occ files found. Check --occ_dir and --indices.") + return + + n = len(paths) + fig, axes = plt.subplots(n, 3, figsize=(12, 4 * n)) + if n == 1: + axes = [axes] + + print(f"\nVoxel statistics:") + for i, path in enumerate(paths): + visualize_occ(path, axes[i]) + + # Legend + legend_patches = [ + mpatches.Patch(color=CLASS_COLORS[c], label=CLASS_NAMES[c], linewidth=0.5, + edgecolor="black") + for c in [255, 0, 1, 2, 3, 4, 5] + ] + fig.legend(handles=legend_patches, loc="lower center", ncol=7, + fontsize=8, bbox_to_anchor=(0.5, 0.0)) + + plt.suptitle("Occupancy grid sanity check\n" + "Left: XY slice at Z=0m | Middle: XY max-Z projection | Right: class counts", + fontsize=10) + plt.tight_layout(rect=[0, 0.05, 1, 1]) + plt.savefig(args.out, dpi=150, bbox_inches="tight") + print(f"\nSaved to {args.out}") + + +if __name__ == "__main__": + main() + diff --git a/src/camera-based-e2e/Point_cloud/vim_occ_3d.py b/src/camera-based-e2e/Point_cloud/vim_occ_3d.py new file mode 100644 index 0000000..48710d7 --- /dev/null +++ b/src/camera-based-e2e/Point_cloud/vim_occ_3d.py @@ -0,0 +1,65 @@ +import numpy as np +import plotly.graph_objects as go +import argparse + +VOX_XY_RANGE = 25.0 +VOX_XY_RES = 0.5 +VOX_Z_MIN = -3.0 +VOX_Z_RES = 0.5 + +COLORS = { + 0: 'lightgrey', # free + 1: 'red', # vehicle + 2: 'blue', # pedestrian + 3: 'orange', # cyclist + 4: 'darkgrey', # road + 5: 'green', # static +} +NAMES = {0:'free', 1:'vehicle', 2:'pedestrian', 3:'cyclist', 4:'road', 5:'static'} + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--occ_path", type=str, required=True) + parser.add_argument("--show_free", action="store_true") + args = parser.parse_args() + + grid = np.load(args.occ_path) + ix, iy, iz = np.where(grid != 255) + classes = grid[ix, iy, iz] + + x = VOX_XY_RANGE - ix * VOX_XY_RES + y = VOX_XY_RANGE - iy * VOX_XY_RES + z = VOX_Z_MIN + iz * VOX_Z_RES + + traces = [] + for cls in range(0 if args.show_free else 1, 6): + mask = classes == cls + if mask.sum() == 0: + continue + traces.append(go.Scatter3d( + x=x[mask], y=y[mask], z=z[mask], + mode='markers', + marker=dict(size=3, color=COLORS[cls], opacity=0.8), + name=NAMES[cls] + )) + + fig = go.Figure(data=traces) + fig.update_layout( + title="Occupancy Grid 3D View", + scene=dict( + xaxis_title="X (m)", + yaxis_title="Y (m)", + zaxis_title="Z (m)", + camera=dict( + up=dict(x=0, y=0, z=1), + eye=dict(x=0, y=0, z=2.5) # top-down by default + ) + ) + ) + out = args.occ_path.replace('.npy', '.html') + fig.write_html(out) + print(f"Saved to {out} — open in browser") + +if __name__ == "__main__": + main() + diff --git a/src/camera-based-e2e/Point_cloud/viz_seg.py b/src/camera-based-e2e/Point_cloud/viz_seg.py new file mode 100644 index 0000000..c8b27e7 --- /dev/null +++ b/src/camera-based-e2e/Point_cloud/viz_seg.py @@ -0,0 +1,91 @@ +""" +Run SegFormer on extracted frame images and save colored segmentation maps. +Usage: + python viz_seg.py --idx 5000 +""" +import argparse +import os +import numpy as np +import cv2 +import torch +import torch.nn.functional as F +from transformers import AutoImageProcessor, SegformerForSemanticSegmentation +from pathlib import Path + +SEG_MODEL = "nvidia/segformer-b2-finetuned-cityscapes-1024-1024" + +# Cityscapes colors for each class +CITYSCAPES_COLORS = [ + (128, 64,128), # road + (244, 35,232), # sidewalk + ( 70, 70, 70), # building + (102,102,156), # wall + (190,153,153), # fence + (153,153,153), # pole + (250,170, 30), # traffic light + (220,220, 0), # traffic sign + (107,142, 35), # vegetation + (152,251,152), # terrain + ( 70,130,180), # sky + (220, 20, 60), # person + (255, 0, 0), # rider + ( 0, 0,142), # car + ( 0, 0, 70), # truck + ( 0, 60,100), # bus + ( 0, 80,100), # train + ( 0, 0,230), # motorcycle + (119, 11, 32), # bicycle +] + +CITYSCAPES_NAMES = [ + 'road','sidewalk','building','wall','fence','pole', + 'traffic light','traffic sign','vegetation','terrain','sky', + 'person','rider','car','truck','bus','train','motorcycle','bicycle' +] + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--idx", type=int, default=5000) + parser.add_argument("--img_dir", type=str, default="./frame_images") + parser.add_argument("--out_dir", type=str, default="./seg_output") + args = parser.parse_args() + + Path(args.out_dir).mkdir(exist_ok=True) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + print(f"Loading SegFormer...") + processor = AutoImageProcessor.from_pretrained(SEG_MODEL) + model = SegformerForSemanticSegmentation.from_pretrained(SEG_MODEL).to(device) + model.eval() + + for cam in range(1, 9): + img_path = os.path.join(args.img_dir, f"frame_{args.idx:07d}_cam{cam}.jpg") + if not os.path.exists(img_path): + continue + + rgb = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB) + H, W = rgb.shape[:2] + + inputs = processor(images=rgb, return_tensors="pt") + inputs = {k: v.to(device) for k, v in inputs.items()} + + with torch.no_grad(): + logits = model(**inputs).logits + + pred = F.interpolate(logits.float(), size=(H, W), mode="bilinear", align_corners=False) + class_map = pred.argmax(dim=1).squeeze().cpu().numpy() + + # Color the segmentation + seg_colored = np.zeros((H, W, 3), dtype=np.uint8) + for cls_id, color in enumerate(CITYSCAPES_COLORS): + seg_colored[class_map == cls_id] = color + + # Side by side: original + segmentation + combined = np.concatenate([rgb[:,:,::-1], seg_colored[:,:,::-1]], axis=1) + out_path = os.path.join(args.out_dir, f"seg_{args.idx:07d}_cam{cam}.jpg") + cv2.imwrite(out_path, combined) + print(f" Saved cam{cam}: {out_path}") + +if __name__ == "__main__": + main() + diff --git a/src/camera-based-e2e/check_occ.py b/src/camera-based-e2e/check_occ.py new file mode 100644 index 0000000..4443ec8 --- /dev/null +++ b/src/camera-based-e2e/check_occ.py @@ -0,0 +1,145 @@ +""" +OCC sanity checker — run from the camera-based-e2e root: + + python check_occ.py + +Paths are taken directly from train_occ.slurm. +""" + +import pickle +import os +import numpy as np + +SCRIPT_DIR = "/scratch/gilbreth/kumar753/robotvision/robotvision/src/camera-based-e2e" +OCC_ROOT = "/scratch/gilbreth/kumar753/waymo_occ_new" +INDEX_TRAIN = os.path.join(SCRIPT_DIR, "index_train.pkl") +INDEX_VAL = os.path.join(SCRIPT_DIR, "index_val.pkl") +N_SAMPLES = 20 + + +def load_indexes(pkl_path): + with open(pkl_path, "rb") as f: + raw = pickle.load(f) + return [(i, item) for i, item in enumerate(raw)] + + +def check_split(split, indexes, n_samples, label): + print(f"\n{'='*60}") + print(f" {label} ({split}) — checking {n_samples} of {len(indexes)} samples") + print(f"{'='*60}") + + sample_idxs = list(range(min(n_samples, len(indexes)))) + missing, bad_shape, bad_dtype = [], [], [] + all_unique = set() + + for i in sample_idxs: + orig_idx, _ = indexes[i] + path = os.path.join(OCC_ROOT, split, f"occ_{orig_idx:07d}.npy") + + if not os.path.exists(path): + missing.append((i, orig_idx, path)) + continue + + arr = np.load(path) + + if arr.shape != (100, 100, 16): + bad_shape.append((i, orig_idx, arr.shape)) + + if arr.dtype != np.uint8: + bad_dtype.append((i, orig_idx, str(arr.dtype))) + + all_unique.update(arr.flatten().tolist()) + + # ---- report ---- + print(f"\n[1] File existence") + if missing: + print(f" MISSING {len(missing)} files!") + for i, orig, p in missing[:5]: + print(f" sample[{i}] orig_idx={orig} -> {p}") + if len(missing) > 5: + print(f" ... and {len(missing)-5} more") + else: + print(f" OK — all {len(sample_idxs)} files found") + + print(f"\n[2] Shape (expected (100, 100, 16))") + if bad_shape: + print(f" BAD SHAPES on {len(bad_shape)} files!") + for i, orig, sh in bad_shape[:5]: + print(f" sample[{i}] orig_idx={orig} -> shape={sh}") + else: + print(f" OK — all shapes correct") + + print(f"\n[3] Dtype (expected uint8)") + if bad_dtype: + print(f" BAD DTYPE on {len(bad_dtype)} files!") + for i, orig, dt in bad_dtype[:5]: + print(f" sample[{i}] orig_idx={orig} -> dtype={dt}") + else: + print(f" OK — all uint8") + + print(f"\n[4] Unique values across sampled files") + sorted_vals = sorted(all_unique) + print(f" {sorted_vals}") + if 255 in all_unique: + print(f" (255 present — used as ignore_index in cross_entropy, OK)") + class_vals = [v for v in sorted_vals if v != 255] + n_classes = len(class_vals) + print(f" Class range: [{min(class_vals) if class_vals else 'N/A'}, {max(class_vals) if class_vals else 'N/A'}] => {n_classes} distinct classes") + if n_classes != 6: + print(f" WARNING: occ_head outputs 6 classes but ground truth has {n_classes} — mismatch!") + else: + print(f" OK — matches occ_head output of 6 classes") + + print(f"\n[5] Cross-entropy smoke test") + try: + import torch + import torch.nn.functional as F + orig_idx, _ = indexes[0] + path = os.path.join(OCC_ROOT, split, f"occ_{orig_idx:07d}.npy") + if os.path.exists(path): + arr = np.load(path).astype(np.int64) + occ_gt = torch.from_numpy(arr).unsqueeze(0) # (1, 100, 100, 16) + pred = torch.randn(1, 6, 100, 100, 16) # matches occ_head output + loss = F.cross_entropy(pred, occ_gt, ignore_index=255) + print(f" OK — loss={loss.item():.4f} (random pred, just checking shapes)") + else: + print(f" SKIPPED — first file missing") + except Exception as e: + print(f" ERROR: {e}") + + return len(missing), len(bad_shape), len(bad_dtype) + + +def main(): + print(f"\nOCC_ROOT = {OCC_ROOT}") + print(f"INDEX_TRAIN = {INDEX_TRAIN}") + print(f"INDEX_VAL = {INDEX_VAL}") + + train_idxs = load_indexes(INDEX_TRAIN) + val_idxs = load_indexes(INDEX_VAL) + + print(f"\n[6] Index alignment check") + print(f" train orig_idx range: 0 .. {len(train_idxs)-1}") + print(f" val orig_idx range: 0 .. {len(val_idxs)-1}") + print(f" Both restart from 0 — occ files must live in separate subdirs:") + print(f" {OCC_ROOT}/train/occ_XXXXXXX.npy") + print(f" {OCC_ROOT}/val/occ_XXXXXXX.npy") + train_dir = os.path.join(OCC_ROOT, "train") + val_dir = os.path.join(OCC_ROOT, "val") + print(f" train dir exists: {os.path.isdir(train_dir)}") + print(f" val dir exists: {os.path.isdir(val_dir)}") + + tm, ts, td = check_split("train", train_idxs, N_SAMPLES, "TRAIN") + vm, vs, vd = check_split("val", val_idxs, N_SAMPLES, "VAL") + + total_issues = tm + ts + td + vm + vs + vd + print(f"\n{'='*60}") + if total_issues == 0: + print(" All checks passed!") + else: + print(f" {total_issues} issue(s) found — see above") + print(f"{'='*60}\n") + + +if __name__ == "__main__": + main() diff --git a/src/camera-based-e2e/check_past_stats.py b/src/camera-based-e2e/check_past_stats.py new file mode 100644 index 0000000..5b46566 --- /dev/null +++ b/src/camera-based-e2e/check_past_stats.py @@ -0,0 +1,88 @@ +""" +Check past state statistics from the dataset. +Run from camera-based-e2e root: + + python check_past_stats.py +""" + +import pickle +import os +import numpy as np +import sys + +sys.path.insert(0, "/scratch/gilbreth/kumar753/robotvision/robotvision/src/camera-based-e2e") + +from protos import e2e_pb2 + +INDEX_FILE = "/scratch/gilbreth/kumar753/robotvision/robotvision/src/camera-based-e2e/index_train.pkl" +DATA_DIR = "/scratch/gilbreth/kumar753/robotvision/waymo_end_to_end_camera_v1_0_0/waymo_open_dataset_end_to_end_camera_v_1_0_0" +N_SAMPLES = 200 + +FEATURE_NAMES = ["pos_x", "pos_y", "vel_x", "vel_y", "accel_x", "accel_y"] + +def main(): + with open(INDEX_FILE, "rb") as f: + indexes = pickle.load(f) + + print(f"Checking {N_SAMPLES} samples from {INDEX_FILE}\n") + + all_past = [] + current_file, fh = None, None + + for i in range(min(N_SAMPLES, len(indexes))): + filename, start_byte, byte_length = indexes[i] + path = os.path.join(DATA_DIR, filename) + + if current_file != path: + if fh: + fh.close() + fh = open(path, "rb") + current_file = path + + fh.seek(start_byte) + frame = e2e_pb2.E2EDFrame() + frame.ParseFromString(fh.read(byte_length)) + + past = np.stack([ + frame.past_states.pos_x, + frame.past_states.pos_y, + frame.past_states.vel_x, + frame.past_states.vel_y, + frame.past_states.accel_x, + frame.past_states.accel_y, + ], axis=-1).astype(np.float32) # (16, 6) + + all_past.append(past) + + if fh: + fh.close() + + all_past = np.stack(all_past, axis=0) # (N, 16, 6) + + print(f"{'Feature':<12} {'mean':>10} {'std':>10} {'min':>10} {'max':>10}") + print("-" * 55) + for i, name in enumerate(FEATURE_NAMES): + col = all_past[:, :, i] + print(f"{name:<12} {col.mean():>10.3f} {col.std():>10.3f} {col.min():>10.3f} {col.max():>10.3f}") + + print("\nOverall (all features):") + print(f" mean={all_past.mean():.3f} std={all_past.std():.3f} min={all_past.min():.3f} max={all_past.max():.3f}") + + print("\nDiagnosis:") + pos = all_past[:, :, :2] + if np.abs(pos).max() > 50: + print(" WARNING: pos_x/pos_y are in absolute world coordinates (large values).") + print(" If your supervisor normalized these, that explains the ADE gap.") + print(" Consider subtracting the last observed position:") + print(" past[:, :, :2] -= past[:, -1:, :2]") + else: + print(" OK: positions look relative/small — probably already normalized.") + + vel = all_past[:, :, 2:4] + if np.abs(vel).max() > 50: + print(" WARNING: vel_x/vel_y have very large values — check units (m/s vs km/h?).") + else: + print(" OK: velocities look reasonable.") + +if __name__ == "__main__": + main() diff --git a/src/camera-based-e2e/index_test.pkl b/src/camera-based-e2e/index_test.pkl index a39c96d..30b0424 100644 Binary files a/src/camera-based-e2e/index_test.pkl and b/src/camera-based-e2e/index_test.pkl differ diff --git a/src/camera-based-e2e/index_train.pkl b/src/camera-based-e2e/index_train.pkl index 4d9f7af..355cf94 100644 Binary files a/src/camera-based-e2e/index_train.pkl and b/src/camera-based-e2e/index_train.pkl differ diff --git a/src/camera-based-e2e/index_val.pkl b/src/camera-based-e2e/index_val.pkl index e007004..df12e8f 100644 Binary files a/src/camera-based-e2e/index_val.pkl and b/src/camera-based-e2e/index_val.pkl differ diff --git a/src/camera-based-e2e/loader_baseline.py b/src/camera-based-e2e/loader_baseline.py new file mode 100644 index 0000000..cbe4574 --- /dev/null +++ b/src/camera-based-e2e/loader_baseline.py @@ -0,0 +1,125 @@ +import torch +from torch.utils.data import Dataset +from protos import e2e_pb2 +import pickle +import os +import numpy as np +from typing import Optional +import random + +devices = ['cuda:0', 'cuda:1'] + +random.seed(42) # Deterministic + + +class WaymoE2E(Dataset): + def __init__( + self, + indexFile='index.pkl', + data_dir='./dataset', + n_items: Optional[int] = None, + seed: Optional[int] = None, + ): + self.data_dir = data_dir + self.seed = seed + + self.filename = "" + self.file = None + + with open(indexFile, 'rb') as f: + # NOTE: test does not have reference trajectories + # We train on train and validate on val set + self.indexes = pickle.load(f) + + # Use the first n_items in dataset order + if n_items is not None and n_items < len(self.indexes): + self.indexes = self.indexes[:n_items] + + def __len__(self): + return len(self.indexes) + + def __getitem__(self, idx): + frame = e2e_pb2.E2EDFrame() # type: ignore + filename, start_byte, byte_length = self.indexes[idx] + + if self.filename != filename: + if self.file: + self.file.close() + del self.file + self.file = open(os.path.join(self.data_dir, filename), 'rb') + self.filename = filename + + self.file.seek(start_byte) # type: ignore + protobuf = self.file.read(byte_length) # type: ignore + frame.ParseFromString(protobuf) + + past = np.stack( + [ + frame.past_states.pos_x, + frame.past_states.pos_y, + frame.past_states.vel_x, + frame.past_states.vel_y, + frame.past_states.accel_x, + frame.past_states.accel_y, + ], + axis=-1, + ) + + future = np.stack( + [frame.future_states.pos_x, frame.future_states.pos_y], axis=-1 + ) + + past = np.array(past, dtype=np.float32) + future = np.array(future, dtype=np.float32) + + # For submission to waymo evaluation server + name = frame.frame.context.name + + # Return JPEG images as torch uint8 tensors so DataLoader can use shared memory. + jpeg_tensors = [ + torch.from_numpy(np.frombuffer(img.image, dtype=np.uint8).copy()) + for img in frame.frame.images + ] + + return { + 'PAST': past, + 'FUTURE': future, + 'IMAGES_JPEG': jpeg_tensors, + 'INTENT': frame.intent, + 'NAME': name, + } + + +def collate_with_images(batch): + """Collate that keeps IMAGES_JPEG as a list-of-lists (variable-size JPEG + bytes cannot be stacked) and delegates everything else to default_collate.""" + from torch.utils.data.dataloader import default_collate + + images = [sample.pop('IMAGES_JPEG') for sample in batch] + collated = default_collate(batch) + collated['IMAGES_JPEG'] = images # list[list[Tensor]], one inner list per sample + return collated + + +if __name__ == "__main__": + from torch.utils.data import DataLoader + from tqdm import tqdm + + # NOTE: Replace with your path + DATA_DIR = '/anvil/scratch/x-mgagvani/wod/waymo_end_to_end_camera_v1_0_0/waymo_open_dataset_end_to_end_camera_v_1_0_0' + BATCH_SIZE = 256 + dataset = WaymoE2E(indexFile="index_train.pkl", data_dir=DATA_DIR) + loader = DataLoader( + dataset, + batch_size=BATCH_SIZE, + num_workers=0, + collate_fn=collate_with_images, + pin_memory=True, # causes error + ) + + def main(): + for batch_of_frames in tqdm(loader): + pass + + import cProfile + main() diff --git a/src/camera-based-e2e/loader_occ.py b/src/camera-based-e2e/loader_occ.py new file mode 100644 index 0000000..b73210c --- /dev/null +++ b/src/camera-based-e2e/loader_occ.py @@ -0,0 +1,162 @@ +import torch +from torch.utils.data import Dataset +from protos import e2e_pb2 +import pickle +import os +import numpy as np +from typing import Optional +import random + +devices = ['cuda:0', 'cuda:1'] + +random.seed(42) # Deterministic + +class WaymoE2E(Dataset): + def __init__( + self, + indexFile = 'index.pkl', + data_dir='./dataset', + n_items: Optional[int] = None, + seed: Optional[int] = None, + occ_root: Optional[str] = None, + ): + self.data_dir = data_dir + self.seed = seed + + self.filename = "" + self.file = None + + with open(indexFile, 'rb') as f: + # NOTE: test does not have reference trajectories + # We train on train and validate on val set + raw_indexes = pickle.load(f) + + self.indexes = [(orig_idx, item) for orig_idx, item in enumerate(raw_indexes)] + + # TODO: Determine how to sample specific subsets of the data that we care about. + if n_items is not None and n_items < len(self.indexes): + self.indexes = self.indexes[:n_items] + + self.occ_root = occ_root + self.occ_index = None + self.split = "train" if "train" in os.path.basename(indexFile) else "val" if "val" in os.path.basename(indexFile) else "test" + + def __len__(self): + return len(self.indexes) + + def __getitem__(self, idx): + frame = e2e_pb2.E2EDFrame() # type: ignore + orig_idx, (filename, start_byte, byte_length) = self.indexes[idx] + + if self.filename != filename: + if self.file: + self.file.close() + del self.file + self.file = open(os.path.join(self.data_dir, filename), 'rb') + self.filename = filename + + self.file.seek(start_byte) # type: ignore + protobuf = self.file.read(byte_length) # type: ignore + frame.ParseFromString(protobuf) + + past = np.stack( + [ + frame.past_states.pos_x, + frame.past_states.pos_y, + frame.past_states.vel_x, + frame.past_states.vel_y, + frame.past_states.accel_x, + frame.past_states.accel_y, + ], + axis=-1, + ) + + future = np.stack([frame.future_states.pos_x, frame.future_states.pos_y], axis=-1) + + past = np.array(past, dtype=np.float32) # ensure consistent dtype + future = np.array(future, dtype=np.float32) + + # For submission to waymo evaluation server + name = frame.frame.context.name + + # Return JPEG images as torch uint8 tensors so DataLoader can use shared memory. + jpeg_tensors = [ + torch.from_numpy(np.frombuffer(img.image, dtype=np.uint8).copy()) + for img in frame.frame.images + ] + + occ = None + if self.occ_root is not None: + occ_path = os.path.join(self.occ_root, self.split, f"occ_{orig_idx:07d}.npy") + if not os.path.exists(occ_path): + raise FileNotFoundError(f"Missing OCC file for index {orig_idx}: {occ_path}") + occ = np.load(occ_path).astype(np.uint8) + + return {'PAST': past, 'FUTURE': future, 'IMAGES_JPEG': jpeg_tensors, 'INTENT': frame.intent, 'NAME': name, 'OCC': occ} + + +def collate_with_images(batch): + """Collate that keeps IMAGES_JPEG as a list-of-lists (variable-size JPEG + bytes cannot be stacked) and delegates everything else to default_collate + when OCC is not present.""" + from torch.utils.data.dataloader import default_collate + + if batch[0].get("OCC", None) is None: + images = [sample.pop('IMAGES_JPEG') for sample in batch] + for sample in batch: + sample.pop('OCC', None) + collated = default_collate(batch) + collated['IMAGES_JPEG'] = images # list[list[Tensor]], one inner list per sample + return collated + + past = [torch.as_tensor(b["PAST"], dtype=torch.float32) for b in batch] + future = [torch.as_tensor(b["FUTURE"], dtype=torch.float32) for b in batch] + intent = torch.as_tensor([b["INTENT"] for b in batch]) + names = [b["NAME"] for b in batch] + + cams = list(zip(*[b["IMAGES_JPEG"] for b in batch])) # per-camera tuples + images_jpeg = [list(cam_imgs) for cam_imgs in cams] # stay on CPU + + out = { + "PAST": torch.stack(past, dim=0), + "FUTURE": torch.stack(future, dim=0), + "INTENT": intent, + "IMAGES_JPEG": images_jpeg, + "NAME": names, + "OCC": torch.stack( + [torch.as_tensor(b["OCC"], dtype=torch.long) for b in batch], + dim=0, + ), + } + + return out + + +if __name__ == "__main__": + + from torch.utils.data import DataLoader + import time + from tqdm import tqdm + # NOTE: Replace with your path + DATA_DIR = '/anvil/scratch/x-mgagvani/wod/waymo_end_to_end_camera_v1_0_0/waymo_open_dataset_end_to_end_camera_v_1_0_0' + BATCH_SIZE = 256 + dataset = WaymoE2E(indexFile="index_train.pkl", data_dir = DATA_DIR) + loader = DataLoader( + dataset, + batch_size=BATCH_SIZE, + num_workers=0, + collate_fn=collate_with_images, + pin_memory=True, # causes error + ) + # next(iter(loader)) + + def main(): + # start = time.time() + for batch_of_frames in tqdm(loader): + # print(batch_of_frames["INTENT"]) + # print(batch_of_frames.keys(), [b.shape for b in batch_of_frames.values() if isinstance(b, torch.Tensor)]) + pass + # print("Total Time:", time.time()-start) + + import cProfile + main() diff --git a/src/camera-based-e2e/models/base_model_baseline.py b/src/camera-based-e2e/models/base_model_baseline.py new file mode 100644 index 0000000..17782e6 --- /dev/null +++ b/src/camera-based-e2e/models/base_model_baseline.py @@ -0,0 +1,386 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import pytorch_lightning as pl +import torchvision +from dataclasses import asdict, is_dataclass + +from .losses.depth_loss import DepthLoss + +class BaseModel(nn.Module): + def __init__(self, in_dim, out_dim): + super(BaseModel, self).__init__() + + # This is literally just linear regression = y_hat = Wx + b + self.nn = nn.Sequential( + nn.Linear(in_dim, out_dim) + ) + + def forward(self, x: dict) -> torch.Tensor: + past, images, intent = x['PAST'], x['IMAGES'], x['INTENT'] + x = past.reshape(past.size(0), -1) # Flatten to (B, 16 * 6) = (B, 96) + return self.nn(x) + +class LitModel(pl.LightningModule): + def __init__(self, model: nn.Module, lr: float, lr_vision: float | None = None, rfs_weight: float = 0.0): + super(LitModel, self).__init__() + self.model = model + + # If we are using ScorerModel, which has a cfg, then save the attributes of the cfg as hparams, so they go into wandb + cfg = getattr(model, "cfg", None) + if cfg is None: + cfg_dict = {} + elif is_dataclass(cfg): + cfg_dict = asdict(cfg) + elif isinstance(cfg, dict): + cfg_dict = dict(cfg) + else: + try: + cfg_dict = dict(vars(cfg)) + except TypeError: + cfg_dict = {"repr": repr(cfg)} + + hparams = { + "lr": lr, + "lr_vision": lr_vision, + "rfs_weight": rfs_weight, + "model_name": model.__class__.__name__, + "model_cfg": cfg_dict, + } + for k, v in cfg_dict.items(): + if isinstance(v, (int, float, str, bool)) or v is None: + hparams[f"model_cfg_{k}"] = v + + self.example_input_array = ({ + 'PAST': torch.zeros((1, 16, 6)), # PAST + 'IMAGES': [torch.zeros((1, 3, 1280, 1920)) for _ in range(6)], # IMAGES + 'INTENT': torch.tensor([1.0]), # INTENT + },) + + self.save_hyperparameters(hparams, ignore=["model"]) + + # --- Data Loading ---- + def transfer_batch_to_device(self, batch, device, dataloader_idx): + # if not a dict, it's not actually proper training data, so delegate this to the super()class + if not isinstance(batch, dict): + return super().transfer_batch_to_device(batch, device, dataloader_idx) + + # don't move images_jpeg to gpu, move the decoded images to gpu + if "IMAGES_JPEG" in batch: + images_jpeg = batch["IMAGES_JPEG"] + batch_wo_jpeg = dict(batch) + batch_wo_jpeg.pop("IMAGES_JPEG", None) + moved = super().transfer_batch_to_device(batch_wo_jpeg, device, dataloader_idx) + moved["IMAGES"] = self.decode_batch_jpeg(images_jpeg, device=device) + return moved + + return super().transfer_batch_to_device(batch, device, dataloader_idx) + + + def decode_batch_jpeg( + self, + images_jpeg: list[list[torch.Tensor]], + device: torch.device | None = None, + ) -> list[torch.Tensor]: + decode_device = self.device if device is None else device + # Flatten cameras + flat_encoded, cam_sizes = [], [] + for cam in images_jpeg: + cam_sizes.append(len(cam)) + for jpg in cam: + t = jpg if isinstance(jpg, torch.Tensor) else torch.frombuffer(memoryview(jpg), dtype=torch.uint8) + # decode_jpeg requires the raw jpeg bytes to be on cpu + if t.device.type != "cpu": + t = t.cpu() + flat_encoded.append(t) + + flat_decoded = torchvision.io.decode_jpeg( + flat_encoded, + mode=torchvision.io.ImageReadMode.UNCHANGED, + device=decode_device, + ) # list of (C, H, W) gpu tensors + + out = [] + idx = 0 + for n in cam_sizes: + cam_list = flat_decoded[idx: idx+n] + idx += n + out.append(torch.stack(cam_list, dim=0)) # (B, C, H, W) + return out + + def on_fit_start(self) -> None: + super().on_fit_start() + self.depth_loss = DepthLoss(self.device) + + # ---- Metrics ---- + def ade_loss(self, pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor: + """ + Average Displacement Error -> L2 Norm -> Average Euclidean Distance between predicted and ground truth future trajectory + """ + return torch.mean(torch.norm(pred - gt, dim=-1)) + + def time_thresholds(self, t_idx): + # Time-based thresholds at 3s and 5s. + lat = torch.where(t_idx <= 3, 1.0, 1.8) + lng = torch.where(t_idx <= 3, 4.0, 7.2) + return lat, lng + + def speed_scale(self, v): + # Speed-based scaling copied from RFS paper. + return torch.where( + v < 1.4, + 0.5, + torch.where( + v < 11.0, + 0.5 + 0.5 * (v - 1.4) / (11.0 - 1.4), + 1.0 + ) + ) + + def compute_direction(self, trajectory): + # Pad with first point so displacement stays (B, T, 2). + padded = torch.cat([trajectory[:, :1], trajectory], dim=1) + displacement = padded[:, 1:] - padded[:, :-1] + lng_dir = F.normalize(displacement, p=2, dim=-1, eps=1e-6) + lat_dir = torch.stack([-lng_dir[..., 1], lng_dir[..., 0]], dim=-1) + return lng_dir, lat_dir + + def rfs_loss(self, pred, gt, lng_dir, lat_dir, speed, t_idx): + """ + pred, gt: (B, T, 2) + speed: (B,) or (B, T) + t_idx: (T,) or (B, T) + """ + delta = pred - gt + delta_lng = (delta * lng_dir).sum(dim=-1).abs() + delta_lat = (delta * lat_dir).sum(dim=-1).abs() + + tau_lat_raw, tau_lng_raw = self.time_thresholds(t_idx) + scale = self.speed_scale(speed) + if scale.dim() == 1: + scale = scale.unsqueeze(1) + + tau_lat = tau_lat_raw * scale + tau_lng = tau_lng_raw * scale + + deviation = torch.max( + delta_lat / tau_lat, + delta_lng / tau_lng, + ) + score = torch.where( + deviation <= 1, + torch.ones_like(deviation), + torch.pow(0.1, deviation - 1) + ) + return (1.0 - score).mean() + + def _prepare_rfs_inputs(self, past, future, pred_future): + speed = torch.norm(past[..., 2:4], dim=-1)[:, -1] # (B,), speed at last observed time step + full_lng_dir, full_lat_dir = self.compute_direction(future) + indices = [11, 19] # 3s and 5s into the future + + pred_slice = pred_future[:, indices, :] + gt_slice = future[:, indices, :] + lng_dir_slice = full_lng_dir[:, indices, :] + lat_dir_slice = full_lat_dir[:, indices, :] + t_idx = torch.tensor([3.0, 5.0], device=future.device).unsqueeze(0).expand(future.size(0), -1) + return pred_slice, gt_slice, lng_dir_slice, lat_dir_slice, speed, t_idx + + # ---- optimizers ---- + def configure_optimizers(self): + # NOTE: This can be extended and tuned, LR especially will differ and have an impact. + # vision encoder, if trainable, should have 1/10 the LR of the rest of the model + if hasattr(self.model, "features"): + encoder_params = [p for p in self.model.features.parameters() if p.requires_grad] + other_params = [ + p for n, p in self.model.named_parameters() + if not n.startswith("features.") and p.requires_grad + ] + if encoder_params: + encoder_lr = self.hparams.lr * 0.1 if self.hparams.lr_vision is None else self.hparams.lr_vision + return torch.optim.Adam( + [ + {"params": other_params, "lr": self.hparams.lr}, + {"params": encoder_params, "lr": encoder_lr}, + ] + ) + if other_params: + return torch.optim.Adam(other_params, lr=self.hparams.lr) + + optimizer = torch.optim.Adam(self.model.parameters(), lr=self.hparams.lr) + return optimizer + + # ---- forward / step ---- + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.model(x) + + def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor: + past, future, intent = batch['PAST'], batch['FUTURE'], batch['INTENT'] + + if "IMAGES" in batch: + images = batch["IMAGES"] + elif "IMAGES_JPEG" in batch: + images_jpeg = batch["IMAGES_JPEG"] + images = self.decode_batch_jpeg(images_jpeg) + else: + raise KeyError("Batch must contain either 'IMAGES_JPEG' or 'IMAGES' key.") + + # `past` is our input (B, 16, 6) e.g. Batch x Time x (x, y, v_x, v_y, a_x, a_y) + # and `future` is our output (B, 20, 2) e.g. Batch x Time x (x, y) + + # create all input data that we are allowed to give to a model + model_inputs = {'PAST': past, 'IMAGES': images, 'INTENT': intent} + + pred_future = self.forward(model_inputs) # (B, T*2) + pred_depth = None + pred_scores: torch.Tensor = None + if isinstance(pred_future, dict): + pred_future, pred_depth, pred_scores = pred_future["trajectory"], pred_future.get("depth", None), pred_future.get("scores", None) + + pred = pred_future + t_steps = future.shape[1] + t2 = t_steps * 2 + + if pred.ndim != 2: + raise ValueError(f"Unexpected pred shape {pred.shape}; expected 2D (B, K*T*2).") + + if pred.shape[1] % t2 != 0: + raise ValueError(f"pred dim1={pred.shape[1]} is not divisible by T*2={t2}.") + k_modes = pred.shape[1] // t2 + pred = pred.view(pred.size(0), k_modes, t_steps, 2) + + if pred_scores is not None and pred.size(1) > 1: + rfs_pred_idx = pred_scores.argmin(dim=1) + else: + rfs_pred_idx = torch.zeros(pred.size(0), dtype=torch.long, device=pred.device) + pred_for_rfs = pred[torch.arange(pred.size(0), device=pred.device), rfs_pred_idx] + + pred_slice, gt_slice, lng_dir_slice, lat_dir_slice, speed, t_idx = self._prepare_rfs_inputs( + past, + future, + pred_for_rfs, + ) + rfs_unweighted = self.rfs_loss(pred_slice, gt_slice, lng_dir_slice, lat_dir_slice, speed, t_idx) + rfs_weight = getattr(self.hparams, "rfs_weight", 0.0) + loss_rfs = rfs_weight * rfs_unweighted + + loss_type = getattr(self.hparams, "model_cfg_loss_type", "mse") + + # ADE per mode: (B, K) + dist = torch.norm(pred - future[:, torch.newaxis, :, :], dim=-1) # (B, K, T) + ade_per_mode = dist.mean(dim=-1) + + # Top-M WTA for trajectory loss. Here, we have an "oracle" that picks the best mode + # so, our loss is calculated on the mean of the top n trajectories. + top_m = min(getattr(self.hparams, "model_cfg_loss_top_n", 5), ade_per_mode.size(1)) + loss_ade = ade_per_mode.topk(top_m, largest=False, dim=1).values.mean() + + # oracle ade is best of all proposals, since we have the GT data during training + oracle_ade = ade_per_mode.min(dim=1).values.mean() + ade_pred = None + # pred_scores is now the predicted ADE of each trajectory / expectation loss + if pred_scores is not None and k_modes > 1: + pred_idx = pred_scores.argmin(dim=1) + ade_pred = ade_per_mode[torch.arange(pred.size(0), device=pred.device), pred_idx].mean() + elif k_modes == 1: + ade_pred = ade_per_mode.squeeze(1).mean() + regret = (ade_pred - oracle_ade) if ade_pred is not None else None + + # Scorer Losses -> encourage ranking of predicted scores to match true ranking of ades that are generated + if k_modes > 1 and pred_scores is not None: + ade = ade_per_mode.detach() # (B, K) + if loss_type == "mse": + loss_score = F.mse_loss(pred_scores, ade) + elif loss_type == "reinforce": + tau_base = getattr(self.hparams, "model_cfg_loss_tau_base", 1.0) + decay_factor = getattr(self.hparams, "model_cfg_loss_tau_decay", 0.95) + entropy_lambda = getattr(self.hparams, "model_cfg_loss_entropy_lambda", 0.01) + # p_k = softmax(s_k / tau) + logits = -pred_scores / max(0.1, tau_base * (decay_factor ** self.current_epoch)) + p = F.softmax(logits, dim=1) + # Loss = sum_k {p_k * e_k} (e.g., expectation of e_k) + loss_selection = (p * ade).sum(dim=1).mean() + # H(p) = -sum(p) * log(p) + entropy = -(p * (p + 1e-8).log()).sum(dim=1).mean() + loss_score = loss_selection - entropy_lambda * entropy + else: + raise NotImplementedError(f"Loss {loss_type} is not implemented") + pred_idx = pred_scores.argmin(dim=1) + ade_pred = ade_per_mode[torch.arange(pred.size(0), device=pred.device), pred_idx].mean() + else: + loss_score = torch.tensor(0.0, device=self.device) + + # Scorer Metrics + scorer_metrics = {} + if k_modes > 1 and pred_scores is not None: + # Rank of the scorer's top-1 pick among oracle-sorted proposals + oracle_ranking = ade_per_mode.argsort(dim=1) # (B, K) indices sorted by true ADE + oracle_rank_of = oracle_ranking.argsort(dim=1) # (B, K) rank of each proposal + scorer_pick = pred_scores.argmin(dim=1) # (B,) scorer's best + pick_rank = oracle_rank_of[torch.arange(pred.size(0), device=pred.device), scorer_pick].float() # (B,) + scorer_metrics[f"{stage}_scorer_mean_rank"] = pick_rank.mean() + for topk in (1, 5, 10): + if topk <= k_modes: + scorer_metrics[f"{stage}_scorer_top{topk}_acc"] = (pick_rank < topk).float().mean() + # Spearman rank correlation + K = ade_per_mode.size(1) + oracle_ranks = oracle_rank_of.float() # (B, K) + scorer_ranks = pred_scores.argsort(dim=1).argsort(dim=1).float() # (B, K) + d = oracle_ranks - scorer_ranks + rho = 1 - 6 * (d ** 2).sum(dim=1) / (K * (K ** 2 - 1)) + scorer_metrics[f"{stage}_scorer_spearman"] = rho.mean() + + # Depth Loss + if pred_depth is not None: + front_img = images[1] # front camera + depth_in = F.interpolate(front_img, size=(128, 128), mode='nearest') + loss_depth = self.depth_loss(depth_in, pred_depth, loss_fn=F.l1_loss) + else: + loss_depth = torch.tensor(0.0, device=self.device) + + loss_depth *= 0.1 # slightly enabled + loss_ade *= 1.0 # TODO: tune loss terms + loss_score *= 1.0 + total_loss = loss_ade + loss_depth + loss_score + loss_rfs + # TODO: improve logging both to disk and to console + log_payload = { + f"{stage}_loss_ade": loss_ade, + f"{stage}_loss_score": loss_score, + f"{stage}_loss_depth": loss_depth, + f"{stage}_loss_rfs": loss_rfs, + f"{stage}_rfs_unweighted": rfs_unweighted, + f"{stage}_loss": total_loss, + } + if ade_pred is not None: + log_payload[f"{stage}_ade_pred"] = ade_pred + log_payload[f"{stage}_ade_oracle"] = oracle_ade + log_payload[f"{stage}_ade_regret"] = regret + log_payload.update(scorer_metrics) + self.log_dict(log_payload, prog_bar=True, logger=True, + batch_size=past.size(0), + sync_dist=(stage == "val")) + + return total_loss + + def training_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor: + return self._shared_step(batch, "train") + + def validation_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor: + return self._shared_step(batch, "val") + +def collate_with_images(batch): + past = [torch.as_tensor(b["PAST"], dtype=torch.float32) for b in batch] + future = [torch.as_tensor(b["FUTURE"], dtype=torch.float32) for b in batch] + intent = torch.as_tensor([b["INTENT"] for b in batch]) + names = [b["NAME"] for b in batch] + + cams = list(zip(*[b["IMAGES_JPEG"] for b in batch])) # per-camera tuples + images_jpeg = [list(cam_imgs) for cam_imgs in cams] # stay on CPU + + return { + "PAST": torch.stack(past, dim=0), + "FUTURE": torch.stack(future, dim=0), + "INTENT": intent, + "IMAGES_JPEG": images_jpeg, + "NAME": names, + } diff --git a/src/camera-based-e2e/models/base_model_occ.py b/src/camera-based-e2e/models/base_model_occ.py new file mode 100644 index 0000000..2bc1ab5 --- /dev/null +++ b/src/camera-based-e2e/models/base_model_occ.py @@ -0,0 +1,407 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import pytorch_lightning as pl +import torchvision +from dataclasses import asdict, is_dataclass + +from .losses.depth_loss import DepthLoss + +class BaseModel(nn.Module): + def __init__(self, in_dim, out_dim): + super(BaseModel, self).__init__() + + # This is literally just linear regression = y_hat = Wx + b + self.nn = nn.Sequential( + nn.Linear(in_dim, out_dim) + ) + + def forward(self, x: dict) -> torch.Tensor: + past, images, intent = x['PAST'], x['IMAGES'], x['INTENT'] + x = past.reshape(past.size(0), -1) # Flatten to (B, 16 * 6) = (B, 96) + return self.nn(x) + +class LitModel(pl.LightningModule): + def __init__(self, model: nn.Module, lr: float, lr_vision: float | None = None, rfs_weight: float = 0.0): + super(LitModel, self).__init__() + self.model = model + + # If we are using ScorerModel, which has a cfg, then save the attributes of the cfg as hparams, so they go into wandb + cfg = getattr(model, "cfg", None) + if cfg is None: + cfg_dict = {} + elif is_dataclass(cfg): + cfg_dict = asdict(cfg) + elif isinstance(cfg, dict): + cfg_dict = dict(cfg) + else: + try: + cfg_dict = dict(vars(cfg)) + except TypeError: + cfg_dict = {"repr": repr(cfg)} + + hparams = { + "lr": lr, + "lr_vision": lr_vision, + "rfs_weight": rfs_weight, + "model_name": model.__class__.__name__, + "model_cfg": cfg_dict, + } + for k, v in cfg_dict.items(): + if isinstance(v, (int, float, str, bool)) or v is None: + hparams[f"model_cfg_{k}"] = v + + self.example_input_array = ({ + 'PAST': torch.zeros((1, 16, 6)), # PAST + 'IMAGES': [torch.zeros((1, 3, 1280, 1920)) for _ in range(6)], # IMAGES + 'INTENT': torch.tensor([1.0]), # INTENT + },) + + self.save_hyperparameters(hparams, ignore=["model"]) + + # --- Data Loading ---- + def transfer_batch_to_device(self, batch, device, dataloader_idx): + # if not a dict, it's not actually proper training data, so delegate this to the super()class + if not isinstance(batch, dict): + return super().transfer_batch_to_device(batch, device, dataloader_idx) + + # don't move images_jpeg to gpu, move the decoded images to gpu + if "IMAGES_JPEG" in batch: + images_jpeg = batch["IMAGES_JPEG"] + batch_wo_jpeg = dict(batch) + batch_wo_jpeg.pop("IMAGES_JPEG", None) + moved = super().transfer_batch_to_device(batch_wo_jpeg, device, dataloader_idx) + moved["IMAGES"] = self.decode_batch_jpeg(images_jpeg, device=device) + return moved + + return super().transfer_batch_to_device(batch, device, dataloader_idx) + + + def decode_batch_jpeg( + self, + images_jpeg: list[list[torch.Tensor]], + device: torch.device | None = None, + ) -> list[torch.Tensor]: + decode_device = self.device if device is None else device + # Flatten cameras + flat_encoded, cam_sizes = [], [] + for cam in images_jpeg: + cam_sizes.append(len(cam)) + for jpg in cam: + t = jpg if isinstance(jpg, torch.Tensor) else torch.frombuffer(memoryview(jpg), dtype=torch.uint8) + # decode_jpeg requires the raw jpeg bytes to be on cpu + if t.device.type != "cpu": + t = t.cpu() + flat_encoded.append(t) + + flat_decoded = torchvision.io.decode_jpeg( + flat_encoded, + mode=torchvision.io.ImageReadMode.UNCHANGED, + device=decode_device, + ) # list of (C, H, W) gpu tensors + + out = [] + idx = 0 + for n in cam_sizes: + cam_list = flat_decoded[idx: idx+n] + idx += n + out.append(torch.stack(cam_list, dim=0)) # (B, C, H, W) + return out + + def on_fit_start(self) -> None: + super().on_fit_start() + self.depth_loss = DepthLoss(self.device) + + # ---- Metrics ---- + def ade_loss(self, pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor: + """ + Average Displacement Error -> L2 Norm -> Average Euclidean Distance between predicted and ground truth future trajectory + """ + return torch.mean(torch.norm(pred - gt, dim=-1)) + + def time_thresholds(self, t_idx): + # Time-based thresholds at 3s and 5s. + lat = torch.where(t_idx <= 3, 1.0, 1.8) + lng = torch.where(t_idx <= 3, 4.0, 7.2) + return lat, lng + + def speed_scale(self, v): + # Speed-based scaling copied from RFS paper. + return torch.where( + v < 1.4, + 0.5, + torch.where( + v < 11.0, + 0.5 + 0.5 * (v - 1.4) / (11.0 - 1.4), + 1.0 + ) + ) + + def compute_direction(self, trajectory): + # Pad with first point so displacement stays (B, T, 2). + padded = torch.cat([trajectory[:, :1], trajectory], dim=1) + displacement = padded[:, 1:] - padded[:, :-1] + lng_dir = F.normalize(displacement, p=2, dim=-1, eps=1e-6) + lat_dir = torch.stack([-lng_dir[..., 1], lng_dir[..., 0]], dim=-1) + return lng_dir, lat_dir + + def rfs_loss(self, pred, gt, lng_dir, lat_dir, speed, t_idx): + """ + pred, gt: (B, T, 2) + speed: (B,) or (B, T) + t_idx: (T,) or (B, T) + """ + delta = pred - gt + delta_lng = (delta * lng_dir).sum(dim=-1).abs() + delta_lat = (delta * lat_dir).sum(dim=-1).abs() + + tau_lat_raw, tau_lng_raw = self.time_thresholds(t_idx) + scale = self.speed_scale(speed) + if scale.dim() == 1: + scale = scale.unsqueeze(1) + + tau_lat = tau_lat_raw * scale + tau_lng = tau_lng_raw * scale + + deviation = torch.max( + delta_lat / tau_lat, + delta_lng / tau_lng, + ) + score = torch.where( + deviation <= 1, + torch.ones_like(deviation), + torch.pow(0.1, deviation - 1) + ) + return (1.0 - score).mean() + + def _prepare_rfs_inputs(self, past, future, pred_future): + speed = torch.norm(past[..., 2:4], dim=-1)[:, -1] # (B,), speed at last observed time step + full_lng_dir, full_lat_dir = self.compute_direction(future) + indices = [11, 19] # 3s and 5s into the future + + pred_slice = pred_future[:, indices, :] + gt_slice = future[:, indices, :] + lng_dir_slice = full_lng_dir[:, indices, :] + lat_dir_slice = full_lat_dir[:, indices, :] + t_idx = torch.tensor([3.0, 5.0], device=future.device).unsqueeze(0).expand(future.size(0), -1) + return pred_slice, gt_slice, lng_dir_slice, lat_dir_slice, speed, t_idx + + # ---- optimizers ---- + def configure_optimizers(self): + # NOTE: This can be extended and tuned, LR especially will differ and have an impact. + # vision encoder, if trainable, should have 1/10 the LR of the rest of the model + if hasattr(self.model, "features"): + encoder_params = [p for p in self.model.features.parameters() if p.requires_grad] + other_params = [ + p for n, p in self.model.named_parameters() + if not n.startswith("features.") and p.requires_grad + ] + if encoder_params: + encoder_lr = self.hparams.lr * 0.1 if self.hparams.lr_vision is None else self.hparams.lr_vision + return torch.optim.Adam( + [ + {"params": other_params, "lr": self.hparams.lr}, + {"params": encoder_params, "lr": encoder_lr}, + ] + ) + if other_params: + return torch.optim.Adam(other_params, lr=self.hparams.lr) + + optimizer = torch.optim.Adam(self.model.parameters(), lr=self.hparams.lr) + return optimizer + + # ---- forward / step ---- + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.model(x) + + def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor: + past, future, intent = batch['PAST'], batch['FUTURE'], batch['INTENT'] + + if "IMAGES" in batch: + images = batch["IMAGES"] + elif "IMAGES_JPEG" in batch: + images_jpeg = batch["IMAGES_JPEG"] + images = self.decode_batch_jpeg(images_jpeg) + else: + raise KeyError("Batch must contain either 'IMAGES_JPEG' or 'IMAGES' key.") + + # `past` is our input (B, 16, 6) e.g. Batch x Time x (x, y, v_x, v_y, a_x, a_y) + # and `future` is our output (B, 20, 2) e.g. Batch x Time x (x, y) + + # create all input data that we are allowed to give to a model + model_inputs = {'PAST': past, 'IMAGES': images, 'INTENT': intent} + + pred_future = self.forward(model_inputs) # (B, T*2) + pred_depth = None + pred_scores: torch.Tensor = None + pred_occ = None + + if isinstance(pred_future, dict): + pred_depth = pred_future.get("depth", None) + pred_scores = pred_future.get("scores", None) + pred_occ = pred_future.get("occ", None) + pred_future = pred_future["trajectory"] + + pred = pred_future + t_steps = future.shape[1] + t2 = t_steps * 2 + + if pred.ndim != 2: + raise ValueError(f"Unexpected pred shape {pred.shape}; expected 2D (B, K*T*2).") + + if pred.shape[1] % t2 != 0: + raise ValueError(f"pred dim1={pred.shape[1]} is not divisible by T*2={t2}.") + k_modes = pred.shape[1] // t2 + pred = pred.view(pred.size(0), k_modes, t_steps, 2) + + if pred_scores is not None and pred.size(1) > 1: + rfs_pred_idx = pred_scores.argmin(dim=1) + else: + rfs_pred_idx = torch.zeros(pred.size(0), dtype=torch.long, device=pred.device) + pred_for_rfs = pred[torch.arange(pred.size(0), device=pred.device), rfs_pred_idx] + + pred_slice, gt_slice, lng_dir_slice, lat_dir_slice, speed, t_idx = self._prepare_rfs_inputs( + past, + future, + pred_for_rfs, + ) + rfs_unweighted = self.rfs_loss(pred_slice, gt_slice, lng_dir_slice, lat_dir_slice, speed, t_idx) + rfs_weight = getattr(self.hparams, "rfs_weight", 0.0) + loss_rfs = rfs_weight * rfs_unweighted + + loss_type = getattr(self.hparams, "model_cfg_loss_type", "mse") + + # ADE per mode: (B, K) + dist = torch.norm(pred - future[:, torch.newaxis, :, :], dim=-1) # (B, K, T) + ade_per_mode = dist.mean(dim=-1) + + # Top-M WTA for trajectory loss. Here, we have an "oracle" that picks the best mode + # so, our loss is calculated on the mean of the top n trajectories. + top_m = min(getattr(self.hparams, "model_cfg_loss_top_n", 5), ade_per_mode.size(1)) + loss_ade = ade_per_mode.topk(top_m, largest=False, dim=1).values.mean() + + # oracle ade is best of all proposals, since we have the GT data during training + oracle_ade = ade_per_mode.min(dim=1).values.mean() + ade_pred = None + # pred_scores is now the predicted ADE of each trajectory / expectation loss + if pred_scores is not None and k_modes > 1: + pred_idx = pred_scores.argmin(dim=1) + ade_pred = ade_per_mode[torch.arange(pred.size(0), device=pred.device), pred_idx].mean() + elif k_modes == 1: + ade_pred = ade_per_mode.squeeze(1).mean() + regret = (ade_pred - oracle_ade) if ade_pred is not None else None + + # Scorer Losses -> encourage ranking of predicted scores to match true ranking of ades that are generated + if k_modes > 1 and pred_scores is not None: + ade = ade_per_mode.detach() # (B, K) + if loss_type == "mse": + loss_score = F.mse_loss(pred_scores, ade) + elif loss_type == "reinforce": + tau_base = getattr(self.hparams, "model_cfg_loss_tau_base", 1.0) + decay_factor = getattr(self.hparams, "model_cfg_loss_tau_decay", 0.95) + entropy_lambda = getattr(self.hparams, "model_cfg_loss_entropy_lambda", 0.01) + # p_k = softmax(s_k / tau) + logits = -pred_scores / max(0.1, tau_base * (decay_factor ** self.current_epoch)) + p = F.softmax(logits, dim=1) + # Loss = sum_k {p_k * e_k} (e.g., expectation of e_k) + loss_selection = (p * ade).sum(dim=1).mean() + # H(p) = -sum(p) * log(p) + entropy = -(p * (p + 1e-8).log()).sum(dim=1).mean() + loss_score = loss_selection - entropy_lambda * entropy + else: + raise NotImplementedError(f"Loss {loss_type} is not implemented") + pred_idx = pred_scores.argmin(dim=1) + ade_pred = ade_per_mode[torch.arange(pred.size(0), device=pred.device), pred_idx].mean() + else: + loss_score = torch.tensor(0.0, device=self.device) + + # Scorer Metrics + scorer_metrics = {} + if k_modes > 1 and pred_scores is not None: + # Rank of the scorer's top-1 pick among oracle-sorted proposals + oracle_ranking = ade_per_mode.argsort(dim=1) # (B, K) indices sorted by true ADE + oracle_rank_of = oracle_ranking.argsort(dim=1) # (B, K) rank of each proposal + scorer_pick = pred_scores.argmin(dim=1) # (B,) scorer's best + pick_rank = oracle_rank_of[torch.arange(pred.size(0), device=pred.device), scorer_pick].float() # (B,) + scorer_metrics[f"{stage}_scorer_mean_rank"] = pick_rank.mean() + for topk in (1, 5, 10): + if topk <= k_modes: + scorer_metrics[f"{stage}_scorer_top{topk}_acc"] = (pick_rank < topk).float().mean() + # Spearman rank correlation + K = ade_per_mode.size(1) + oracle_ranks = oracle_rank_of.float() # (B, K) + scorer_ranks = pred_scores.argsort(dim=1).argsort(dim=1).float() # (B, K) + d = oracle_ranks - scorer_ranks + rho = 1 - 6 * (d ** 2).sum(dim=1) / (K * (K ** 2 - 1)) + scorer_metrics[f"{stage}_scorer_spearman"] = rho.mean() + + # Depth Loss + if pred_depth is not None: + front_img = images[1] # front camera + depth_in = F.interpolate(front_img, size=(128, 128), mode='nearest') + loss_depth = self.depth_loss(depth_in, pred_depth, loss_fn=F.l1_loss) + else: + loss_depth = torch.tensor(0.0, device=self.device) + + if pred_occ is not None and "OCC" in batch: + occ_gt = batch["OCC"].long() # (B, 100, 100, 16) + loss_occ = F.cross_entropy(pred_occ, occ_gt, ignore_index=255) + else: + loss_occ = torch.tensor(0.0, device=self.device) + + loss_depth *= 0.1 # slightly enabled + loss_occ *= 0.1 + loss_ade *= 1.0 # TODO: tune loss terms + loss_score *= 1.0 + total_loss = loss_ade + loss_depth + loss_score + loss_rfs + loss_occ + # TODO: improve logging both to disk and to console + log_payload = { + f"{stage}_loss_ade": loss_ade, + f"{stage}_loss_score": loss_score, + f"{stage}_loss_depth": loss_depth, + f"{stage}_loss_rfs": loss_rfs, + f"{stage}_rfs_unweighted": rfs_unweighted, + f"{stage}_loss": total_loss, + f"{stage}_loss_occ": loss_occ, + } + if ade_pred is not None: + log_payload[f"{stage}_ade_pred"] = ade_pred + log_payload[f"{stage}_ade_oracle"] = oracle_ade + log_payload[f"{stage}_ade_regret"] = regret + log_payload.update(scorer_metrics) + self.log_dict(log_payload, prog_bar=True, logger=True, + batch_size=past.size(0), + sync_dist=(stage == "val")) + + return total_loss + + def training_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor: + return self._shared_step(batch, "train") + + def validation_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor: + return self._shared_step(batch, "val") + +def collate_with_images(batch): + past = [torch.as_tensor(b["PAST"], dtype=torch.float32) for b in batch] + future = [torch.as_tensor(b["FUTURE"], dtype=torch.float32) for b in batch] + intent = torch.as_tensor([b["INTENT"] for b in batch]) + names = [b["NAME"] for b in batch] + + cams = list(zip(*[b["IMAGES_JPEG"] for b in batch])) # per-camera tuples + images_jpeg = [list(cam_imgs) for cam_imgs in cams] # stay on CPU + + out = { + "PAST": torch.stack(past, dim=0), + "FUTURE": torch.stack(future, dim=0), + "INTENT": intent, + "IMAGES_JPEG": images_jpeg, + "NAME": names, + } + + if batch[0].get("OCC", None) is not None: + out["OCC"] = torch.stack( + [torch.as_tensor(b["OCC"], dtype=torch.long) for b in batch], + dim=0, + ) # (B, 100, 100, 16) + + return out diff --git a/src/camera-based-e2e/models/monocular_baseline.py b/src/camera-based-e2e/models/monocular_baseline.py new file mode 100644 index 0000000..c310027 --- /dev/null +++ b/src/camera-based-e2e/models/monocular_baseline.py @@ -0,0 +1,231 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from math import sqrt + +from .blocks import TransformerBlock + +class MonocularModel(nn.Module): + def __init__( + self, + in_dim: int, + out_dim: int, + feature_extractor: nn.Module + ): + # out_dim: (B, 40) which gets reshaped to (B, 20, 2) later + super(MonocularModel, self).__init__() + self.features = feature_extractor + + # attention + self.feature_dim = sum(self.features.dims) # works for both DINO and SAM + self.key_projection = nn.Linear(in_features=self.feature_dim, out_features=self.feature_dim) # project into "key" space + self.value_projection = nn.Linear(in_features=self.feature_dim, out_features=self.feature_dim) + + # condition the query on intent (B,) and past (B, 16, 6) + query_input_dim = 3 + 16 * 6 # one hot -- concat -- flattened + self.query = nn.Sequential( + nn.Linear(query_input_dim, self.feature_dim), + nn.LeakyReLU(), + nn.Linear(self.feature_dim, self.feature_dim), + ) + + # learnable positional encoding + self.n_tokens = self.features.data_config["input_size"][1] // self.features.patch_size * (self.features.data_config["input_size"][2] // self.features.patch_size) + self.positional_encoding = nn.Parameter(nn.init.trunc_normal_(torch.zeros((1, self.n_tokens, self.feature_dim)), std=0.02)) # (1, N, C) + + # MLP at end rather than directly using softmax as final output + self.decoder = nn.Sequential( + nn.Linear(self.feature_dim, self.feature_dim), + nn.LeakyReLU(), + nn.Linear(self.feature_dim, out_dim), + ) + + # LayerNorms + self.token_norm = nn.LayerNorm(self.feature_dim) + self.query_norm = nn.LayerNorm(self.feature_dim) + self.attn_norm = nn.LayerNorm(self.feature_dim) + + + def forward(self, x: dict) -> torch.Tensor: + # past: (B, 16, 6), intent: int + past, images, intent = x['PAST'], x['IMAGES'], x['INTENT'] + + # Ref: https://github.com/waymo-research/waymo-open-dataset/blob/5f8a1cd42491210e7de629b6f8fc09b65e0cbe99/src/waymo_open_dataset/dataset.proto#L50%20%20order%20=%20[2,%201,%203] + front_cam = images[1] + with torch.no_grad(): + feats = self.features(front_cam) # list or tensor + + # tokens: handle list of features or single tensor + if isinstance(feats, (list, tuple)): + tokens = torch.cat([f.flatten(2) for f in feats], dim=1) # (B, C_total, N) + else: + tokens = feats.flatten(2) # (B, C, N) + tokens = torch.permute(tokens, (0, 2, 1)) + self.positional_encoding # (B, N, C_total) + tokens = self.token_norm(tokens) + + # attention + key = self.key_projection(tokens) # (B, 256, 1152) + value = self.value_projection(tokens) # (B, 256, 40) + + intent_onehot = F.one_hot((intent - 1).long(), num_classes=3).float() # (B, 3). minus 1 --> 0, 1, 2 + past_flat = past.view(past.size(0), -1) # (B, 96) + query = self.query(torch.cat([intent_onehot, past_flat], dim=1)).unsqueeze(1) # (B, 1, 256) + query = self.query_norm(query) + + scores = query @ key.permute((0, 2, 1)) # (B, T, N) + attention = F.softmax(scores / sqrt(key.shape[2]), dim=2) @ value # (B, 1, 40) + attention = self.attn_norm(attention) + return self.decoder(attention.squeeze(1)) # (B, 40) + +class DeepMonocularModel(nn.Module): + def __init__( + self, + feature_extractor, + out_dim, + n_blocks=1, + n_proposals=50, + dt: float = 0.25, + max_accel: float = 8.0, + max_omega: float = 1.0, + ): + super().__init__() + self.features = feature_extractor + self.feature_dim = sum(self.features.dims) + if out_dim % 2 != 0: + raise ValueError(f"out_dim must be even for (x,y) rollout, got {out_dim}") + self.horizon = out_dim // 2 + self.dt = dt + self.max_accel = max_accel + self.max_omega = max_omega + + # Initial Query Projection (Intent + Past -> C) + query_input_dim = 3 + 16 * 6 + self.query_init = nn.Linear(query_input_dim, self.feature_dim) + + # Instead of fine-tuning feature extractor, project w/ conv + self.visual_adapter = nn.Sequential( + nn.Conv2d(self.feature_dim, self.feature_dim, 3, padding=1), + nn.GELU(), + nn.Conv2d(self.feature_dim, self.feature_dim, 3, padding=1), + ) + + # learnable positional encoding + self.n_tokens = self.features.data_config["input_size"][1] // self.features.patch_size * (self.features.data_config["input_size"][2] // self.features.patch_size) + self.positional_encoding = nn.Parameter(nn.init.trunc_normal_(torch.zeros((1, self.n_tokens, self.feature_dim)), std=0.02)) # (1, N, C) + + # Deep network rather than single attention in MonocularModel + self.blocks = nn.ModuleList([ + TransformerBlock(self.feature_dim, num_heads=8, mlp_dim=self.feature_dim*4) + for _ in range(n_blocks) + ]) + + # For Supervised Depth Loss -> (B, 128, 128) + self.depth_gen = nn.Sequential( + nn.Conv2d(self.feature_dim, 64, 3, padding=1), + nn.GELU(), + nn.Upsample(scale_factor=2, mode='nearest'), + nn.Conv2d(64, 32, 3, padding=1), + nn.GELU(), + nn.Upsample(scale_factor=2, mode='nearest'), + nn.Conv2d(32, 1, 1) + ) + + self.n_proposals = n_proposals + self.traj_decoder = nn.Sequential( + nn.Linear(self.feature_dim, self.feature_dim), + nn.GELU(), + nn.Linear(self.feature_dim, self.feature_dim), + nn.GELU(), + nn.Linear(self.feature_dim, out_dim * self.n_proposals), + ) + self.traj_features = nn.Sequential( + nn.Linear(out_dim, self.feature_dim), + nn.GELU(), + nn.Linear(self.feature_dim, self.feature_dim), + nn.GELU(), + ) + self.score_decoder = nn.Sequential( + nn.Linear(self.feature_dim * 2, self.feature_dim), + nn.GELU(), + nn.Linear(self.feature_dim, self.feature_dim), + nn.GELU(), + nn.Linear(self.feature_dim, 1), + ) # no softmax, since we use cross entropy later + + def bicycle_model(self, control_pred: torch.Tensor, past: torch.Tensor) -> torch.Tensor: + accel = torch.tanh(control_pred[..., 0]) * self.max_accel # (B, K, T) + omega = torch.tanh(control_pred[..., 1]) * self.max_omega # (B, K, T) + + x_state = past[:, -1, 0].unsqueeze(1).expand(-1, self.n_proposals).clone() + y_state = past[:, -1, 1].unsqueeze(1).expand(-1, self.n_proposals).clone() + vx0 = past[:, -1, 2] + vy0 = past[:, -1, 3] + speed_state = torch.sqrt(vx0 * vx0 + vy0 * vy0 + 1e-6).unsqueeze(1).expand(-1, self.n_proposals).clone() + heading_state = torch.atan2(vy0, vx0).unsqueeze(1).expand(-1, self.n_proposals).clone() + + xy_steps = [] + for t in range(self.horizon): + x_state = x_state + speed_state * torch.cos(heading_state) * self.dt + y_state = y_state + speed_state * torch.sin(heading_state) * self.dt + xy_steps.append(torch.stack([x_state, y_state], dim=-1)) + + heading_state = heading_state + omega[:, :, t] * self.dt + speed_state = torch.clamp_min(speed_state + accel[:, :, t] * self.dt, 0.0) + + traj_xy = torch.stack(xy_steps, dim=2) # (B, K, T, 2) + return traj_xy, traj_xy.reshape(traj_xy.size(0), -1), accel, omega # (B, K*T*2) + + def forward(self, x): + # Copied from MonocularModel + # past: (B, 16, 6), intent: int + past, images, intent = x['PAST'], x['IMAGES'], x['INTENT'] + + # Ref: https://github.com/waymo-research/waymo-open-dataset/blob/5f8a1cd42491210e7de629b6f8fc09b65e0cbe99/src/waymo_open_dataset/dataset.proto#L50%20%20order%20=%20[2,%201,%203] + front_cam = images[1] + + # Doesn't need no_grad b/c DINO/SAMFeatures will freeze if needed + feats_vit = self.features(front_cam) # list or tensor + + if len(feats_vit) == 1 and isinstance(feats_vit, list): + feats_vit = feats_vit[0] + + feats = self.visual_adapter(feats_vit) # (B, C, H, W) + + # Depth Supervision + output_depth = F.softplus(self.depth_gen(feats).squeeze(1)) # (B, 128, 128) + + # tokens: handle list of features or single tensor + # TODO: is this made redundant by if statement above? + if isinstance(feats, (list, tuple)): + tokens = torch.cat([f.flatten(2) for f in feats], dim=1) # (B, C_total, N) + else: + tokens = feats.flatten(2) # (B, C, N) + tokens = torch.permute(tokens, (0, 2, 1)) + self.positional_encoding # (B, N, C_total) + + # copy procedure to build query_0 from MonocularModel + intent_onehot = F.one_hot((intent - 1).long(), num_classes=3).float() + past_flat = past.view(past.size(0), -1) + query: torch.Tensor = self.query_init(torch.cat([intent_onehot, past_flat], dim=1)).unsqueeze(1) + + for block in self.blocks: + query = block(query, tokens) + + # predict (acceleration, angular velocity) for each timestep + # and roll it out using the kinematic bicycle model + control_pred = self.traj_decoder(query.squeeze(1)).view( + query.size(0), self.n_proposals, self.horizon, 2 + ) # (B, K, T, 2) + traj_xy, traj_pred, accel, omega = self.bicycle_model(control_pred, past) # (B, K, T*2) + + traj_pred_flat = traj_xy.reshape(traj_xy.size(0), self.n_proposals, -1) # (B, K, T*2) + traj_feat: torch.Tensor = self.traj_features(traj_pred_flat.detach()) # (B, K, C) + query_for_score = query.squeeze(1).detach()[:, torch.newaxis, :].expand(-1, self.n_proposals, -1) # (B, K, C) + score_in = torch.cat([query_for_score, traj_feat], dim=-1) # (B, K, 2C) + score_pred = self.score_decoder(score_in).squeeze(-1) # (B, K) + + return { + "trajectory": traj_pred, + "scores": score_pred, + "depth": output_depth, + "controls": torch.stack([accel, omega], dim=-1).reshape(query.size(0), -1), + } diff --git a/src/camera-based-e2e/models/monocular_occ.py b/src/camera-based-e2e/models/monocular_occ.py new file mode 100644 index 0000000..50a9484 --- /dev/null +++ b/src/camera-based-e2e/models/monocular_occ.py @@ -0,0 +1,242 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from math import sqrt + +from .blocks import TransformerBlock + +class MonocularModel(nn.Module): + def __init__( + self, + in_dim: int, + out_dim: int, + feature_extractor: nn.Module + ): + # out_dim: (B, 40) which gets reshaped to (B, 20, 2) later + super(MonocularModel, self).__init__() + self.features = feature_extractor + + # attention + self.feature_dim = sum(self.features.dims) # works for both DINO and SAM + self.key_projection = nn.Linear(in_features=self.feature_dim, out_features=self.feature_dim) # project into "key" space + self.value_projection = nn.Linear(in_features=self.feature_dim, out_features=self.feature_dim) + + # condition the query on intent (B,) and past (B, 16, 6) + query_input_dim = 3 + 16 * 6 # one hot -- concat -- flattened + self.query = nn.Sequential( + nn.Linear(query_input_dim, self.feature_dim), + nn.LeakyReLU(), + nn.Linear(self.feature_dim, self.feature_dim), + ) + + # learnable positional encoding + self.n_tokens = self.features.data_config["input_size"][1] // self.features.patch_size * (self.features.data_config["input_size"][2] // self.features.patch_size) + self.positional_encoding = nn.Parameter(nn.init.trunc_normal_(torch.zeros((1, self.n_tokens, self.feature_dim)), std=0.02)) # (1, N, C) + + # MLP at end rather than directly using softmax as final output + self.decoder = nn.Sequential( + nn.Linear(self.feature_dim, self.feature_dim), + nn.LeakyReLU(), + nn.Linear(self.feature_dim, out_dim), + ) + + # LayerNorms + self.token_norm = nn.LayerNorm(self.feature_dim) + self.query_norm = nn.LayerNorm(self.feature_dim) + self.attn_norm = nn.LayerNorm(self.feature_dim) + + + def forward(self, x: dict) -> torch.Tensor: + # past: (B, 16, 6), intent: int + past, images, intent = x['PAST'], x['IMAGES'], x['INTENT'] + + # Ref: https://github.com/waymo-research/waymo-open-dataset/blob/5f8a1cd42491210e7de629b6f8fc09b65e0cbe99/src/waymo_open_dataset/dataset.proto#L50%20%20order%20=%20[2,%201,%203] + front_cam = images[1] + with torch.no_grad(): + feats = self.features(front_cam) # list or tensor + + # tokens: handle list of features or single tensor + if isinstance(feats, (list, tuple)): + tokens = torch.cat([f.flatten(2) for f in feats], dim=1) # (B, C_total, N) + else: + tokens = feats.flatten(2) # (B, C, N) + tokens = torch.permute(tokens, (0, 2, 1)) + self.positional_encoding # (B, N, C_total) + tokens = self.token_norm(tokens) + + # attention + key = self.key_projection(tokens) # (B, 256, 1152) + value = self.value_projection(tokens) # (B, 256, 40) + + intent_onehot = F.one_hot((intent - 1).long(), num_classes=3).float() # (B, 3). minus 1 --> 0, 1, 2 + past_flat = past.view(past.size(0), -1) # (B, 96) + query = self.query(torch.cat([intent_onehot, past_flat], dim=1)).unsqueeze(1) # (B, 1, 256) + query = self.query_norm(query) + + scores = query @ key.permute((0, 2, 1)) # (B, T, N) + attention = F.softmax(scores / sqrt(key.shape[2]), dim=2) @ value # (B, 1, 40) + attention = self.attn_norm(attention) + return self.decoder(attention.squeeze(1)) # (B, 40) + +class DeepMonocularModel(nn.Module): + def __init__( + self, + feature_extractor, + out_dim, + n_blocks=1, + n_proposals=50, + dt: float = 0.25, + max_accel: float = 8.0, + max_omega: float = 1.0, + ): + super().__init__() + self.features = feature_extractor + self.feature_dim = sum(self.features.dims) + if out_dim % 2 != 0: + raise ValueError(f"out_dim must be even for (x,y) rollout, got {out_dim}") + self.horizon = out_dim // 2 + self.dt = dt + self.max_accel = max_accel + self.max_omega = max_omega + + # Initial Query Projection (Intent + Past -> C) + query_input_dim = 3 + 16 * 6 + self.query_init = nn.Linear(query_input_dim, self.feature_dim) + + # Instead of fine-tuning feature extractor, project w/ conv + self.visual_adapter = nn.Sequential( + nn.Conv2d(self.feature_dim, self.feature_dim, 3, padding=1), + nn.GELU(), + nn.Conv2d(self.feature_dim, self.feature_dim, 3, padding=1), + ) + + # learnable positional encoding + self.n_tokens = self.features.data_config["input_size"][1] // self.features.patch_size * (self.features.data_config["input_size"][2] // self.features.patch_size) + self.positional_encoding = nn.Parameter(nn.init.trunc_normal_(torch.zeros((1, self.n_tokens, self.feature_dim)), std=0.02)) # (1, N, C) + + # Deep network rather than single attention in MonocularModel + self.blocks = nn.ModuleList([ + TransformerBlock(self.feature_dim, num_heads=8, mlp_dim=self.feature_dim*4) + for _ in range(n_blocks) + ]) + + # For Supervised Depth Loss -> (B, 128, 128) + self.depth_gen = nn.Sequential( + nn.Conv2d(self.feature_dim, 64, 3, padding=1), + nn.GELU(), + nn.Upsample(scale_factor=2, mode='nearest'), + nn.Conv2d(64, 32, 3, padding=1), + nn.GELU(), + nn.Upsample(scale_factor=2, mode='nearest'), + nn.Conv2d(32, 1, 1) + ) + self.occ_head = nn.Sequential( + nn.Conv2d(self.feature_dim, 256, 3, padding=1), + nn.GELU(), + nn.Conv2d(256, 16 * 6, 1), + ) + + self.n_proposals = n_proposals + self.traj_decoder = nn.Sequential( + nn.Linear(self.feature_dim, self.feature_dim), + nn.GELU(), + nn.Linear(self.feature_dim, self.feature_dim), + nn.GELU(), + nn.Linear(self.feature_dim, out_dim * self.n_proposals), + ) + self.traj_features = nn.Sequential( + nn.Linear(out_dim, self.feature_dim), + nn.GELU(), + nn.Linear(self.feature_dim, self.feature_dim), + nn.GELU(), + ) + self.score_decoder = nn.Sequential( + nn.Linear(self.feature_dim * 2, self.feature_dim), + nn.GELU(), + nn.Linear(self.feature_dim, self.feature_dim), + nn.GELU(), + nn.Linear(self.feature_dim, 1), + ) # no softmax, since we use cross entropy later + + def bicycle_model(self, control_pred: torch.Tensor, past: torch.Tensor) -> torch.Tensor: + accel = torch.tanh(control_pred[..., 0]) * self.max_accel # (B, K, T) + omega = torch.tanh(control_pred[..., 1]) * self.max_omega # (B, K, T) + + x_state = past[:, -1, 0].unsqueeze(1).expand(-1, self.n_proposals).clone() + y_state = past[:, -1, 1].unsqueeze(1).expand(-1, self.n_proposals).clone() + vx0 = past[:, -1, 2] + vy0 = past[:, -1, 3] + speed_state = torch.sqrt(vx0 * vx0 + vy0 * vy0 + 1e-6).unsqueeze(1).expand(-1, self.n_proposals).clone() + heading_state = torch.atan2(vy0, vx0).unsqueeze(1).expand(-1, self.n_proposals).clone() + + xy_steps = [] + for t in range(self.horizon): + x_state = x_state + speed_state * torch.cos(heading_state) * self.dt + y_state = y_state + speed_state * torch.sin(heading_state) * self.dt + xy_steps.append(torch.stack([x_state, y_state], dim=-1)) + + heading_state = heading_state + omega[:, :, t] * self.dt + speed_state = torch.clamp_min(speed_state + accel[:, :, t] * self.dt, 0.0) + + traj_xy = torch.stack(xy_steps, dim=2) # (B, K, T, 2) + return traj_xy, traj_xy.reshape(traj_xy.size(0), -1), accel, omega # (B, K*T*2) + + def forward(self, x): + # Copied from MonocularModel + # past: (B, 16, 6), intent: int + past, images, intent = x['PAST'], x['IMAGES'], x['INTENT'] + + # Ref: https://github.com/waymo-research/waymo-open-dataset/blob/5f8a1cd42491210e7de629b6f8fc09b65e0cbe99/src/waymo_open_dataset/dataset.proto#L50%20%20order%20=%20[2,%201,%203] + front_cam = images[1] + + # Doesn't need no_grad b/c DINO/SAMFeatures will freeze if needed + feats_vit = self.features(front_cam) # list or tensor + + if len(feats_vit) == 1 and isinstance(feats_vit, list): + feats_vit = feats_vit[0] + + feats = self.visual_adapter(feats_vit) # (B, C, H, W) + + # Depth Supervision + output_depth = F.softplus(self.depth_gen(feats).squeeze(1)) # (B, 128, 128) + occ_logits = self.occ_head(feats) # (B, 96, h, w) + occ_logits = F.interpolate(occ_logits, size=(100, 100), mode="bilinear", align_corners=False) + + B = occ_logits.size(0) + occ_logits = occ_logits.view(B, 6, 16, 100, 100) + occ_logits = occ_logits.permute(0, 1, 3, 4, 2).contiguous() # (B, 6, 100, 100, 16) + # tokens: handle list of features or single tensor + # TODO: is this made redundant by if statement above? + if isinstance(feats, (list, tuple)): + tokens = torch.cat([f.flatten(2) for f in feats], dim=1) # (B, C_total, N) + else: + tokens = feats.flatten(2) # (B, C, N) + tokens = torch.permute(tokens, (0, 2, 1)) + self.positional_encoding # (B, N, C_total) + + # copy procedure to build query_0 from MonocularModel + intent_onehot = F.one_hot((intent - 1).long(), num_classes=3).float() + past_flat = past.view(past.size(0), -1) + query: torch.Tensor = self.query_init(torch.cat([intent_onehot, past_flat], dim=1)).unsqueeze(1) + + for block in self.blocks: + query = block(query, tokens) + + # predict (acceleration, angular velocity) for each timestep + # and roll it out using the kinematic bicycle model + control_pred = self.traj_decoder(query.squeeze(1)).view( + query.size(0), self.n_proposals, self.horizon, 2 + ) # (B, K, T, 2) + traj_xy, traj_pred, accel, omega = self.bicycle_model(control_pred, past) # (B, K, T*2) + + traj_pred_flat = traj_xy.reshape(traj_xy.size(0), self.n_proposals, -1) # (B, K, T*2) + traj_feat: torch.Tensor = self.traj_features(traj_pred_flat.detach()) # (B, K, C) + query_for_score = query.squeeze(1).detach()[:, torch.newaxis, :].expand(-1, self.n_proposals, -1) # (B, K, C) + score_in = torch.cat([query_for_score, traj_feat], dim=-1) # (B, K, 2C) + score_pred = self.score_decoder(score_in).squeeze(-1) # (B, K) + + return { + "trajectory": traj_pred, + "scores": score_pred, + "depth": output_depth, + "occ": occ_logits, + "controls": torch.stack([accel, omega], dim=-1).reshape(query.size(0), -1), + } diff --git a/src/camera-based-e2e/train_baseline.py b/src/camera-based-e2e/train_baseline.py new file mode 100644 index 0000000..6140455 --- /dev/null +++ b/src/camera-based-e2e/train_baseline.py @@ -0,0 +1,330 @@ +import argparse +import itertools +import math +import random + +import pytorch_lightning as pl +from pytorch_lightning.callbacks import ModelCheckpoint +from pytorch_lightning.profilers import SimpleProfiler + +import torch +from pathlib import Path +import os +from torch.utils.data import BatchSampler + + +from models.base_model_baseline import LitModel, collate_with_images +from models.monocular_baseline import DeepMonocularModel +from models.feature_extractors import SAMFeatures + + +class HomogeneousConcatBatchSampler(BatchSampler): + """Emit batches from one ConcatDataset source at a time. + + Designed for ConcatDataset([waymo, nuscenes]) so every batch is + from one source. Works with DDP + """ + + def __init__( + self, + dataset_lengths: tuple[int, int], + batch_size: int, + rank: int | None = None, + world_size: int | None = None, + drop_last: bool = False, + shuffle: bool = True, + seed: int = 42, + source_ratio: tuple[int, int] = (1, 1), + **kwargs, + ): + if len(dataset_lengths) != 2: + raise ValueError(f"Expected exactly 2 datasets, got {len(dataset_lengths)}") + if batch_size <= 0: + raise ValueError(f"batch_size must be > 0, got {batch_size}") + if rank is None: + rank = int(os.environ.get("RANK", "0")) + if world_size is None: + world_size = int(os.environ.get("WORLD_SIZE", "1")) + + if world_size <= 0: + raise ValueError(f"world_size must be > 0, got {world_size}") + if rank < 0 or rank >= world_size: + raise ValueError( + f"Invalid rank/world_size pair: rank={rank}, world_size={world_size}" + ) + + self.lengths = dataset_lengths + self.batch_size = batch_size + self.rank = rank + self.world_size = world_size + self.drop_last = drop_last + self.shuffle = shuffle + self.seed = seed + self.epoch = 0 + self.source_ratio = source_ratio + + self.offset0 = 0 + self.offset1 = dataset_lengths[0] + + def set_epoch(self, epoch: int): + self.epoch = int(epoch) + + def _num_samples_per_rank(self, length: int) -> int: + if length <= 0: + return 0 + if self.drop_last and length % self.world_size != 0: + return math.ceil((length - self.world_size) / self.world_size) + return math.ceil(length / self.world_size) + + def _make_rank_indices(self, start: int, length: int, rng: random.Random): + idx = list(range(start, start + length)) + if self.shuffle: + rng.shuffle(idx) + num_samples = self._num_samples_per_rank(length) + total_size = num_samples * self.world_size + + if self.drop_last: + idx = idx[:total_size] + elif total_size > len(idx): + if not idx: + return [] + padding_size = total_size - len(idx) + repeats = math.ceil(padding_size / len(idx)) + idx += (idx * repeats)[:padding_size] + + return idx[self.rank : total_size : self.world_size] + + def _chunk(self, indices: list[int]) -> list[list[int]]: + if self.drop_last: + n_full = len(indices) // self.batch_size + return [ + indices[i * self.batch_size : (i + 1) * self.batch_size] + for i in range(n_full) + ] + + out = [] + for i in range(0, len(indices), self.batch_size): + out.append(indices[i : i + self.batch_size]) + return out + + def __iter__(self): + rng = random.Random(self.seed + self.epoch) + + idx0 = self._make_rank_indices(self.offset0, self.lengths[0], rng) + idx1 = self._make_rank_indices(self.offset1, self.lengths[1], rng) + + b0 = self._chunk(idx0) + b1 = self._chunk(idx1) + + w0 = max(0, int(self.source_ratio[0])) + w1 = max(0, int(self.source_ratio[1])) + if w0 == 0 and w1 == 0: + w0, w1 = 1, 1 + + pattern = [0] * w0 + [1] * w1 + if not pattern: + pattern = [0, 1] + + i0, i1 = 0, 0 + for source in itertools.cycle(pattern): + if i0 >= len(b0) and i1 >= len(b1): + break + if source == 0: + if i0 < len(b0): + yield b0[i0] + i0 += 1 + else: + if i1 < len(b1): + yield b1[i1] + i1 += 1 + + def __len__(self): + def n_batches(length: int): + per_rank = self._num_samples_per_rank(length) + if self.drop_last: + return per_rank // self.batch_size + return math.ceil(per_rank / self.batch_size) + + return n_batches(self.lengths[0]) + n_batches(self.lengths[1]) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--data_dir", type=str, required=True, help="Path to data directory" + ) + parser.add_argument( + "--batch_size", type=int, default=16, help="Batch size for training" + ) + parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate") + parser.add_argument( + "--max_epochs", type=int, default=10, help="Number of epochs to train" + ) + parser.add_argument( + "--compile", + action="store_true", + help="Whether to compile the model with torch.compile", + ) + parser.add_argument( + "--profile", action="store_true", help="Whether to run the profiler" + ) + parser.add_argument( + "--dataset", + type=str, + default="waymo", + choices=["waymo", "nuscenes", "all"], + help="Which dataset to train on", + ) + args = parser.parse_args() + + pl.seed_everything(42, workers=True) + + # Data + if args.dataset == "waymo": + from loader_baseline import WaymoE2E + + train_dataset = WaymoE2E( + indexFile="index_train.pkl", data_dir=args.data_dir, n_items=250_000 + ) + test_dataset = WaymoE2E( + indexFile="index_val.pkl", data_dir=args.data_dir, n_items=20_000 + ) + nw = 0 + elif args.dataset == "nuscenes": + from nuscenes_loader import NuScenesDataset + + train_dataset = NuScenesDataset( + data_dir=args.data_dir, split="train", n_items=250_000 + ) + test_dataset = NuScenesDataset( + data_dir=args.data_dir, split="val", n_items=25_000 + ) + nw = 16 + elif args.dataset == "all": + from loader_baseline import WaymoE2E + from nuscenes_loader import NuScenesDataset + + waymo_dir, nuscenes_dir = os.getenv("WAYMO_DATA_DIR"), os.getenv( + "NUSCENES_DATA_DIR" + ) + + waymo_train = WaymoE2E( + indexFile="index_train.pkl", data_dir=waymo_dir, n_items=125_000 + ) + waymo_test = WaymoE2E( + indexFile="index_val.pkl", data_dir=waymo_dir, n_items=12_500 + ) + + nuscenes_train = NuScenesDataset( + data_dir=nuscenes_dir, split="train", n_items=125_000 + ) + nuscenes_test = NuScenesDataset( + data_dir=nuscenes_dir, split="val", n_items=12_500 + ) + + train_dataset = torch.utils.data.ConcatDataset([waymo_train, nuscenes_train]) + test_dataset = torch.utils.data.ConcatDataset([waymo_test, nuscenes_test]) + nw = 16 + + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + + train_batch_sampler = HomogeneousConcatBatchSampler( + dataset_lengths=(len(waymo_train), len(nuscenes_train)), + batch_size=args.batch_size, + rank=rank, + world_size=world_size, + drop_last=False, + shuffle=True, + seed=42, + source_ratio=(1, 1), + ) + val_batch_sampler = HomogeneousConcatBatchSampler( + dataset_lengths=(len(waymo_test), len(nuscenes_test)), + batch_size=args.batch_size, + rank=rank, + world_size=world_size, + drop_last=False, + shuffle=False, + seed=42, + source_ratio=(1, 1), + ) + else: + raise ValueError(f"Unsupported dataset: {args.dataset}") + + if args.dataset == "all": + train_loader = torch.utils.data.DataLoader( + train_dataset, + num_workers=4, + batch_sampler=train_batch_sampler, + collate_fn=collate_with_images, + persistent_workers=False, + pin_memory=False, + ) + val_loader = torch.utils.data.DataLoader( + test_dataset, + num_workers=4, + batch_sampler=val_batch_sampler, + collate_fn=collate_with_images, + persistent_workers=False, + pin_memory=False, + ) + else: + train_loader = torch.utils.data.DataLoader( + train_dataset, + batch_size=args.batch_size, + num_workers=4, + collate_fn=collate_with_images, + persistent_workers=False, + pin_memory=False, + ) + val_loader = torch.utils.data.DataLoader( + test_dataset, + batch_size=args.batch_size, + num_workers=4, + collate_fn=collate_with_images, + persistent_workers=False, + pin_memory=False, + ) + + # Model + out_dim = 20 * 2 + + model = DeepMonocularModel( + feature_extractor=SAMFeatures( + model_name="timm/vit_pe_spatial_small_patch16_512.fb", frozen=True + ), + out_dim=out_dim, + n_blocks=4, + ) + if args.compile: + model = torch.compile(model, mode="max-autotune") + lit_model = LitModel(model=model, lr=args.lr) + + base_path = Path(args.data_dir).parent + ckpt_dir = base_path / "checkpoints_baseline" + ckpt_dir.mkdir(parents=True, exist_ok=True) + + strategy = "ddp" if torch.cuda.device_count() > 1 else "auto" + use_distributed_sampler = args.dataset != "all" + torch.set_float32_matmul_precision("medium") + trainer = pl.Trainer( + max_epochs=args.max_epochs, + logger=False, + strategy=strategy, + use_distributed_sampler=use_distributed_sampler, + precision="bf16-mixed" if torch.cuda.is_bf16_supported() else 16, + profiler=SimpleProfiler(extended=True) if args.profile else None, + callbacks=[ + ModelCheckpoint( + monitor="val_ade_pred", + mode="min", + save_top_k=1, + dirpath=ckpt_dir, + filename="camera-e2e-{epoch:02d}-{val_ade_pred:.2f}", + ), + ], + gradient_clip_val=1.0, + ) + + trainer.fit(lit_model, train_loader, val_loader) diff --git a/src/camera-based-e2e/train_baseline.slurm b/src/camera-based-e2e/train_baseline.slurm new file mode 100644 index 0000000..32aedfe --- /dev/null +++ b/src/camera-based-e2e/train_baseline.slurm @@ -0,0 +1,26 @@ +#!/bin/bash +#SBATCH --job-name=waymo_baseline_5ep +#SBATCH --output=logs/%x_%j.out +#SBATCH --error=logs/%x_%j.err +#SBATCH --partition=a10 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=72:00:00 +#SBATCH --account=csso + +PYTHON=/scratch/gilbreth/kumar753/conda_envs/robo_env_310_new/bin/python +SCRIPT_DIR=/scratch/gilbreth/kumar753/robotvision/robotvision/src/camera-based-e2e +DATA_DIR=/scratch/gilbreth/kumar753/robotvision/waymo_end_to_end_camera_v1_0_0/waymo_open_dataset_end_to_end_camera_v_1_0_0 + +# make sure logs folder exists +mkdir -p ${SCRIPT_DIR}/logs + +cd ${SCRIPT_DIR} + +$PYTHON train_baseline.py \ + --data_dir "$DATA_DIR" \ + --batch_size 4 \ + --lr 5e-5 \ + --max_epochs 10 + diff --git a/src/camera-based-e2e/train_occ.py b/src/camera-based-e2e/train_occ.py new file mode 100644 index 0000000..c6fc47b --- /dev/null +++ b/src/camera-based-e2e/train_occ.py @@ -0,0 +1,345 @@ +import argparse +import itertools +import math +import random + +import pytorch_lightning as pl +from pytorch_lightning.callbacks import ModelCheckpoint +from pytorch_lightning.profilers import SimpleProfiler + +import torch +from pathlib import Path +import os +from torch.utils.data import BatchSampler + + +from models.base_model_occ import LitModel, collate_with_images +from models.monocular_occ import DeepMonocularModel +from models.feature_extractors import SAMFeatures + + +class HomogeneousConcatBatchSampler(BatchSampler): + """Emit batches from one ConcatDataset source at a time. + + Designed for ConcatDataset([waymo, nuscenes]) so every batch is + from one source. Works with DDP + """ + + def __init__( + self, + dataset_lengths: tuple[int, int], + batch_size: int, + rank: int | None = None, + world_size: int | None = None, + drop_last: bool = False, + shuffle: bool = True, + seed: int = 42, + source_ratio: tuple[int, int] = (1, 1), + **kwargs, + ): + if len(dataset_lengths) != 2: + raise ValueError(f"Expected exactly 2 datasets, got {len(dataset_lengths)}") + if batch_size <= 0: + raise ValueError(f"batch_size must be > 0, got {batch_size}") + if rank is None: + rank = int(os.environ.get("RANK", "0")) + if world_size is None: + world_size = int(os.environ.get("WORLD_SIZE", "1")) + + if world_size <= 0: + raise ValueError(f"world_size must be > 0, got {world_size}") + if rank < 0 or rank >= world_size: + raise ValueError( + f"Invalid rank/world_size pair: rank={rank}, world_size={world_size}" + ) + + self.lengths = dataset_lengths + self.batch_size = batch_size + self.rank = rank + self.world_size = world_size + self.drop_last = drop_last + self.shuffle = shuffle + self.seed = seed + self.epoch = 0 + self.source_ratio = source_ratio + + self.offset0 = 0 + self.offset1 = dataset_lengths[0] + + def set_epoch(self, epoch: int): + self.epoch = int(epoch) + + def _num_samples_per_rank(self, length: int) -> int: + if length <= 0: + return 0 + if self.drop_last and length % self.world_size != 0: + return math.ceil((length - self.world_size) / self.world_size) + return math.ceil(length / self.world_size) + + def _make_rank_indices(self, start: int, length: int, rng: random.Random): + idx = list(range(start, start + length)) + if self.shuffle: + rng.shuffle(idx) + num_samples = self._num_samples_per_rank(length) + total_size = num_samples * self.world_size + + if self.drop_last: + idx = idx[:total_size] + elif total_size > len(idx): + if not idx: + return [] + padding_size = total_size - len(idx) + repeats = math.ceil(padding_size / len(idx)) + idx += (idx * repeats)[:padding_size] + + return idx[self.rank : total_size : self.world_size] + + def _chunk(self, indices: list[int]) -> list[list[int]]: + if self.drop_last: + n_full = len(indices) // self.batch_size + return [ + indices[i * self.batch_size : (i + 1) * self.batch_size] + for i in range(n_full) + ] + + out = [] + for i in range(0, len(indices), self.batch_size): + out.append(indices[i : i + self.batch_size]) + return out + + def __iter__(self): + rng = random.Random(self.seed + self.epoch) + + idx0 = self._make_rank_indices(self.offset0, self.lengths[0], rng) + idx1 = self._make_rank_indices(self.offset1, self.lengths[1], rng) + + b0 = self._chunk(idx0) + b1 = self._chunk(idx1) + + w0 = max(0, int(self.source_ratio[0])) + w1 = max(0, int(self.source_ratio[1])) + if w0 == 0 and w1 == 0: + w0, w1 = 1, 1 + + pattern = [0] * w0 + [1] * w1 + if not pattern: + pattern = [0, 1] + + i0, i1 = 0, 0 + for source in itertools.cycle(pattern): + if i0 >= len(b0) and i1 >= len(b1): + break + if source == 0: + if i0 < len(b0): + yield b0[i0] + i0 += 1 + else: + if i1 < len(b1): + yield b1[i1] + i1 += 1 + + def __len__(self): + def n_batches(length: int): + per_rank = self._num_samples_per_rank(length) + if self.drop_last: + return per_rank // self.batch_size + return math.ceil(per_rank / self.batch_size) + + return n_batches(self.lengths[0]) + n_batches(self.lengths[1]) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--data_dir", type=str, required=True, help="Path to data directory" + ) + parser.add_argument( + "--batch_size", type=int, default=16, help="Batch size for training" + ) + parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate") + parser.add_argument( + "--max_epochs", type=int, default=10, help="Number of epochs to train" + ) + parser.add_argument( + "--compile", + action="store_true", + help="Whether to compile the model with torch.compile", + ) + parser.add_argument( + "--profile", action="store_true", help="Whether to run the profiler" + ) + parser.add_argument( + "--dataset", + type=str, + default="waymo", + choices=["waymo", "nuscenes", "all"], + help="Which dataset to train on", + ) + parser.add_argument( + "--occ_root", type=str, required=True, help="Path to occupancy root directory" + ) + args = parser.parse_args() + + pl.seed_everything(42, workers=True) + + # Data + if args.dataset == "waymo": + from loader_occ import WaymoE2E + + train_dataset = WaymoE2E( + indexFile="index_train.pkl", + data_dir=args.data_dir, + n_items=250_000, + occ_root=args.occ_root, + ) + test_dataset = WaymoE2E( + indexFile="index_val.pkl", + data_dir=args.data_dir, + n_items=20_000, + occ_root=args.occ_root, + ) + nw = 0 + elif args.dataset == "nuscenes": + from nuscenes_loader import NuScenesDataset + + train_dataset = NuScenesDataset( + data_dir=args.data_dir, split="train", n_items=250_000 + ) + test_dataset = NuScenesDataset( + data_dir=args.data_dir, split="val", n_items=25_000 + ) + nw = 16 + elif args.dataset == "all": + from loader_occ import WaymoE2E + from nuscenes_loader import NuScenesDataset + + waymo_dir, nuscenes_dir = os.getenv("WAYMO_DATA_DIR"), os.getenv( + "NUSCENES_DATA_DIR" + ) + + waymo_train = WaymoE2E( + indexFile="index_train.pkl", + data_dir=waymo_dir, + n_items=125_000, + occ_root=args.occ_root, + ) + waymo_test = WaymoE2E( + indexFile="index_val.pkl", + data_dir=waymo_dir, + n_items=12_500, + occ_root=args.occ_root, + ) + + nuscenes_train = NuScenesDataset( + data_dir=nuscenes_dir, split="train", n_items=125_000 + ) + nuscenes_test = NuScenesDataset( + data_dir=nuscenes_dir, split="val", n_items=12_500 + ) + + train_dataset = torch.utils.data.ConcatDataset([waymo_train, nuscenes_train]) + test_dataset = torch.utils.data.ConcatDataset([waymo_test, nuscenes_test]) + nw = 16 + + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + + train_batch_sampler = HomogeneousConcatBatchSampler( + dataset_lengths=(len(waymo_train), len(nuscenes_train)), + batch_size=args.batch_size, + rank=rank, + world_size=world_size, + drop_last=False, + shuffle=True, + seed=42, + source_ratio=(1, 1), + ) + val_batch_sampler = HomogeneousConcatBatchSampler( + dataset_lengths=(len(waymo_test), len(nuscenes_test)), + batch_size=args.batch_size, + rank=rank, + world_size=world_size, + drop_last=False, + shuffle=False, + seed=42, + source_ratio=(1, 1), + ) + else: + raise ValueError(f"Unsupported dataset: {args.dataset}") + + if args.dataset == "all": + train_loader = torch.utils.data.DataLoader( + train_dataset, + num_workers=4, + batch_sampler=train_batch_sampler, + collate_fn=collate_with_images, + persistent_workers=False, + pin_memory=False, + ) + val_loader = torch.utils.data.DataLoader( + test_dataset, + num_workers=4, + batch_sampler=val_batch_sampler, + collate_fn=collate_with_images, + persistent_workers=False, + pin_memory=False, + ) + else: + train_loader = torch.utils.data.DataLoader( + train_dataset, + batch_size=args.batch_size, + num_workers=4, + collate_fn=collate_with_images, + persistent_workers=False, + pin_memory=False, + ) + val_loader = torch.utils.data.DataLoader( + test_dataset, + batch_size=args.batch_size, + num_workers=4, + collate_fn=collate_with_images, + persistent_workers=False, + pin_memory=False, + ) + + # Model + out_dim = 20 * 2 + + model = DeepMonocularModel( + feature_extractor=SAMFeatures( + model_name="timm/vit_pe_spatial_small_patch16_512.fb", frozen=True + ), + out_dim=out_dim, + n_blocks=4, + ) + if args.compile: + model = torch.compile(model, mode="max-autotune") + lit_model = LitModel(model=model, lr=args.lr) + + base_path = Path(args.data_dir).parent + ckpt_dir = base_path / "checkpoints_occ" + ckpt_dir.mkdir(parents=True, exist_ok=True) + + strategy = "ddp" if torch.cuda.device_count() > 1 else "auto" + use_distributed_sampler = args.dataset != "all" + torch.set_float32_matmul_precision("medium") + trainer = pl.Trainer( + max_epochs=args.max_epochs, + logger=False, + strategy=strategy, + use_distributed_sampler=use_distributed_sampler, + precision="bf16-mixed" if torch.cuda.is_bf16_supported() else 16, + profiler=SimpleProfiler(extended=True) if args.profile else None, + callbacks=[ + ModelCheckpoint( + monitor="val_ade_pred", + mode="min", + save_top_k=1, + dirpath=ckpt_dir, + filename="camera-e2e-{epoch:02d}-{val_ade_pred:.2f}", + ), + ], + gradient_clip_val=1.0, + ) + + trainer.fit(lit_model, train_loader, val_loader) diff --git a/src/camera-based-e2e/train_occ.slurm b/src/camera-based-e2e/train_occ.slurm new file mode 100644 index 0000000..5e645b4 --- /dev/null +++ b/src/camera-based-e2e/train_occ.slurm @@ -0,0 +1,26 @@ +#!/bin/bash +#SBATCH --job-name=waymo_occ_10ep +#SBATCH --output=logs/%x_%j.out +#SBATCH --error=logs/%x_%j.err +#SBATCH --partition=a10 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=72:00:00 +#SBATCH --account=csso + +PYTHON=/scratch/gilbreth/kumar753/conda_envs/robo_env_310_new/bin/python +SCRIPT_DIR=/scratch/gilbreth/kumar753/robotvision/robotvision/src/camera-based-e2e +DATA_DIR=/scratch/gilbreth/kumar753/robotvision/waymo_end_to_end_camera_v1_0_0/waymo_open_dataset_end_to_end_camera_v_1_0_0 +OCC_ROOT=/scratch/gilbreth/kumar753/waymo_occ_new + +mkdir -p ${SCRIPT_DIR}/logs +cd ${SCRIPT_DIR} + +$PYTHON train_occ.py \ + --data_dir "$DATA_DIR" \ + --occ_root "$OCC_ROOT" \ + --batch_size 4 \ + --lr 5e-5 \ + --max_epochs 10 + diff --git a/src/camera-based-e2e/train_og.py b/src/camera-based-e2e/train_og.py new file mode 100644 index 0000000..4db3136 --- /dev/null +++ b/src/camera-based-e2e/train_og.py @@ -0,0 +1,298 @@ +import argparse +from datetime import datetime +import itertools +import math +import random + +import pytorch_lightning as pl +from pytorch_lightning.callbacks import ModelCheckpoint +from pytorch_lightning.profilers import SimpleProfiler + +import torch +from pathlib import Path +import os +from torch.utils.data import BatchSampler + +# Replace with your model defined in models/ +from models.base_model import LitModel, collate_with_images +from models.monocular import DeepMonocularModel +from models.feature_extractors import SAMFeatures + + +class HomogeneousConcatBatchSampler(BatchSampler): + """Emit batches from one ConcatDataset source at a time.""" + + def __init__( + self, + dataset_lengths: tuple[int, int], + batch_size: int, + rank: int | None = None, + world_size: int | None = None, + drop_last: bool = False, + shuffle: bool = True, + seed: int = 42, + source_ratio: tuple[int, int] = (1, 1), + **kwargs, + ): + if len(dataset_lengths) != 2: + raise ValueError(f"Expected exactly 2 datasets, got {len(dataset_lengths)}") + if batch_size <= 0: + raise ValueError(f"batch_size must be > 0, got {batch_size}") + if rank is None: + rank = int(os.environ.get("RANK", "0")) + if world_size is None: + world_size = int(os.environ.get("WORLD_SIZE", "1")) + + if world_size <= 0: + raise ValueError(f"world_size must be > 0, got {world_size}") + if rank < 0 or rank >= world_size: + raise ValueError( + f"Invalid rank/world_size pair: rank={rank}, world_size={world_size}" + ) + + self.lengths = dataset_lengths + self.batch_size = batch_size + self.rank = rank + self.world_size = world_size + self.drop_last = drop_last + self.shuffle = shuffle + self.seed = seed + self.epoch = 0 + self.source_ratio = source_ratio + + self.offset0 = 0 + self.offset1 = dataset_lengths[0] + + def set_epoch(self, epoch: int): + self.epoch = int(epoch) + + def _num_samples_per_rank(self, length: int) -> int: + if length <= 0: + return 0 + if self.drop_last and length % self.world_size != 0: + return math.ceil((length - self.world_size) / self.world_size) + return math.ceil(length / self.world_size) + + def _make_rank_indices(self, start: int, length: int, rng: random.Random): + idx = list(range(start, start + length)) + if self.shuffle: + rng.shuffle(idx) + num_samples = self._num_samples_per_rank(length) + total_size = num_samples * self.world_size + + if self.drop_last: + idx = idx[:total_size] + elif total_size > len(idx): + if not idx: + return [] + padding_size = total_size - len(idx) + repeats = math.ceil(padding_size / len(idx)) + idx += (idx * repeats)[:padding_size] + + return idx[self.rank : total_size : self.world_size] + + def _chunk(self, indices: list[int]) -> list[list[int]]: + if self.drop_last: + n_full = len(indices) // self.batch_size + return [ + indices[i * self.batch_size : (i + 1) * self.batch_size] + for i in range(n_full) + ] + + out = [] + for i in range(0, len(indices), self.batch_size): + out.append(indices[i : i + self.batch_size]) + return out + + def __iter__(self): + rng = random.Random(self.seed + self.epoch) + + idx0 = self._make_rank_indices(self.offset0, self.lengths[0], rng) + idx1 = self._make_rank_indices(self.offset1, self.lengths[1], rng) + + b0 = self._chunk(idx0) + b1 = self._chunk(idx1) + + w0 = max(0, int(self.source_ratio[0])) + w1 = max(0, int(self.source_ratio[1])) + if w0 == 0 and w1 == 0: + w0, w1 = 1, 1 + + pattern = [0] * w0 + [1] * w1 + if not pattern: + pattern = [0, 1] + + i0, i1 = 0, 0 + for source in itertools.cycle(pattern): + if i0 >= len(b0) and i1 >= len(b1): + break + if source == 0: + if i0 < len(b0): + yield b0[i0] + i0 += 1 + else: + if i1 < len(b1): + yield b1[i1] + i1 += 1 + + def __len__(self): + def n_batches(length: int): + per_rank = self._num_samples_per_rank(length) + if self.drop_last: + return per_rank // self.batch_size + return math.ceil(per_rank / self.batch_size) + + return n_batches(self.lengths[0]) + n_batches(self.lengths[1]) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data_dir", type=str, required=True) + parser.add_argument("--batch_size", type=int, default=16) + parser.add_argument("--lr", type=float, default=1e-4) + parser.add_argument("--max_epochs", type=int, default=10) + parser.add_argument("--compile", action="store_true") + parser.add_argument("--profile", action="store_true") + parser.add_argument( + "--dataset", + type=str, + default="waymo", + choices=["waymo", "nuscenes", "all"], + ) + args = parser.parse_args() + + pl.seed_everything(42, workers=True) + + # Data + if args.dataset == "waymo": + from loader import WaymoE2E + + train_dataset = WaymoE2E( + indexFile="index_train.pkl", data_dir=args.data_dir, n_items=250_000 + ) + test_dataset = WaymoE2E( + indexFile="index_val.pkl", data_dir=args.data_dir, n_items=25_000 + ) + nw = 0 + + elif args.dataset == "nuscenes": + from nuscenes_loader import NuScenesDataset + + train_dataset = NuScenesDataset( + data_dir=args.data_dir, split="train", n_items=250_000 + ) + test_dataset = NuScenesDataset( + data_dir=args.data_dir, split="val", n_items=25_000 + ) + nw = 16 + + elif args.dataset == "all": + from loader import WaymoE2E + from nuscenes_loader import NuScenesDataset + + waymo_dir = os.getenv("WAYMO_DATA_DIR") + nuscenes_dir = os.getenv("NUSCENES_DATA_DIR") + + waymo_train = WaymoE2E( + indexFile="index_train.pkl", data_dir=waymo_dir, n_items=125_000 + ) + waymo_test = WaymoE2E( + indexFile="index_val.pkl", data_dir=waymo_dir, n_items=12_500 + ) + + nuscenes_train = NuScenesDataset( + data_dir=nuscenes_dir, split="train", n_items=125_000 + ) + nuscenes_test = NuScenesDataset( + data_dir=nuscenes_dir, split="val", n_items=12_500 + ) + + train_dataset = torch.utils.data.ConcatDataset([waymo_train, nuscenes_train]) + test_dataset = torch.utils.data.ConcatDataset([waymo_test, nuscenes_test]) + nw = 16 + + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + + train_batch_sampler = HomogeneousConcatBatchSampler( + dataset_lengths=(len(waymo_train), len(nuscenes_train)), + batch_size=args.batch_size, + rank=rank, + world_size=world_size, + ) + val_batch_sampler = HomogeneousConcatBatchSampler( + dataset_lengths=(len(waymo_test), len(nuscenes_test)), + batch_size=args.batch_size, + rank=rank, + world_size=world_size, + shuffle=False, + ) + + else: + raise ValueError(f"Unsupported dataset: {args.dataset}") + + if args.dataset == "all": + train_loader = torch.utils.data.DataLoader( + train_dataset, + num_workers=nw, + batch_sampler=train_batch_sampler, + collate_fn=collate_with_images, + ) + val_loader = torch.utils.data.DataLoader( + test_dataset, + num_workers=nw, + batch_sampler=val_batch_sampler, + collate_fn=collate_with_images, + ) + else: + train_loader = torch.utils.data.DataLoader( + train_dataset, + batch_size=args.batch_size, + num_workers=nw, + collate_fn=collate_with_images, + ) + val_loader = torch.utils.data.DataLoader( + test_dataset, + batch_size=args.batch_size, + num_workers=nw, + collate_fn=collate_with_images, + ) + + # Model + in_dim = 16 * 6 + out_dim = 20 * 2 + + model = DeepMonocularModel( + feature_extractor=SAMFeatures( + model_name="timm/vit_pe_spatial_small_patch16_512.fb", frozen=True + ), + out_dim=out_dim, + n_blocks=4, + ) + + if args.compile: + model = torch.compile(model, mode="max-autotune") + + lit_model = LitModel(model=model, lr=args.lr) + + strategy = "ddp" if torch.cuda.device_count() > 1 else "auto" + use_distributed_sampler = args.dataset != "all" + + trainer = pl.Trainer( + max_epochs=args.max_epochs, + strategy=strategy, + use_distributed_sampler=use_distributed_sampler, + precision="bf16-mixed" if torch.cuda.is_bf16_supported() else 16, + profiler=SimpleProfiler(extended=True) if args.profile else None, + callbacks=[ + ModelCheckpoint( + monitor="val_loss", + mode="min", + save_top_k=1, + dirpath=Path(args.data_dir).parent.as_posix() + "/checkpoints", + filename="camera-e2e-{epoch:02d}-{val_loss:.2f}", + ), + ], + ) + + trainer.fit(lit_model, train_loader, val_loader) diff --git a/src/camera-based-e2e/train_og.slurm b/src/camera-based-e2e/train_og.slurm new file mode 100644 index 0000000..5cd69ce --- /dev/null +++ b/src/camera-based-e2e/train_og.slurm @@ -0,0 +1,26 @@ +#!/bin/bash +#SBATCH --job-name=waymo_og_5ep +#SBATCH --output=logs/%x_%j.out +#SBATCH --error=logs/%x_%j.err +#SBATCH --partition=a10 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=48G +#SBATCH --time=72:00:00 +#SBATCH --account=csso + +PYTHON=/scratch/gilbreth/kumar753/conda_envs/robo_env_310_new/bin/python +SCRIPT_DIR=/scratch/gilbreth/kumar753/robotvision/robotvision/src/camera-based-e2e +DATA_DIR=/scratch/gilbreth/kumar753/robotvision/waymo_end_to_end_camera_v1_0_0/waymo_open_dataset_end_to_end_camera_v_1_0_0 + +# make sure logs folder exists +mkdir -p ${SCRIPT_DIR}/logs + +cd ${SCRIPT_DIR} + +$PYTHON train_og.py \ + --data_dir "$DATA_DIR" \ + --batch_size 4 \ + --lr 5e-5 \ + --max_epochs 10 + diff --git a/src/camera-based-e2e/viz_occ.py b/src/camera-based-e2e/viz_occ.py new file mode 100644 index 0000000..3731927 --- /dev/null +++ b/src/camera-based-e2e/viz_occ.py @@ -0,0 +1,158 @@ +"""Quick sanity-check visualizer for occupancy grids. + +Loads a few occ_{idx}.npy files and shows: + 1. Top-down 2D slice (XY at ground level) — colored by class + 2. Side view (XZ slice through center) + 3. Class distribution bar chart + +Does NOT require Open3D. Uses matplotlib only. + +Usage (on login node — no GPU needed): + python viz_occ.py \ + --occ_dir /scratch/gilbreth/kumar753/waymo_occ/train \ + --indices 0 100 500 1000 5000 + +Classes: + 0 = free (white) + 1 = vehicle (red) + 2 = pedestrian (blue) + 3 = cyclist (orange) + 4 = road (gray) + 5 = static (green) + 255 = unknown (black) +""" + +import argparse +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches +import os + +# Class colors (RGB, normalized) +CLASS_COLORS = { + 255: (0.0, 0.0, 0.0), # unknown — black + 0: (1.0, 1.0, 1.0), # free — white + 1: (0.9, 0.1, 0.1), # vehicle — red + 2: (0.1, 0.3, 0.9), # pedestrian — blue + 3: (1.0, 0.6, 0.0), # cyclist — orange + 4: (0.5, 0.5, 0.5), # road — gray + 5: (0.2, 0.7, 0.2), # static — green +} +CLASS_NAMES = {255: "unknown", 0: "free", 1: "vehicle", 2: "pedestrian", + 3: "cyclist", 4: "road", 5: "static"} + +VOX_XY_SIZE = 100 +VOX_Z_SIZE = 16 +VOX_Z_MIN = -3.0 +VOX_Z_RES = 0.5 + + +def grid_to_rgb(slice_2d): + """Convert a 2D (H, W) uint8 class grid to (H, W, 3) RGB image.""" + h, w = slice_2d.shape + rgb = np.zeros((h, w, 3), dtype=np.float32) + for cls, color in CLASS_COLORS.items(): + mask = slice_2d == cls + rgb[mask] = color + return rgb + + +def ground_level_z(): + """Return the voxel Z index closest to Z=0 (ground level).""" + return int((0.0 - VOX_Z_MIN) / VOX_Z_RES) + + +def visualize_occ(occ_path, ax_row): + """Visualize one occupancy grid on a row of 3 axes.""" + grid = np.load(occ_path) # (100, 100, 16) + idx = os.path.basename(occ_path).replace("occ_", "").replace(".npy", "") + + # ── Slice 1: top-down XY at ground level ── + z_ground = ground_level_z() + xy_slice = grid[:, :, z_ground] # (100, 100) + ax_row[0].imshow(grid_to_rgb(xy_slice), origin="upper") + ax_row[0].set_title(f"idx={idx} XY@Z=0m", fontsize=9) + ax_row[0].axis("off") + + # ── Slice 2: top-down XY — max occupied class along Z ── + # For each XY cell take the most common non-unknown, non-free class + best = np.full((VOX_XY_SIZE, VOX_XY_SIZE), 255, dtype=np.uint8) + for z in range(VOX_Z_SIZE): + layer = grid[:, :, z] + occupied = (layer >= 1) & (layer <= 5) + best[occupied] = layer[occupied] + ax_row[1].imshow(grid_to_rgb(best), origin="upper") + ax_row[1].set_title(f"idx={idx} XY max-Z", fontsize=9) + ax_row[1].axis("off") + + # ── Slice 3: class distribution ── + classes = [0, 1, 2, 3, 4, 5, 255] + counts = [int((grid == c).sum()) for c in classes] + colors = [CLASS_COLORS[c] for c in classes] + labels = [CLASS_NAMES[c] for c in classes] + bars = ax_row[2].bar(labels, counts, color=colors, edgecolor="black", linewidth=0.5) + ax_row[2].set_title(f"idx={idx} voxel counts", fontsize=9) + ax_row[2].tick_params(axis="x", labelsize=7, rotation=30) + ax_row[2].tick_params(axis="y", labelsize=7) + + # Print summary + total = grid.size + known = (grid != 255).sum() + free = (grid == 0).sum() + occupied_sum = sum((grid == c).sum() for c in range(1, 6)) + print(f" idx={idx}: known={known/total:.1%} free={free/total:.1%} " + f"occupied={occupied_sum/total:.1%} unknown={((grid==255).sum())/total:.1%}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--occ_dir", type=str, required=True, + help="Directory containing occ_{idx}.npy files") + parser.add_argument("--indices", type=int, nargs="+", default=[0, 100, 500, 1000, 5000], + help="Frame indices to visualize") + parser.add_argument("--out", type=str, default="occ_viz.png", + help="Output image path") + args = parser.parse_args() + + # Filter to indices that actually exist + paths = [] + for idx in args.indices: + p = os.path.join(args.occ_dir, f"occ_{idx:07d}.npy") + if os.path.exists(p): + paths.append(p) + else: + print(f" Warning: {p} not found, skipping") + + if not paths: + print("No valid occ files found. Check --occ_dir and --indices.") + return + + n = len(paths) + fig, axes = plt.subplots(n, 3, figsize=(12, 4 * n)) + if n == 1: + axes = [axes] + + print(f"\nVoxel statistics:") + for i, path in enumerate(paths): + visualize_occ(path, axes[i]) + + # Legend + legend_patches = [ + mpatches.Patch(color=CLASS_COLORS[c], label=CLASS_NAMES[c], linewidth=0.5, + edgecolor="black") + for c in [255, 0, 1, 2, 3, 4, 5] + ] + fig.legend(handles=legend_patches, loc="lower center", ncol=7, + fontsize=8, bbox_to_anchor=(0.5, 0.0)) + + plt.suptitle("Occupancy grid sanity check\n" + "Left: XY slice at Z=0m | Middle: XY max-Z projection | Right: class counts", + fontsize=10) + plt.tight_layout(rect=[0, 0.05, 1, 1]) + plt.savefig(args.out, dpi=150, bbox_inches="tight") + print(f"\nSaved to {args.out}") + + +if __name__ == "__main__": + main() + diff --git a/src/camera-based-e2e/viz_occ_3d.py b/src/camera-based-e2e/viz_occ_3d.py new file mode 100644 index 0000000..48710d7 --- /dev/null +++ b/src/camera-based-e2e/viz_occ_3d.py @@ -0,0 +1,65 @@ +import numpy as np +import plotly.graph_objects as go +import argparse + +VOX_XY_RANGE = 25.0 +VOX_XY_RES = 0.5 +VOX_Z_MIN = -3.0 +VOX_Z_RES = 0.5 + +COLORS = { + 0: 'lightgrey', # free + 1: 'red', # vehicle + 2: 'blue', # pedestrian + 3: 'orange', # cyclist + 4: 'darkgrey', # road + 5: 'green', # static +} +NAMES = {0:'free', 1:'vehicle', 2:'pedestrian', 3:'cyclist', 4:'road', 5:'static'} + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--occ_path", type=str, required=True) + parser.add_argument("--show_free", action="store_true") + args = parser.parse_args() + + grid = np.load(args.occ_path) + ix, iy, iz = np.where(grid != 255) + classes = grid[ix, iy, iz] + + x = VOX_XY_RANGE - ix * VOX_XY_RES + y = VOX_XY_RANGE - iy * VOX_XY_RES + z = VOX_Z_MIN + iz * VOX_Z_RES + + traces = [] + for cls in range(0 if args.show_free else 1, 6): + mask = classes == cls + if mask.sum() == 0: + continue + traces.append(go.Scatter3d( + x=x[mask], y=y[mask], z=z[mask], + mode='markers', + marker=dict(size=3, color=COLORS[cls], opacity=0.8), + name=NAMES[cls] + )) + + fig = go.Figure(data=traces) + fig.update_layout( + title="Occupancy Grid 3D View", + scene=dict( + xaxis_title="X (m)", + yaxis_title="Y (m)", + zaxis_title="Z (m)", + camera=dict( + up=dict(x=0, y=0, z=1), + eye=dict(x=0, y=0, z=2.5) # top-down by default + ) + ) + ) + out = args.occ_path.replace('.npy', '.html') + fig.write_html(out) + print(f"Saved to {out} — open in browser") + +if __name__ == "__main__": + main() + diff --git a/src/camera-based-e2e/viz_seg.py b/src/camera-based-e2e/viz_seg.py new file mode 100644 index 0000000..c8b27e7 --- /dev/null +++ b/src/camera-based-e2e/viz_seg.py @@ -0,0 +1,91 @@ +""" +Run SegFormer on extracted frame images and save colored segmentation maps. +Usage: + python viz_seg.py --idx 5000 +""" +import argparse +import os +import numpy as np +import cv2 +import torch +import torch.nn.functional as F +from transformers import AutoImageProcessor, SegformerForSemanticSegmentation +from pathlib import Path + +SEG_MODEL = "nvidia/segformer-b2-finetuned-cityscapes-1024-1024" + +# Cityscapes colors for each class +CITYSCAPES_COLORS = [ + (128, 64,128), # road + (244, 35,232), # sidewalk + ( 70, 70, 70), # building + (102,102,156), # wall + (190,153,153), # fence + (153,153,153), # pole + (250,170, 30), # traffic light + (220,220, 0), # traffic sign + (107,142, 35), # vegetation + (152,251,152), # terrain + ( 70,130,180), # sky + (220, 20, 60), # person + (255, 0, 0), # rider + ( 0, 0,142), # car + ( 0, 0, 70), # truck + ( 0, 60,100), # bus + ( 0, 80,100), # train + ( 0, 0,230), # motorcycle + (119, 11, 32), # bicycle +] + +CITYSCAPES_NAMES = [ + 'road','sidewalk','building','wall','fence','pole', + 'traffic light','traffic sign','vegetation','terrain','sky', + 'person','rider','car','truck','bus','train','motorcycle','bicycle' +] + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--idx", type=int, default=5000) + parser.add_argument("--img_dir", type=str, default="./frame_images") + parser.add_argument("--out_dir", type=str, default="./seg_output") + args = parser.parse_args() + + Path(args.out_dir).mkdir(exist_ok=True) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + print(f"Loading SegFormer...") + processor = AutoImageProcessor.from_pretrained(SEG_MODEL) + model = SegformerForSemanticSegmentation.from_pretrained(SEG_MODEL).to(device) + model.eval() + + for cam in range(1, 9): + img_path = os.path.join(args.img_dir, f"frame_{args.idx:07d}_cam{cam}.jpg") + if not os.path.exists(img_path): + continue + + rgb = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB) + H, W = rgb.shape[:2] + + inputs = processor(images=rgb, return_tensors="pt") + inputs = {k: v.to(device) for k, v in inputs.items()} + + with torch.no_grad(): + logits = model(**inputs).logits + + pred = F.interpolate(logits.float(), size=(H, W), mode="bilinear", align_corners=False) + class_map = pred.argmax(dim=1).squeeze().cpu().numpy() + + # Color the segmentation + seg_colored = np.zeros((H, W, 3), dtype=np.uint8) + for cls_id, color in enumerate(CITYSCAPES_COLORS): + seg_colored[class_map == cls_id] = color + + # Side by side: original + segmentation + combined = np.concatenate([rgb[:,:,::-1], seg_colored[:,:,::-1]], axis=1) + out_path = os.path.join(args.out_dir, f"seg_{args.idx:07d}_cam{cam}.jpg") + cv2.imwrite(out_path, combined) + print(f" Saved cam{cam}: {out_path}") + +if __name__ == "__main__": + main() +