From c7f936a8bf7f014c7d0106fa03dace64d4f87539 Mon Sep 17 00:00:00 2001 From: Croopan Date: Fri, 16 Jan 2026 13:10:29 -0500 Subject: [PATCH 1/3] RFS model training --- .../dataset/export_dataset.py | 7 +- src/camera-based-e2e/loader.py | 27 ++-- src/camera-based-e2e/models/base_model.py | 25 +++- src/camera-based-e2e/models/monocular.py | 16 ++- src/camera-based-e2e/preference_cache.pkl | Bin 0 -> 10942 bytes src/camera-based-e2e/preference_cache.py | 55 ++++++++ src/camera-based-e2e/preference_loader.py | 132 ++++++++++++++++++ src/camera-based-e2e/protos/dataset.proto | 1 - src/camera-based-e2e/protos/e2e.proto | 2 +- src/camera-based-e2e/testRead.py | 44 ++++++ src/camera-based-e2e/train.py | 47 ++++++- src/camera-based-e2e/train.slurm | 20 +++ 12 files changed, 345 insertions(+), 31 deletions(-) create mode 100644 src/camera-based-e2e/preference_cache.pkl create mode 100644 src/camera-based-e2e/preference_cache.py create mode 100644 src/camera-based-e2e/preference_loader.py create mode 100644 src/camera-based-e2e/testRead.py create mode 100644 src/camera-based-e2e/train.slurm diff --git a/src/camera-based-e2e/dataset/export_dataset.py b/src/camera-based-e2e/dataset/export_dataset.py index 274de66..a9853b0 100644 --- a/src/camera-based-e2e/dataset/export_dataset.py +++ b/src/camera-based-e2e/dataset/export_dataset.py @@ -15,17 +15,18 @@ def safe_name(name: str) -> str: # Keep filesystem-friendly ASCII. - return re.sub(r"[^A-Za-z0-9_.-]+", "_", name) + return re.sub(r"[^A-Za-z0-9_.-]+", "_", name) #substitutes any non ascii chars with _ -def camera_name(name_enum: int) -> str: +def camera_name(name_enum: int) -> str: #which camera (front, left, etc) return dataset_pb2.CameraName.Name.Name(name_enum) + #access dataset_pb2 class, then CameraName nested class, then Name enum nested class, then Name constructor method to get string name def intent_name(intent_enum: int) -> str: return e2e_pb2.EgoIntent.Intent.Name(intent_enum) - +#NOT SURE ABOUT THIS FUNCTION def transform_list(transform_msg) -> list: return list(transform_msg.transform) diff --git a/src/camera-based-e2e/loader.py b/src/camera-based-e2e/loader.py index 4cbc140..27c1e6f 100644 --- a/src/camera-based-e2e/loader.py +++ b/src/camera-based-e2e/loader.py @@ -1,6 +1,6 @@ import torch from torch.utils.data import IterableDataset -from waymo_open_dataset.protos import end_to_end_driving_data_pb2 as e2e_pb2 +from protos import e2e_pb2 import torchvision import pickle import struct @@ -19,7 +19,7 @@ class WaymoE2E(IterableDataset): def __init__( self, - indexFile = 'index.pkl', + indexFile = 'preference_cache.pkl', data_dir='./dataset', images = True, n_items: Optional[int] = None, @@ -27,7 +27,7 @@ def __init__( ): self.images = images self.data_dir = data_dir - self.seed = seed + self.seed = seed self.filename = "" self.file = None @@ -37,7 +37,6 @@ def __init__( # We train on train and validate on val set self.indexes = pickle.load(f) - # 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): total = len(self.indexes) # pick a deterministic contiguous block when a seed is provided @@ -64,14 +63,16 @@ def __len__(self): return len(self.indexes) def __iter__(self): + # spawns a new worker, assigns worker.id and num_workers worker = torch.utils.data.get_worker_info() if worker is None: start, step = 0, 1 else: start, step = worker.id, worker.num_workers + # this code makes each worker read a subset of teh data (worker 1 = 0, 16, 32..., worker 2 = 1, 17, 33...) for idx in range(start, len(self.indexes), step): - frame = e2e_pb2.E2EDFrame() # type: ignore + frame = e2e_pb2.E2EDFrame() filename, start_byte, byte_length = self.indexes[idx] if self.filename != filename: @@ -81,13 +82,12 @@ def __iter__(self): 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 + self.file.seek(start_byte) + protobuf = self.file.read(byte_length) 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 @@ -96,6 +96,7 @@ def __iter__(self): # For submission to waymo evaluation server name = frame.frame.context.name + # returns dictionary of past, future, images, intent yield {'PAST': past, 'FUTURE': future, 'IMAGES': [self.decode_img(images.image) for images in frame.frame.images], 'INTENT': frame.intent, 'NAME': name} if __name__ == "__main__": @@ -104,15 +105,18 @@ def __iter__(self): import time from tqdm import tqdm # NOTE: Replace with your path - DATA_DIR = '/scratch/gilbreth/mgagvani/wod/waymo_open_dataset_end_to_end_camera_v_1_0_0/' + DATA_DIR = '/scratch/gilbreth/svelmuru/waymo_end_to_end_dataset/waymo_open_dataset_end_to_end_camera_v_1_0_0/' BATCH_SIZE = 32 - dataset = WaymoE2E(indexFile="index_train.pkl", data_dir = DATA_DIR, images=True) + dataset = WaymoE2E(indexFile="preference_cache.pkl", data_dir = DATA_DIR, images=True) loader = DataLoader( dataset, batch_size=BATCH_SIZE, num_workers=16, ) + #now we have this laoder class that I can iterate over like an array + #it will return dictionaries with keys PAST, FUTURE, IMAGES, INTENT + def main(): # start = time.time() for batch_of_frames in tqdm(loader): @@ -120,6 +124,7 @@ def main(): # 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) + print(len(loader)) import cProfile - main() + main() \ No newline at end of file diff --git a/src/camera-based-e2e/models/base_model.py b/src/camera-based-e2e/models/base_model.py index 9096ea2..38d2eb2 100644 --- a/src/camera-based-e2e/models/base_model.py +++ b/src/camera-based-e2e/models/base_model.py @@ -27,6 +27,7 @@ def __init__(self, model: nn.Module, lr: float): '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 + 'PREF_TRAJ': torch.zeros((1, 21, 2)), # PREF_TRAJ },) # ---- Metrics ---- @@ -48,15 +49,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor: past, future, images, intent = batch['PAST'], batch['FUTURE'], batch['IMAGES'], batch['INTENT'] + pref_traj, pref_score = batch["PREF_TRAJ"], batch["PREF_SCORE"] # `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} + model_inputs = {'PAST': past, 'IMAGES': images, 'INTENT': intent, "PREF_TRAJ": pref_traj} - pred_future = self.forward(model_inputs) # (B, T*2) - loss = self.ade_loss(pred_future.reshape_as(future), future) # reshape to (B, T, 2 + pred_score = self.forward(model_inputs).squeeze(-1) # (B,) rather than (B, 1) + loss = F.mse_loss(pred_score, pref_score) # TODO: improve logging both to disk and to console self.log_dict({ @@ -71,12 +73,26 @@ def training_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor: def validation_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor: return self._shared_step(batch, "val") +# helps to batch data with images 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] + padded_trajs = [] + + for b in batch: + traj = torch.as_tensor(b["PREF_TRAJ"], dtype=torch.float32) + for _ in range(len(traj), 21): + traj = F.pad(traj, (0,0,0,1), value=0.0) # pad to length 20 + + padded_trajs.append(traj) + + # Extra values for preference trajectories + pref_traj = padded_trajs + pref_score = torch.as_tensor([b["PREF_SCORE"] for b in batch], dtype=torch.float32) + cams = list(zip(*[b["IMAGES"] for b in batch])) # per-camera tuples images = [torch.stack(cam_imgs, dim=0) for cam_imgs in cams] # stay on CPU @@ -86,5 +102,6 @@ def collate_with_images(batch): "INTENT": intent, "IMAGES": images, "NAME": names, + "PREF_TRAJ": torch.stack(pref_traj, dim=0), + "PREF_SCORE": pref_score, } - diff --git a/src/camera-based-e2e/models/monocular.py b/src/camera-based-e2e/models/monocular.py index 2edd6cc..ee5089a 100644 --- a/src/camera-based-e2e/models/monocular.py +++ b/src/camera-based-e2e/models/monocular.py @@ -73,11 +73,13 @@ def __init__( # attention self.feature_dim = sum(self.features.dims) # works for both DINO and SAM + + #converts tokens into key, value spaces 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 + query_input_dim = 3 + 16 * 6 + 21 * 2 # one hot -- concat -- flattened self.query = nn.Sequential( nn.Linear(query_input_dim, self.feature_dim), nn.LeakyReLU(), @@ -104,6 +106,8 @@ def __init__( def forward(self, x: dict) -> torch.Tensor: # past: (B, 16, 6), intent: int past, images, intent = x['PAST'], x['IMAGES'], x['INTENT'] + + pref_traj = x['PREF_TRAJ'] # 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] @@ -124,7 +128,9 @@ def forward(self, x: dict) -> torch.Tensor: 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) + pref_traj_flat = pref_traj.view(pref_traj.size(0), -1) # (B, 42) + #print(pref_traj_flat.shape) + query = self.query(torch.cat([intent_onehot, past_flat, pref_traj_flat], dim=1)).unsqueeze(1) query = self.query_norm(query) scores = query @ key.permute((0, 2, 1)) # (B, T, N) @@ -140,7 +146,7 @@ def __init__(self, feature_extractor, out_dim, n_layers=1): self.feature_dim = sum(self.features.dims) # Initial Query Projection (Intent + Past -> C) - query_input_dim = 3 + 16 * 6 + query_input_dim = 3 + 16 * 6 + 21 * 2 self.query_init = nn.Linear(query_input_dim, self.feature_dim) # learnable positional encoding @@ -163,6 +169,7 @@ def forward(self, x): # Copied from MonocularModel # past: (B, 16, 6), intent: int past, images, intent = x['PAST'], x['IMAGES'], x['INTENT'] + pref_traj = x['PREF_TRAJ'] # 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] @@ -179,7 +186,8 @@ def forward(self, x): # 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 = self.query_init(torch.cat([intent_onehot, past_flat], dim=1)).unsqueeze(1) + pref_traj_flat = pref_traj.view(pref_traj.size(0), -1) # (B, 42) + query = self.query_init(torch.cat([intent_onehot, past_flat, pref_traj_flat], dim=1)).unsqueeze(1) for block in self.blocks: query = block(query, tokens) diff --git a/src/camera-based-e2e/preference_cache.pkl b/src/camera-based-e2e/preference_cache.pkl new file mode 100644 index 0000000000000000000000000000000000000000..abd77c574f2d8585db23b6bb62ccce52d3ac584b GIT binary patch literal 10942 zcma)Cd0b8DAHLCPBT=D}rKDB4x7#LENXsK>AzOAeDU}w6Y7&VYWy!uy)*?&&Fo_|Q zeGOyH5;eB0SsLr_yx;4i^WVMm*PN&Ke3$3DzURtNU$M_tLjR3kZk=O2ef*TMj#5Ww zsgt8j=H?`KNRCX5m>8cp$zCdzx=QWiBiaA-lrP7TlsH3dTGL4~e|eM|nxEo3;qBFq zil@_gnl@P5r>o+r?MI^_2p?jpcsg*_Y$ANk*n3LUIXa2d`G@W60{a&@j^b?)y!wnTY1Qm2dihJoJr zBcIgi^4#FI)(a&tKWD96s7x$0q4bdr0*f7(oo!$|NT)k4&8UhG z!3tM%#Z!vOj&Qv5AphF%Tt@4FObOKcGINwr=qy&R{w1UZriYm_p+*>eevbjt3`n67 zDwDoiA!Z#ZG=fR|Q(yGDLh&=g#B(205xTTHy~lIL?>(1GP;|jwxq3I5D1yz$t=Gho zy%d+u7+ImK3qKD^iq7aZ*`NdLCv{^gIwN-d4QF(@`Saw0DGU7&E-_L3ZSU|KDR}42 zSX1J=T&&#c)E5qMKPZBh*m{40BQh+zFp-w{e8oHldrj>WPh~F$#o(qj*TR}&M zV(rXxyL;cCFTvL(ZIyfM;wBPPRhW=~4gt(7B}DZUYd>fcr-mEbJ20o*(MNvA1Y=IH zToRs=xI?DUZcaqGqfNv%EA+cccDh6FcD)O1K2t@yV{;FEZ*&^-#k!M9z zx%1qYqBFSyH9r)mdCtS@l|bFmi!9rf7%P*DeB5PpPlEkLWUMz@HtVehX#yGKjnx%D zY_TMUV&x5s{4yWByha9iW5YvtKiFDP&U$0Hn%p1v^M0P}e>Y_iQKWJ(I&S!D43>;07YFm)+LVR!B)I)ARlD%~&S*lm1bgdLV9rgZ}VY6Df{vm+g`6cw9@(b23sNTo{9GdXzGg5&Rel7`mOjos9e!-)BHShU#Q8 z&()33UM;~L3srm&cefXjW?FEGCadEy^K7yk?%gBx@jMr}w4U{mBFv{ zk=R|z*iOQ{NoU$(-Zv_XB(!o}X^NXN@+k>3a%Y*rLY>qn0dA}0(CkVrGYPe;csGm; zrpPCu+q^^JXm*z(pM*ZkzbC@93vrQzrTm*@gl$sDdtp1C3HhcV+c3;o~JdLns@T=7)= zYol8J3IQU8SDbZBM0@i=?^1UdYlTQ*se;uhI{nN-4a#D7rmwF!XePN}1_d z7KIaU6jFHZsK2_9Ymy}Dthqclnuo~H>2^y%F zK<2Clf7B)j^Va+^Irabr$!zgl;w3$UIL}Z9+q~y&i5|H(>GIgeb28c5o>rsy1P7s8bI-dqYHR z!myCc1gL*tfl^|%qd3)S^q=UX)`~fv+8@+?i;2(OHTMTA?$rE1goJl;?KsF3Oi+)p6#% z5`Ue<3THLO>LKWamjto9z{W!eAoiB)m>8i-mj=71>Fi zyZ2?9um!0%V`r66;V7SNJnu&-syM57=QZ_WA*UK(UjrL0TNk{lO;Z$NN*%=>#;X=w-*le zCN}G#)wlys1O1yIwud~fw)oa42gc@oiE{>lbiT;x*RT~j)uJ#A4s#LJqO`Hvq zHjmn)PkW-_EBLyu9dV!y72{XLY+0X1^XRyBv6~g`tMJu=8f$l;plu-A};2nvZ8BijNwT*$F!Xc@G?Vq5F2aJ1ms@H zcR^AuV@&}hKRY=Rd6mpB1rTZUZyMyQnQ8)f&Zl&baH%aG#hR3o2)T&A^W3X{XNy}c zNm5H@P4OkYL0A&W_>CqY`SI%EIQzQYFThzn=Qn&Er@|H`6>I8xcbgW{Df-cEn5GNM zljI2c#zunzNL(h|(+8|!_j?6!YW&o}NHu0HS^+q0dN>Vv|FVQ7fah*c9Vbkwp@!^x zmlD5a;=#Y}WvK>6`V#3z%$gVRi)A>!nO!iV325Hy?m(=OQp6d+b5q+V2plguLZMSY zx`k!J4fw}rRy6b)&%OMjCv0nFq3o0ri(JHgwPSpCGmNOANOWcv9k{(mxD}rzX9?KR z`k4Wy&tp9-nt&0y8@eN9F)JtqFkz5n8a{6&1KA0l>uTziAq4g5FR?r~XYmfJ5oK1)z{c#J*66T{L@R*8%l^1{qrD2Ai%a^*x|w#+C?Tuu$m$l7 zf~69?SMn6RRFjQPm37=~SamOm(Z;kWQ%NSoe-5pyc<-{jpyC7r=Ad zY>Toau#0EKrbL&kxNpbm-IxwBB1OC5s&JonGuUc%kS#YY7BzMUeF z`uPiyp*VHpj`~Uo^n6w5n2us+y*@0`!|r~>vnMmaykMdYT6&T;0yfsCc;cWVGf|0W za&aqGpSQsrL!HR8{!DOXZK52{PLa(7n4M}DkHBh{g-QgwiR(SDqunqx%va?#v7M!i^FnkyhxN1 zP%xq{0c&QF=LDR2`6L-%tXMZx;<=N!Y|qd5(*W;!QDTNLO(Pp_+GAKLl`R1aI$V;W z-L0k**#kOF#P&$)JnRI|Sr1zxB+5G@6%wVVc)fA))#Fgxi}GSPQ{Jli>V9y#!j{fx z0%CPt`U@?ch#;VWKR*!8RVp-eC-HbGe%NgUE-a;XJ(_7+`t|3#OAe zEBr3hHHV)sNs1&%8gCX0MceN<%|HdDyDDDj-#WP}Ebgx2MA(@=gu_%<>U7b}?&OTv zcCcASVJ2YyM?-r&^P~<>K+N-fa?H3--I4&yyupbmDWu#XAh=*tI+hiyGVH`lP-sn< z9_ncMi%wv+Ms30Ra=jw?ranBiX_Lo=ZKMAdGseGRjj0w8B*`e9=i?r$V!K z6Z71<=!G_16+77PO#P{-V=eK-jkFQ)kG#7v+S8|`XaZvW+L^+B6B(!g%CFkvd>kbl z0arWC_rj|bGLHbW_bI?NFUmjyJbkx>VE}zKAfZB={=QwfA?6LEvS(+6q441^q13Ku z{CbV&Zm1P(l;Hb)b`}5mV(r6g0@Q7f3f;g>++a8EH8a7=wUkbCndjk4FIZy+{oW8w zKz{qh{qgN0r8fawM@+)Xr(_QUcrN!ruelPmsU^A;khIlXAvl4nS?Vh3PR`;;d<$^k zu(+NQY9X`7W&AZ41X)m(60qt+XK(0#pf!;I`ML!m`07KH5m2(UZzya}l7R#m`1y^) z=v*?7faBTSQc$wJ>6M%QRjE*KYG+T+So z>ZJsH&3NpH4iZvLfcDeG0DL(}>=KZ)(|-s$`8ocQ;p9;khJGI@t?5KgOw>&5r(c>$ z3(vg@Q4=mYud3Ku_Y*2Qs3Cms@()qL1sCz+qW9fW2bG}dxs|z)ID1A5lwM>XC8G&w zd;XCJZjDx9@so*3KRQYHZRB|&1-Ot2{u*;Q6l*If&j~OHFN(yTKEw?Jcy2ziGBt?y-Aq@(%Z-3W5phU>hwiOGX#G)z%#n#}(_UwjrSQ$2x@`#) z{O*02J-T_5%>=}#9}7lx&n6(*;L0fc7DcM*#LSSSNXT7?2m*R=NmJ3`G+E04p3~j% zHbZC~EmZMaJU_T0X*T#ihIExN^RDSG^FTu&4R-=&cDOthL2}YXz$MSpG(1vV83CtuAn+{z~4Ky9@-)U}$lbe!kqfR7`afL$kKe%P>ua)D%I{TmvJt20#y znpEr$7S?FOYaR8q%gmpHc}dzRA4+43fC}%tj`-7+Dun>;+u>#iy3KxRiYA~w*T)L? zGN>gJ@Ho??H-?`je+cM$`kx?NLz9+e?kgssyZJAnJh$Lmg)lQ#WD-s6CE@q$Lsr7f znAMNMs(>z6r3#nlJ8>!ubg4Kuqxgk}$a+CRxy#&~IOKydB7BKZ0`^>p8i+={rV|=y z7Kqmx)QjkZ*3ercp+AkhVgS!=4>}`U`1`)3maBjknA{dF{8zjFdX48+9uCZrz^Ghw z>Z+RNI^kE&5h`>5x%j$~KBrC>$(Lzx)iF2Hx18^RD;FtR1oU2>Eyq#CcqE|0R5l1L z?~q>v%zE565G$^c=LAd-FiywLNamN4N^lX|a&~Z$59}42NjPb)Jji&)ffZOO19=Ntjg}_L~WY}k^1)+@oiI);b-Nd7!&hddRB0MNK zpP7Lp!vkzF$C4~2Am^=rFxt5{0WI3?4@Hi&2^en3hv8Wv?PK&x=8olK@m5Cc5fIwJ zdj_6wQsHhU6|Y`?osa9{;vyRB&6)|u#T5$w^l0-Z>Y)S_CqJ^nWh+umKz7)Z03_u% zowzx8ObGJ(P-~zQF%ylW@UfhkuY@R9u`Sn=`|IGFI#rc6v*no29(|a!Crt#jGa7Ay zt!hLi0rP8xPdnCch$sT~_dV$iZ!M}70vwKw8iLmosl5=et8m3IJXEwg0^&MYO~j=Q zG{G3abIr7~gqvH {e}") + offset += 8 + 4 + record_len + 4 + print(len(indexes)) + continue + + if len(frame.preference_trajectories) and frame.preference_trajectories[0].preference_score != -1: + indexes.append((file, start_payload_offset, record_len)) + + # update offset for next record + offset += 8 + 4 + record_len + 4 + print(len(indexes)) + +#pickle turns objects into bytes +with open('preference_cache.pkl', 'wb') as f: + pickle.dump(indexes, f) diff --git a/src/camera-based-e2e/preference_loader.py b/src/camera-based-e2e/preference_loader.py new file mode 100644 index 0000000..6016959 --- /dev/null +++ b/src/camera-based-e2e/preference_loader.py @@ -0,0 +1,132 @@ +import torch +from torch.utils.data import IterableDataset +from protos import e2e_pb2 +import torchvision +import pickle +import struct +import os +import numpy as np +from PIL import Image +from io import BytesIO +import cv2 +from typing import Optional +import random +from models.base_model import collate_with_images + +devices = ['cuda:0', 'cuda:1'] + +random.seed(42) # Deterministic + +class WaymoE2E(IterableDataset): + def __init__( + self, + indexFile = 'preference_cache.pkl', + data_dir='./dataset', + images = True, + n_items: Optional[int] = None, + seed: Optional[int] = None, + ): + self.images = images + 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) + + if n_items is not None and n_items < len(self.indexes): + total = len(self.indexes) + # pick a deterministic contiguous block when a seed is provided + rng = random.Random(seed) if seed is not None else random + start = rng.randint(0, total - n_items) + self.indexes = self.indexes[start : start + n_items] + + def decode_img(self, img): + if not self.images: + return np.array([]) + + img_tensor = torch.from_numpy(np.frombuffer(img, dtype=np.uint8).copy()) + gpu_tensors_list = torchvision.io.decode_jpeg( + img_tensor, + mode=torchvision.io.ImageReadMode.UNCHANGED, + device= 'cpu' #['cuda:0', 'cuda:1'][torch.utils.data.get_worker_info().id%2] + ) + # img_array = np.frombuffer(img, np.uint8) + return gpu_tensors_list + + def __len__(self): + return len(self.indexes) + + def __iter__(self): + # spawns a new worker, assigns worker.id and num_workers + worker = torch.utils.data.get_worker_info() + if worker is None: + start, step = 0, 1 + else: + start, step = worker.id, worker.num_workers + + # this code makes each worker read a subset of teh data (worker 1 = 0, 16, 32..., worker 2 = 1, 17, 33...) + for idx in range(start, len(self.indexes), step): + frame = e2e_pb2.E2EDFrame() + 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) + protobuf = self.file.read(byte_length) + 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 + + for traj in frame.preference_trajectories: + # stack x, y + traj_array = np.stack([traj.pos_x, traj.pos_y], axis=-1) # shape (T, 2), T = timesteps + yield {'PAST': past, 'FUTURE': future, 'IMAGES': [self.decode_img(images.image) for images in frame.frame.images], 'INTENT': frame.intent, 'NAME': name, 'PREF_TRAJ': traj_array, 'PREF_SCORE': traj.preference_score} + +if __name__ == "__main__": + + from torch.utils.data import DataLoader + import time + from tqdm import tqdm + # NOTE: Replace with your path + DATA_DIR = '/scratch/gilbreth/svelmuru/waymo_end_to_end_dataset/waymo_open_dataset_end_to_end_camera_v_1_0_0/' + BATCH_SIZE = 32 + dataset = WaymoE2E(indexFile="preference_cache.pkl", data_dir = DATA_DIR, images=True) + loader = DataLoader( + dataset, + batch_size=BATCH_SIZE, + num_workers=16, + collate_fn=collate_with_images + ) + + #now we have this laoder class that I can iterate over like an array + #it will return dictionaries with keys PAST, FUTURE, IMAGES, INTENT + + 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) + print(len(loader)) + + import cProfile + main() \ No newline at end of file diff --git a/src/camera-based-e2e/protos/dataset.proto b/src/camera-based-e2e/protos/dataset.proto index 0615466..11d8f9e 100644 --- a/src/camera-based-e2e/protos/dataset.proto +++ b/src/camera-based-e2e/protos/dataset.proto @@ -58,7 +58,6 @@ message CameraName { REAR_LEFT = 6; REAR = 7; REAR_RIGHT = 8; - } } diff --git a/src/camera-based-e2e/protos/e2e.proto b/src/camera-based-e2e/protos/e2e.proto index 0550d65..f7cbbff 100644 --- a/src/camera-based-e2e/protos/e2e.proto +++ b/src/camera-based-e2e/protos/e2e.proto @@ -55,7 +55,7 @@ message E2EDFrame { // assigned rater scores of -1 or left empty. Valid scores range from [0, 10]. repeated EgoTrajectoryStates preference_trajectories = 8; - + // Numbers here are field numbers (1,5,6,7,8) } message EgoTrajectoryStates { diff --git a/src/camera-based-e2e/testRead.py b/src/camera-based-e2e/testRead.py new file mode 100644 index 0000000..edfa9d1 --- /dev/null +++ b/src/camera-based-e2e/testRead.py @@ -0,0 +1,44 @@ +import os +import struct +import time +import pickle +from tqdm import tqdm +import sys, os +from protos import e2e_pb2 + +DATA_DIR = '/scratch/gilbreth/svelmuru/waymo_end_to_end_dataset/waymo_open_dataset_end_to_end_camera_v_1_0_0/' +indexes = [] + +with open(os.path.join(DATA_DIR, 'val_202504211843.tfrecord-00027-of-00093_.gstmp'), 'rb') as fp: + offset = 0 + while True: + raw_length = fp.read(8) + if len(raw_length) < 8: + break + record_len = struct.unpack('Q', raw_length)[0] # unsigned + + fp.read(4) # skip first CRC + + start_payload_offset = offset + 8 + 4 # where protobuf actually starts + + data = fp.read(record_len) + if len(data) != record_len: + print(f"Incomplete record at {file}:{offset}") + break + + fp.read(4) # skip last CRC + + frame = e2e_pb2.E2EDFrame() + try: + frame.ParseFromString(data) + except Exception as e: + print(f"length of indexes: {len(indexes)}") + print(f"Failed to parse frame at {file}:{offset} -> {e}") + offset += 8 + 4 + record_len + 4 + continue + + if len(frame.preference_trajectories) and frame.preference_trajectories[0].preference_score != -1: + indexes.append((file, start_payload_offset, record_len)) + + # update offset for next record + offset += 8 + 4 + record_len + 4 \ No newline at end of file diff --git a/src/camera-based-e2e/train.py b/src/camera-based-e2e/train.py index 875c942..266813c 100644 --- a/src/camera-based-e2e/train.py +++ b/src/camera-based-e2e/train.py @@ -12,8 +12,10 @@ import torch.nn as nn import torch.nn.functional as F from pathlib import Path +import pickle -from loader import WaymoE2E +from preference_loader import WaymoE2E +from torch.utils.data import random_split, DataLoader # Replace with your model defined in models/ from models.base_model import LitModel, collate_with_images @@ -28,17 +30,48 @@ args = parser.parse_args() # Data - train_dataset = WaymoE2E(batch_size=args.batch_size, indexFile='index_train.pkl', data_dir=args.data_dir, images=True, n_items=250000) - test_dataset = WaymoE2E(batch_size=args.batch_size, indexFile='index_val.pkl', data_dir=args.data_dir, images=True, n_items=50000) + # Initialize the full dataset + full_dataset = WaymoE2E( + indexFile='preference_cache.pkl', + data_dir=args.data_dir, + images=True, + n_items=250000 + ) + + # Decide split sizes + train_size = int(0.8 * len(full_dataset)) # 80% train + val_size = len(full_dataset) - train_size # 20% val + + # Random split + # Load full dataset indexes + with open("preference_cache.pkl", "rb") as f: + indexes = pickle.load(f) + + # Shuffle for randomness + import random + random.seed(42) + random.shuffle(indexes) + # Split 80/20 + split = int(0.8 * len(indexes)) + train_indexes = indexes[:split] + val_indexes = indexes[split:] + + # Create train / val datasets + train_dataset = WaymoE2E(indexFile='preference_cache.pkl', data_dir=args.data_dir, images=True) + train_dataset.indexes = train_indexes # override indexes + + val_dataset = WaymoE2E(indexFile='preference_cache.pkl', data_dir=args.data_dir, images=True) + val_dataset.indexes = val_indexes # override indexes + train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, num_workers=12, 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=12, collate_fn=collate_with_images, persistent_workers=False, pin_memory=False) + val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, num_workers=12, collate_fn=collate_with_images, persistent_workers=False, pin_memory=False) # Model - in_dim = 16 * 6 # Past: (B, 16, 6) - out_dim = 20 * 2 # Future: (B, 20, 2) + in_dim = 16 * 6 + 3 + 21 * 2 + out_dim = 1 - model = DeepMonocularModel(feature_extractor=SAMFeatures(model_name="timm/vit_pe_spatial_tiny_patch16_512.fb"), out_dim=out_dim) + model = DeepMonocularModel(out_dim=out_dim, feature_extractor=SAMFeatures(model_name="timm/vit_pe_spatial_tiny_patch16_512.fb")) lit_model = LitModel(model=torch.compile(model, mode="max-autotune"), lr=args.lr) base_path = Path(args.data_dir).parent.as_posix() diff --git a/src/camera-based-e2e/train.slurm b/src/camera-based-e2e/train.slurm new file mode 100644 index 0000000..4b1c872 --- /dev/null +++ b/src/camera-based-e2e/train.slurm @@ -0,0 +1,20 @@ +#!/bin/bash +#SBATCH --job-name=waymo_e2e +#SBATCH --output=logs/%x_%j.out +#SBATCH --error=logs/%x_%j.err +#SBATCH --partition=v100 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=20 +#SBATCH --mem=32G +#SBATCH --time=24:00:00 +#SBATCH --account=csso + +# Load modules (edit for your cluster) +module load cuda/12.1 +module load conda + +# Activate your environment +conda activate waymo_env + +# Run training +python train.py --data_dir /scratch/gilbreth/svelmuru/waymo_end_to_end_dataset/waymo_open_dataset_end_to_end_camera_v_1_0_0/ \ No newline at end of file From b8f5c4af16968cc0b5b750b6c6de945d95edd5d4 Mon Sep 17 00:00:00 2001 From: Croopan Date: Sat, 24 Jan 2026 12:52:45 -0500 Subject: [PATCH 2/3] added correlation testing --- preference_dataset.py | 118 ++++++++++++++++ src/camera-based-e2e/models/base_model.py | 160 ++++++++++++++-------- src/camera-based-e2e/models/monocular.py | 29 ++-- src/camera-based-e2e/train.py | 23 +--- 4 files changed, 240 insertions(+), 90 deletions(-) create mode 100644 preference_dataset.py diff --git a/preference_dataset.py b/preference_dataset.py new file mode 100644 index 0000000..2032d02 --- /dev/null +++ b/preference_dataset.py @@ -0,0 +1,118 @@ +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import Dataset, DataLoader +import pickle +import os +import numpy as np +import e2e_pb2 +from tqdm import tqdm +import random + + +class PreferenceDataset(Dataset): + def __init__(self, index_file, data_dir, cache_file='preference_cache_full.pkl'): + self.data_dir = data_dir + + if os.path.exists(cache_file): + print(f"Loading filtered samples from cache: {cache_file}") + with open(cache_file, 'rb') as f: + self.valid_samples = pickle.load(f) + print(f"Loaded {len(self.valid_samples)} valid preference trajectory samples from cache") + return + + print("Cache not found. Filtering dataset for valid preference trajectories...") + with open(index_file, 'rb') as f: + all_indexes = pickle.load(f) + + self.valid_samples = [] + + # Group indexes by filename + file_groups = {} + for idx, (filename, start_byte, byte_length) in enumerate(all_indexes): + if filename not in file_groups: + file_groups[filename] = [] + file_groups[filename].append((idx, start_byte, byte_length)) + + # Process each file once + for filename, frame_list in tqdm(file_groups.items(), desc="Processing files"): + file_path = os.path.join(data_dir, filename) + + # Check if file exists + if not os.path.exists(file_path): + print(f"WARNING: File not found: {file_path}") + continue + + with open(file_path, 'rb') as f: + for original_idx, start_byte, byte_length in frame_list: + # Read and parse frame + f.seek(start_byte) + protobuf = f.read(byte_length) + frame = e2e_pb2.E2EDFrame() + frame.ParseFromString(protobuf) + print("num preference trajectories:", len(frame.preference_trajectories)) + + # Check if this frame has valid preference trajectories + for traj_idx, pref_traj in enumerate(frame.preference_trajectories): + if pref_traj.preference_score >= 0: # Valid score + self.valid_samples.append({ + 'file_info': (filename, start_byte, byte_length), + 'traj_idx': traj_idx, + 'original_idx': original_idx + }) + + print(f"Found {len(self.valid_samples)} valid preference trajectory samples") + + # Save to cache + print(f"Saving filtered samples to cache: {cache_file}") + with open(cache_file, 'wb') as f: + pickle.dump(self.valid_samples, f) + print("Cache saved!") + + def __len__(self): + return len(self.valid_samples) + + def __getitem__(self, idx): + sample_info = self.valid_samples[idx] + filename, start_byte, byte_length = sample_info['file_info'] + traj_idx = sample_info['traj_idx'] + + # Read frame + file_path = os.path.join(self.data_dir, filename) + with open(file_path, 'rb') as f: + f.seek(start_byte) + protobuf = f.read(byte_length) + frame = e2e_pb2.E2EDFrame() + frame.ParseFromString(protobuf) + + # Extract past states (T,6) + 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) + + # Extract intent (one-hot encode) + intent = frame.intent + intent_onehot = np.zeros(4, dtype=np.float32) + intent_onehot[intent] = 1.0 + + # Extract preference trajectory (only pos_x, pos_y are populated) + pref_traj = frame.preference_trajectories[traj_idx] + trajectory = np.stack([ + pref_traj.pos_x, + pref_traj.pos_y + ], axis=-1).astype(np.float32) + + # Extract preference score (target) + score = pref_traj.preference_score + + return { + 'past_states': torch.from_numpy(past.flatten()), + 'intent': torch.from_numpy(intent_onehot), + 'trajectory': torch.from_numpy(trajectory.flatten()), + 'score': torch.tensor(score, dtype=torch.float32) + } diff --git a/src/camera-based-e2e/models/base_model.py b/src/camera-based-e2e/models/base_model.py index 38d2eb2..b3a547e 100644 --- a/src/camera-based-e2e/models/base_model.py +++ b/src/camera-based-e2e/models/base_model.py @@ -2,77 +2,119 @@ import torch.nn as nn import torch.nn.functional as F import pytorch_lightning as pl - -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) +from scipy.stats import spearmanr, kendalltau +import torch + class LitModel(pl.LightningModule): - def __init__(self, model: nn.Module, lr: float): - super(LitModel, self).__init__() + def __init__(self, model: nn.Module, lr: float, delta: float = 1.0): + super().__init__() self.model = model - self.hparams.lr = lr - - 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 - 'PREF_TRAJ': torch.zeros((1, 21, 2)), # PREF_TRAJ - },) - - # ---- 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)) - - # ---- optimizers ---- + self.save_hyperparameters(ignore=["model"]) + + self.delta = delta + self.val_errors = [] + self.val_pred = [] + self.val_true = [] + + # ---- optimizer ---- def configure_optimizers(self): - # NOTE: This can be extended and tuned, LR especially will differ and have an impact. - optimizer = torch.optim.Adam(self.model.parameters(), lr=self.hparams.lr) - return optimizer - - # ---- forward / step ---- - def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.optim.Adam(self.model.parameters(), lr=self.hparams.lr) + + # ---- forward ---- + def forward(self, x): return self.model(x) - - def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor: - past, future, images, intent = batch['PAST'], batch['FUTURE'], batch['IMAGES'], batch['INTENT'] - pref_traj, pref_score = batch["PREF_TRAJ"], batch["PREF_SCORE"] - - # `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, "PREF_TRAJ": pref_traj} + # ---- shared step ---- + def _shared_step(self, batch, stage: str): + past = batch["PAST"] + images = batch["IMAGES"] + intent = batch["INTENT"] + pref_traj = batch["PREF_TRAJ"] + pref_score = batch["PREF_SCORE"] + + pred = self( + { + "PAST": past, + "IMAGES": images, + "INTENT": intent, + "PREF_TRAJ": pref_traj, + } + ).squeeze(-1) - pred_score = self.forward(model_inputs).squeeze(-1) # (B,) rather than (B, 1) - loss = F.mse_loss(pred_score, pref_score) + loss = F.smooth_l1_loss(pred, pref_score, beta=self.delta) + + self.log( + f"{stage}_loss", + loss, + batch_size=past.size(0), + prog_bar=True, + on_epoch=True, + on_step=False, + ) - # TODO: improve logging both to disk and to console - self.log_dict({ - f"{stage}_loss": loss, - }, prog_bar=True, logger=True) + if stage == "val": + err = (pred.detach() - pref_score.detach()).abs() + self.val_errors.append(err.cpu()) + # Save predictions and ground truth for correlation + self.val_pred.append(pred.detach().cpu()) + self.val_true.append(pref_score.detach().cpu()) return loss - - def training_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor: + + def training_step(self, batch, batch_idx): return self._shared_step(batch, "train") - - def validation_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor: + + def validation_step(self, batch, batch_idx): return self._shared_step(batch, "val") - + + # ---- Lightning 2.x hook ---- + def on_validation_epoch_end(self): + if not self.trainer.is_global_zero: + return + + import os + import matplotlib.pyplot as plt + import numpy as np + + if len(self.val_errors) == 0: + print("[WARN] No validation errors collected") + return + + # Concatenate validation results + errors = torch.cat(self.val_errors).to(torch.float32).numpy() + preds = torch.cat(self.val_pred).to(torch.float32).numpy() + trues = torch.cat(self.val_true).to(torch.float32).numpy() + + self.val_errors.clear() + self.val_pred.clear() + self.val_true.clear() + + # Histogram + plt.figure(figsize=(6, 4)) + plt.hist(errors, bins=50) + plt.axvline(np.mean(errors), linestyle="--", label="Mean") + plt.xlabel("|pred − gt|") + plt.ylabel("Count") + plt.title(f"Validation Error Histogram (epoch {self.current_epoch})") + plt.legend() + out_path = os.path.join("logs", f"val_error_hist_epoch_{self.current_epoch}.png") + os.makedirs(os.path.dirname(out_path), exist_ok=True) + plt.savefig(out_path) + plt.close() + print(f"[OK] Saved histogram → {out_path}") + + # Spearman and Kendall correlations + spearman_corr = spearmanr(trues, preds).correlation + kendall_corr = kendalltau(trues, preds).correlation + + print(f"[INFO] Spearman correlation: {spearman_corr:.4f}") + print(f"[INFO] Kendall tau correlation: {kendall_corr:.4f}") + + # Log to Lightning + self.log("val_spearman", spearman_corr, prog_bar=True, logger=True) + self.log("val_kendall", kendall_corr, prog_bar=True, logger=True) + # helps to batch data with images def collate_with_images(batch): past = [torch.as_tensor(b["PAST"], dtype=torch.float32) for b in batch] diff --git a/src/camera-based-e2e/models/monocular.py b/src/camera-based-e2e/models/monocular.py index ee5089a..59c626d 100644 --- a/src/camera-based-e2e/models/monocular.py +++ b/src/camera-based-e2e/models/monocular.py @@ -7,28 +7,36 @@ import timm from math import sqrt -from .base_model import BaseModel, LitModel +from .base_model import LitModel from .blocks import TransformerBlock class DINOFeatures(nn.Module): - def __init__(self, model_name: str = "vit_small_plus_patch16_dinov3.lvd1689m", frozen: bool = True): + def __init__(self, model_name: str = "vit_tiny_plus_patch16_dinov3.lvd1689m", frozen: bool = True): super(DINOFeatures, self).__init__() + # load in dino classification model + # features_only returns a list of layer outputs self.dino_model = timm.create_model(model_name, pretrained=True, features_only=True) + + #sets up input preprocessing transforms self.data_config = timm.data.resolve_data_config(model=self.dino_model) self.transforms = timm.data.create_transform(**self.data_config, is_training=False) + + # freezes the weights and doesn't continue training the model if frozen: for param in self.dino_model.parameters(): param.requires_grad = False - self.dims = [384, 384, 384] # feature dims for each layer + self.dims = [384, 384, 384] # feature the number of channels/features extracted at each layer + # intermediate layer outputs are also returned and useful + self.patch_size = 16 # patch size def forward(self, x: torch.Tensor) -> List[torch.Tensor]: # x: (B, 3, H, W) # transforms: resize 256x256, center crop, normalize - x_t = self.transforms(x.float()) # preprocess - features = self.dino_model(x_t) + x_t = self.transforms(x.float()) # preprocess input + features = self.dino_model(x_t) # 3 layers all with 384 channels/features and 16x16 patches return features # 3 x [B, 384, 16, 16] class SAMFeatures(nn.Module): @@ -49,12 +57,13 @@ def __init__( param.requires_grad = False channels = self.sam_model.feature_info.channels() - reductions = self.sam_model.feature_info.reduction() - self.feature_stage = feature_stage + reductions = self.sam_model.feature_info.reduction() #input is shrunk by cumulative factor of reductions (array) at each stage + + self.feature_stage = feature_stage # chooses which feature stage to output (-1 = last) self.dims = [channels[feature_stage]] - self.patch_size = reductions[feature_stage] # effective stride + self.patch_size = reductions[feature_stage] # sets patch size to reduction at that stage - def forward(self, x: torch.Tensor) -> List[torch.Tensor]: + def forward(self, x: torch.Tensor) -> List[torch.Tensor]: # x: (B, 3, H, W) x_t = self.transforms(x.float()) # preprocess feats = self.sam_model(x_t) # list of feature maps @@ -83,6 +92,7 @@ def __init__( self.query = nn.Sequential( nn.Linear(query_input_dim, self.feature_dim), nn.LeakyReLU(), + nn.Dropout(p=0.1), nn.Linear(self.feature_dim, self.feature_dim), ) @@ -94,6 +104,7 @@ def __init__( self.decoder = nn.Sequential( nn.Linear(self.feature_dim, self.feature_dim), nn.LeakyReLU(), + nn.Dropout(p=0.1), nn.Linear(self.feature_dim, out_dim), ) diff --git a/src/camera-based-e2e/train.py b/src/camera-based-e2e/train.py index 266813c..028b888 100644 --- a/src/camera-based-e2e/train.py +++ b/src/camera-based-e2e/train.py @@ -71,7 +71,7 @@ in_dim = 16 * 6 + 3 + 21 * 2 out_dim = 1 - model = DeepMonocularModel(out_dim=out_dim, feature_extractor=SAMFeatures(model_name="timm/vit_pe_spatial_tiny_patch16_512.fb")) + model = MonocularModel(in_dim=in_dim, out_dim=out_dim, feature_extractor=SAMFeatures(model_name="timm/vit_pe_spatial_tiny_patch16_512.fb")) lit_model = LitModel(model=torch.compile(model, mode="max-autotune"), lr=args.lr) base_path = Path(args.data_dir).parent.as_posix() @@ -92,27 +92,6 @@ trainer.fit(lit_model, train_loader, val_loader) - # Export loss graph to visualizations/ - try: - base_path = Path(base_path) - run_dir = sorted((base_path / "logs").glob("camera_e2e_*"))[-1] # newest run - metrics = pd.read_csv(run_dir / "version_0" / "metrics.csv") - train = metrics[metrics["train_loss"].notna()] - val = metrics[metrics["val_loss"].notna()] - - plt.figure() - plt.plot(train["step"], train["train_loss"], label="train_loss") - plt.plot(val["step"], val["val_loss"], label="val_loss") - plt.xlabel("Step") - plt.ylabel("Loss") - plt.legend() - plt.tight_layout() - out = Path("./visualizations") - plt.savefig(out / "loss.png", dpi=200) - except Exception as e: - print(f"Could not save loss plot: {e}") - - From 4319e90361d93856958ce32bd09b222309ea4691 Mon Sep 17 00:00:00 2001 From: Croopan Date: Sun, 22 Feb 2026 20:25:02 -0500 Subject: [PATCH 3/3] 3D Depth map --- src/camera-based-e2e/depthLoss.py | 107 +++++++ src/camera-based-e2e/loader.py | 112 ++++++-- src/camera-based-e2e/point_cloud.py | 73 +++++ src/camera-based-e2e/test.sh | 61 ++++ src/camera-based-e2e/view_pcd.py | 11 + src/camera-based-e2e/waymoModel.py | 428 ++++++++++++++++++++++++++++ 6 files changed, 771 insertions(+), 21 deletions(-) create mode 100644 src/camera-based-e2e/depthLoss.py create mode 100644 src/camera-based-e2e/point_cloud.py create mode 100755 src/camera-based-e2e/test.sh create mode 100644 src/camera-based-e2e/view_pcd.py create mode 100644 src/camera-based-e2e/waymoModel.py diff --git a/src/camera-based-e2e/depthLoss.py b/src/camera-based-e2e/depthLoss.py new file mode 100644 index 0000000..d781f95 --- /dev/null +++ b/src/camera-based-e2e/depthLoss.py @@ -0,0 +1,107 @@ +from transformers import AutoImageProcessor, AutoModelForDepthEstimation +import torch +import torch.nn.functional as F + +DEPTH_MODEL_ID = "depth-anything/Depth-Anything-V2-Small-hf" + +class DepthLoss: + def __init__(self, device): + self.device = device + self.depth_processor = AutoImageProcessor.from_pretrained(DEPTH_MODEL_ID) + self.depth_model = AutoModelForDepthEstimation.from_pretrained(DEPTH_MODEL_ID).to(self.device) + + def get_depth(self, images): + """ + Compute ground truth depth from images + + Args: + images (torch.Tensor): Input images of shape (B, C, H, W). + """ + # Preprocess images + inputs = self.depth_processor(images=images, return_tensors="pt").to(self.device) + + # Forward pass through the depth estimation model + with torch.no_grad(): + outputs = self.depth_model(**inputs) + predicted_depth = outputs.predicted_depth + + height, width = images.shape[2], images.shape[3] + + # Interpolate to original size + prediction = F.interpolate( + predicted_depth.unsqueeze(1), + size=(height, width), + mode="bicubic", + align_corners=False, + ).squeeze(1) # (B, H, W) + + return prediction + + def compute_depth_loss(self, gt_images, pred_depths, loss_fn): + pred_depth = self.get_depth(gt_images) + depth_loss = loss_fn(pred_depth, pred_depths) + + return depth_loss + + def __call__(self, gt_images, pred_depths, loss_fn=F.l1_loss): + return self.compute_depth_loss(gt_images, pred_depths, loss_fn) + + +if __name__ == "__main__": + if not torch.cuda.is_available(): + print("Must run on GPU") + + import sys + from pathlib import Path + # Determine the absolute path to the directory containing loader.py + # depth_loss.py is in models/losses/ + current_dir = Path(__file__).resolve().parent + project_root = current_dir.parent.parent + sys.path.append(str(project_root)) + + from loader import WaymoE2E + loader = WaymoE2E(indexFile="index_val.pkl", 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/", images=True) + data_iterator = iter(torch.utils.data.DataLoader(loader, batch_size=8, num_workers=4)) + + device = torch.device("cuda") + depth_loss_fn = DepthLoss(device) + + from matplotlib import pyplot as plt + import numpy as np + + for _ in range(6): + batch = next(data_iterator) + images = batch["IMAGES"][1].to(device) # front camera + + res = depth_loss_fn.get_depth(images) # (B, H, W) tensor + + fig, ax = plt.subplots(8, 2, figsize=(16, 48)) + for i in range(8): + # Input image + ax[i, 0].set_title(f"Input Image {i+1}") + ax[i, 0].imshow(images[i].permute(1, 2, 0).cpu().numpy().astype(np.uint8)) + ax[i, 0].axis('off') + + # Predicted depth + ax[i, 1].set_title(f"Predicted Depth {i+1}") + depth_img = ax[i, 1].imshow(res[i].cpu().numpy(), cmap='plasma') + ax[i, 1].axis('off') + fig.colorbar(depth_img, ax=ax[i, 1], orientation='vertical', label='Depth') + + plt.tight_layout() + fig.savefig("predicted_depth.png", dpi=150, bbox_inches='tight') + + # throughput test + from time import perf_counter + + start_time = perf_counter() + times = [] + for _ in range(100): + batch = next(data_iterator) + images = batch["IMAGES"][1].to(device) # front camera + t0 = perf_counter() + res = depth_loss_fn.get_depth(images) + times.append(perf_counter() - t0) + end_time = perf_counter() + print(f"Throughput: {100 / (end_time - start_time):.2f} batches/sec") + print(f"Avg depth inference batch/s: {1 / np.mean(times):.2f} FPS") diff --git a/src/camera-based-e2e/loader.py b/src/camera-based-e2e/loader.py index 27c1e6f..c97a14f 100644 --- a/src/camera-based-e2e/loader.py +++ b/src/camera-based-e2e/loader.py @@ -11,6 +11,10 @@ import cv2 from typing import Optional import random +import tqdm +import open3d as o3d +from point_cloud import get_waymo_intrinsics, create_point_cloud +from depthLoss import DepthLoss devices = ['cuda:0', 'cuda:1'] @@ -19,7 +23,7 @@ class WaymoE2E(IterableDataset): def __init__( self, - indexFile = 'preference_cache.pkl', + indexFile = 'index.pkl', data_dir='./dataset', images = True, n_items: Optional[int] = None, @@ -27,7 +31,7 @@ def __init__( ): self.images = images self.data_dir = data_dir - self.seed = seed + self.seed = seed self.filename = "" self.file = None @@ -37,6 +41,7 @@ def __init__( # We train on train and validate on val set self.indexes = pickle.load(f) + # 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): total = len(self.indexes) # pick a deterministic contiguous block when a seed is provided @@ -54,7 +59,7 @@ def decode_img(self, img): gpu_tensors_list = torchvision.io.decode_jpeg( img_tensor, mode=torchvision.io.ImageReadMode.UNCHANGED, - device= 'cpu' #['cuda:0', 'cuda:1'][torch.utils.data.get_worker_info().id%2] + device= 'cpu', #['cuda:0', 'cuda:1'][torch.utils.data.get_worker_info().id%2] ) # img_array = np.frombuffer(img, np.uint8) return gpu_tensors_list @@ -63,16 +68,14 @@ def __len__(self): return len(self.indexes) def __iter__(self): - # spawns a new worker, assigns worker.id and num_workers worker = torch.utils.data.get_worker_info() if worker is None: start, step = 0, 1 else: start, step = worker.id, worker.num_workers - # this code makes each worker read a subset of teh data (worker 1 = 0, 16, 32..., worker 2 = 1, 17, 33...) for idx in range(start, len(self.indexes), step): - frame = e2e_pb2.E2EDFrame() + frame = e2e_pb2.E2EDFrame() # type: ignore filename, start_byte, byte_length = self.indexes[idx] if self.filename != filename: @@ -82,12 +85,13 @@ def __iter__(self): self.file = open(os.path.join(self.data_dir, filename), 'rb') self.filename = filename - self.file.seek(start_byte) - protobuf = self.file.read(byte_length) + 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 @@ -96,8 +100,28 @@ def __iter__(self): # For submission to waymo evaluation server name = frame.frame.context.name - # returns dictionary of past, future, images, intent - yield {'PAST': past, 'FUTURE': future, 'IMAGES': [self.decode_img(images.image) for images in frame.frame.images], 'INTENT': frame.intent, 'NAME': name} + intrinsics_vec = np.zeros(6, dtype=np.float32) + + # Find FRONT camera (Enum 1) in the context + for calib in frame.frame.context.camera_calibrations: + if calib.name == 1: # 1 = FRONT + intrinsics_vec[0] = calib.intrinsic[0] # fx + intrinsics_vec[1] = calib.intrinsic[1] # fy + intrinsics_vec[2] = calib.intrinsic[2] # cx + intrinsics_vec[3] = calib.intrinsic[3] # cy + intrinsics_vec[4] = calib.width + intrinsics_vec[5] = calib.height + break + + decoded_images = [self.decode_img(img.image) for img in frame.frame.images] + + yield { + 'PAST': past, + 'FUTURE': future, + 'IMAGES': decoded_images, + 'INTRINSICS': intrinsics_vec, # <--- Passing this to main + 'NAME': frame.frame.context.name + } if __name__ == "__main__": @@ -107,24 +131,70 @@ def __iter__(self): # NOTE: Replace with your path DATA_DIR = '/scratch/gilbreth/svelmuru/waymo_end_to_end_dataset/waymo_open_dataset_end_to_end_camera_v_1_0_0/' BATCH_SIZE = 32 - dataset = WaymoE2E(indexFile="preference_cache.pkl", data_dir = DATA_DIR, images=True) + dataset = WaymoE2E(indexFile="index_train.pkl", data_dir = DATA_DIR, images=True, n_items= 2) loader = DataLoader( dataset, batch_size=BATCH_SIZE, num_workers=16, ) - #now we have this laoder class that I can iterate over like an array - #it will return dictionaries with keys PAST, FUTURE, IMAGES, INTENT - 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) - print(len(loader)) + device = torch.device("cuda") + depth_model = DepthLoss(device) + output_dir = "visualizations" + + for batch_idx, batch_of_frames in enumerate(tqdm(loader)): + images = batch_of_frames["IMAGES"][1].to(device) # Shape: (B, 3, H, W) + intrinsics_batch = batch_of_frames["INTRINSICS"] + pred_depths = depth_model.get_depth(images) # Shape: (B, H, W) + + batch_size = images.shape[0] # B + + for i in range(batch_size): + # Image: (H, W, 3) uint8 + + img_np = images[i].permute(1, 2, 0).cpu().numpy().copy() + + if img_np.max() <= 1.0: + img_np = (img_np * 255).astype(np.uint8) + else: + img_np = img_np.astype(np.uint8) + + #save img + img_filename = f"batch_{batch_idx:04d}_img_{i:02d}.png" + img_path = os.path.join(output_dir, img_filename) + + img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) + cv2.imwrite(img_path, img_bgr) + + # Depth: (H, W) float32 + depth_np = pred_depths[i].cpu().numpy() + + # set camera intrinsics + vals = intrinsics_batch[i].numpy() + camera_intrinsics = o3d.camera.PinholeCameraIntrinsic() + camera_intrinsics.set_intrinsics( + width=int(vals[4]), + height=int(vals[5]), + fx=vals[0], fy=vals[1], + cx=vals[2], cy=vals[3] + ) + + # Create Point Cloud + pcd = create_point_cloud( + img_np, + depth_np, + intrinsics=camera_intrinsics, + depth_scale=1.0 + ) + + # Save to Disk + filename = f"batch_{batch_idx:04d}_img_{i:02d}.pcd" + file_path = os.path.join(output_dir, filename) + + # non-blocking save command + o3d.io.write_point_cloud(file_path, pcd) + import cProfile main() \ No newline at end of file diff --git a/src/camera-based-e2e/point_cloud.py b/src/camera-based-e2e/point_cloud.py new file mode 100644 index 0000000..e54e60b --- /dev/null +++ b/src/camera-based-e2e/point_cloud.py @@ -0,0 +1,73 @@ +import open3d as o3d +import numpy as np + +import open3d as o3d +from protos import dataset_pb2 # Standard waymo import + +def get_waymo_intrinsics(frame, camera_name_enum): + """ + Extracts intrinsics for a specific camera from a Waymo Frame proto. + + Args: + frame: The Waymo Open Dataset Frame proto. + camera_name_enum: Integer enum for the camera (e.g., 1 for FRONT). + + Returns: + o3d.camera.PinholeCameraIntrinsic object. + """ + # 1. Find the calibration for the requested camera + calib = None + for c in frame.context.camera_calibrations: + if c.name == camera_name_enum: + calib = c + break + + # Proto definition: [f_u, f_v, c_u, c_v, k1, k2, p1, p2, k3] + fx = calib.intrinsic[0] + fy = calib.intrinsic[1] + cx = calib.intrinsic[2] + cy = calib.intrinsic[3] + + width = calib.width + height = calib.height + + # 3. Create Open3D Intrinsic Object + intrinsics = o3d.camera.PinholeCameraIntrinsic() + intrinsics.set_intrinsics(width, height, fx, fy, cx, cy) + + return intrinsics + +def create_point_cloud(rgb_image, depth_map, intrinsics=None, depth_scale=1.0): + """ + Converts a single RGB image and Depth map into a 3D Point Cloud. + + Args: + rgb_image (np.array): Shape (H, W, 3) - uint8 [0-255] + depth_map (np.array): Shape (H, W) - float32 + intrinsics (o3d.camera.PinholeCameraIntrinsic): Optional custom intrinsics. + """ + height, width = depth_map.shape + + # 1. Create Open3D Image objects + o3d_color = o3d.geometry.Image(rgb_image) + o3d_depth = o3d.geometry.Image(depth_map.astype(np.float32)) + + # 2. Create RGBD Image + rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth( + o3d_color, + o3d_depth, + depth_scale=1.0, + depth_trunc=10000.0, + convert_rgb_to_intensity=False + ) + + # 4. Back-project to Point Cloud + pcd = o3d.geometry.PointCloud.create_from_rgbd_image( + rgbd_image, + intrinsics + ) + + # Flip the point cloud because Open3D uses Y-down + pcd.transform([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]]) + + return pcd diff --git a/src/camera-based-e2e/test.sh b/src/camera-based-e2e/test.sh new file mode 100755 index 0000000..08f2ed0 --- /dev/null +++ b/src/camera-based-e2e/test.sh @@ -0,0 +1,61 @@ +#!/bin/bash + +# Configuration +PARTITIONS=("v100" "a10" "a30") +ACCOUNT="csso" +GRES="gpu:1" +MEM="128G" +CPUS="16" +TIME="8:00:00" + +echo "Checking estimated wait times for 8h job..." +echo "--------------------------------------------" + +for PART in "${PARTITIONS[@]}"; do + NOW_TS=$(date +%s) + + # Run sbatch test-only + OUTPUT=$(sbatch --test-only \ + -A $ACCOUNT \ + --gres=$GRES \ + --mem=$MEM \ + --cpus-per-task=$CPUS \ + --partition=$PART \ + -t $TIME \ + --wrap="hostname" 2>&1) + + if [[ "$OUTPUT" == *"allocation available"* ]] || [[ "$OUTPUT" == *"begins now"* ]]; then + echo "${PART}: 0:00:00 delay (Immediate)" + else + # 1. Capture the text after "start at" + EST_TIME_STRING=$(echo "$OUTPUT" | grep -oP 'start at \K.*') + + # 2. FIX: Isolate ONLY the timestamp (take the first "word" of the string) + # This removes "using 16 processors on nodes..." + CLEAN_TIMESTAMP=$(echo "$EST_TIME_STRING" | awk '{print $1}') + + if [[ -z "$CLEAN_TIMESTAMP" ]]; then + echo "${PART}: [Could not parse SLURM output] $OUTPUT" + else + # 3. Convert clean timestamp to Epoch seconds + EST_TS=$(date -d "$CLEAN_TIMESTAMP" +%s 2>/dev/null) + + if [[ -z "$EST_TS" ]]; then + echo "${PART}: [Date parse error] Raw: $CLEAN_TIMESTAMP" + else + DIFF=$((EST_TS - NOW_TS)) + if [ $DIFF -lt 0 ]; then DIFF=0; fi + + # Format output + FORMATTED_DELAY=$(date -u -d @"$DIFF" +'%H:%M:%S') + DAYS=$((DIFF / 86400)) + + if [ $DAYS -gt 0 ]; then + FORMATTED_DELAY="${DAYS}d-${FORMATTED_DELAY}" + fi + + echo "${PART}: ${FORMATTED_DELAY} delay (Est: $CLEAN_TIMESTAMP)" + fi + fi + fi +done \ No newline at end of file diff --git a/src/camera-based-e2e/view_pcd.py b/src/camera-based-e2e/view_pcd.py new file mode 100644 index 0000000..270f2b1 --- /dev/null +++ b/src/camera-based-e2e/view_pcd.py @@ -0,0 +1,11 @@ +import open3d as o3d +import sys + +# Usage: python view_pcd.py filename.pcd +if len(sys.argv) < 2: + print("Please provide a filename.") +else: + filename = sys.argv[1] + print(f"Loading {filename}...") + pcd = o3d.io.read_point_cloud(filename) + o3d.visualization.draw_geometries([pcd]) \ No newline at end of file diff --git a/src/camera-based-e2e/waymoModel.py b/src/camera-based-e2e/waymoModel.py new file mode 100644 index 0000000..2067d48 --- /dev/null +++ b/src/camera-based-e2e/waymoModel.py @@ -0,0 +1,428 @@ +# Copyright 2025 The Waymo Open Dataset Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""WOD E2E Rater Feedback Score.""" + +from typing import Dict, List, Tuple +import numpy as np + + +_THRESHOLD_TIME_SECONDS = np.array([3, 5], dtype=np.int64) +_BASE_THRESHOLDS = np.array([1.0, 1.8], dtype=np.float64) +_MINIMUM_SCORE_OUTSIDE_TRUST_REGION = 4.0 + + +def get_lat_lng_thresholds( + init_speed: np.ndarray, # [B] + lat_lng_threshold_multipliers: Tuple[float, float], +) -> Tuple[np.ndarray, np.ndarray]: + """Get lateral and longitudinal thresholds.""" + # Set and scale thresholds with the initial velocity + lat_threshold_multiplier, lng_threshold_multiplier = ( + lat_lng_threshold_multipliers + ) + + lat_thresholds = _BASE_THRESHOLDS * lat_threshold_multiplier # [2] + lng_thresholds = _BASE_THRESHOLDS * lng_threshold_multiplier # [2] + scale_by_init_speed = np.clip( + 0.5 + 0.5 * (init_speed - 1.4) / (11 - 1.4), 0.5, 1.0 + ) # [B] + lat_thresholds = scale_by_init_speed[..., None] * lat_thresholds # [B, 2] + lng_thresholds = scale_by_init_speed[..., None] * lng_thresholds # [B, 2] + + return lat_thresholds, lng_thresholds + + +def process_rater_specified_trajectories( + trajectory_batches: List[List[np.ndarray]], + trajectory_labels_batches: List[np.ndarray], + target_num_waypoints: int, + target_num_trajectories_per_batch: int, +) -> Tuple[np.ndarray, np.ndarray]: + """Processes rater-specified trajectories by truncating or padding. + + Args: + trajectory_batches: A list where each element is a batch of trajectories. A + trajectory is represented as a NumPy array of waypoints. + trajectory_labels_batches: A list where each element is a NumPy array of + labels corresponding to a batch of trajectories. + target_num_waypoints: The fixed number of waypoints each trajectory should + have after processing. Trajectories longer than this will be truncated, + and shorter ones will be padded by duplicating their last waypoint. + target_num_trajectories_per_batch: The fixed number of trajectories each + batch should have after processing. Batches with more trajectories will be + truncated, and those with fewer will be padded by duplicating their last + trajectory (and its label). + + Returns: + A tuple containing: + - processed_trajectory_batches: The processed trajectory batches after + truncation or padding them at both the trajectory and waypoint levels. + - processed_labels_batches: The corresponding processed label batches. + + Raises: + ValueError: If the number of trajectory batches and label batches do not + match, or if within a batch, the number of trajectories and labels do not + match. + """ + if len(trajectory_batches) != len(trajectory_labels_batches): + raise ValueError( + 'The number of trajectory batches and label batches must be the same.' + ) + + processed_trajectory_batches_list = [] + processed_labels_batches_list = [] + + # Iterate over each batch of trajectories and their corresponding labels + for i in range(len(trajectory_batches)): + current_trajectory_batch = trajectory_batches[i] + current_labels_batch = trajectory_labels_batches[i] + if len(current_trajectory_batch) != len(current_labels_batch): + raise ValueError( + 'In each batch, the number of trajectories and labels must be the' + ' same.' + ) + + # --- Step 1: Truncate or pad the number of trajectories in the + # current batch --- + num_trajectories_in_batch = len(current_trajectory_batch) + + if num_trajectories_in_batch > target_num_trajectories_per_batch: + # Truncate trajectories and labels + processed_batch_trajectories = current_trajectory_batch[ + :target_num_trajectories_per_batch + ] + processed_batch_labels = current_labels_batch[ + :target_num_trajectories_per_batch + ] + elif num_trajectories_in_batch < target_num_trajectories_per_batch: + # Pad trajectories and labels by duplicating the last element + num_to_pad = target_num_trajectories_per_batch - num_trajectories_in_batch + padding_trajectories = [current_trajectory_batch[-1]] * num_to_pad + padding_labels = [current_labels_batch[-1]] * num_to_pad + + processed_batch_trajectories = ( + current_trajectory_batch + padding_trajectories + ) + processed_batch_labels = current_labels_batch.tolist() + padding_labels + else: + # Number of trajectories already matches the target + processed_batch_trajectories = current_trajectory_batch + processed_batch_labels = current_labels_batch + + # --- Step 2: Truncate or pad waypoints for each trajectory + # in the processed batch --- + final_trajectories_for_current_batch = [] + + for trajectory in processed_batch_trajectories: + num_waypoints = len(trajectory) + if num_waypoints > target_num_waypoints: + # Truncate waypoints + processed_trajectory = trajectory[:target_num_waypoints] + elif num_waypoints < target_num_waypoints: + # Pad waypoints by duplicating the last waypoint + last_waypoint = trajectory[-1] + padding_waypoints = np.array( + [last_waypoint] * (target_num_waypoints - num_waypoints) + ) + processed_trajectory = np.concatenate( + (trajectory, padding_waypoints), axis=0 + ) + + else: + # Number of waypoints already matches the target + processed_trajectory = trajectory + final_trajectories_for_current_batch.append(processed_trajectory) + + processed_trajectory_batches_list.append( + np.array(final_trajectories_for_current_batch) + ) + + # Labels are already processed at the batch level (number of trajectories) + # No per-label processing is typically needed unless labels have internal + # structure to pad/truncate + if isinstance(processed_batch_labels, np.ndarray): + processed_labels_batches_list.append(processed_batch_labels) + else: # if it became a list during padding + processed_labels_batches_list.append(np.array(processed_batch_labels)) + + # Convert lists of batches to NumPy arrays + + final_processed_trajectories = np.array(processed_trajectory_batches_list) + final_processed_labels = np.array(processed_labels_batches_list) + + return final_processed_trajectories, final_processed_labels + + +def get_rater_feedback_score( + inference_trajectories: np.ndarray, # [B, I, T, 2] + inference_probs: np.ndarray, # [B, I] + rater_specified_trajectories: List[ + List[np.ndarray] + ], # [[T1, 2], [T2, 2], ...], ...] + rater_feedback_labels: List[np.ndarray], # [[P1], [P2], ...] + init_speed: np.ndarray, # [B] + lat_lng_threshold_multipliers: Tuple[float, float] = (1.0, 4.0), + decay_factor: float = 0.1, + frequency: int = 4, + length_seconds: int = 5, + default_num_of_rater_specified_trajectories: int = 3, + output_trust_region_visualization: bool = False, + minimum_score_outside_trust_region: float = _MINIMUM_SCORE_OUTSIDE_TRUST_REGION, +) -> Dict[str, np.ndarray]: + """Get rater feedback score (https://waymo.com/open/challenges/2025/e2e-driving/). + + Notations: + - B: batch size + - I: number of inference trajectories + - P: number of rater-specified trajectories + - T: number of timesteps + + Args: + inference_trajectories: An array of inference trajectories with shape [B, I, + T, 2] + inference_probs: An array of inference probabilities with shape [B, I] + rater_specified_trajectories: An array of rater-specified trajectories with + shape [B, P, T, 2] + rater_feedback_labels: An array of rater feedback labels (scores between 0 + and 10, both inclusive) with shape [B, P] + init_speed: A batch of initial velocities with shape [B] + lat_lng_threshold_multipliers: A tuple of latitude and longitude threshold + multipliers with shape [2] + decay_factor: A scalar score decay factor outside the trust region + frequency: The frequency (Hz) of trajectories to be considered. + length_seconds: The length (seconds) of trajectories to be considered. + default_num_of_rater_specified_trajectories: The default number of rater + specified trajectories to be used. + output_trust_region_visualization: Whether to output trust region + visualization. + minimum_score_outside_trust_region: The minimum score for inference + trajectories that are not fully within the trust region. + + Returns: + A dictionary of final rater feedback score and output for visualization. + """ + # We first process the rater-specified trajectories and labels by + # truncating or padding them to the same length. + # After processing, the shape of rater_specified_trajectories is + # [B, P, T, 2], and the shape of rater_feedback_labels is [B, P]. + rater_specified_trajectories, rater_feedback_labels = ( + process_rater_specified_trajectories( + rater_specified_trajectories, + rater_feedback_labels, + target_num_waypoints=length_seconds * frequency, + target_num_trajectories_per_batch=default_num_of_rater_specified_trajectories, + ) + ) + + if inference_trajectories.shape[-2] != rater_specified_trajectories.shape[-2]: + raise ValueError( + 'Inference and rater-specified trajectories must have the same number' + ' of timesteps.' + ) + + if ( + inference_trajectories.shape[-2] + < _THRESHOLD_TIME_SECONDS.max() * frequency + ): + raise ValueError( + 'Inference trajectories must have at least' + f' {_THRESHOLD_TIME_SECONDS.max()} timesteps.' + ) + + # Make rater-specified trajectories to include the origin + padded_rater_specified_trajectories = np.pad( + rater_specified_trajectories, + ((0, 0), (0, 0), (1, 0), (0, 0)), + constant_values=0, + ) # [B, P, T + 1, 2] + + # Compute displacement vectors + displacement_vectors = ( + padded_rater_specified_trajectories[..., 1:, :] + - padded_rater_specified_trajectories[..., :-1, :] + ) # [B, P, T, 2] + + # Get unnormalized directions + lng_directions = displacement_vectors # [B, P, T, 2] + + # When displacement is zero, which means the vehicle did not move, we bring + # longitudinal directions from the previous timestep. + lng_magnitudes = np.linalg.norm(lng_directions, axis=-1) # [B, P, T] + + # At the first timestep, we set the longitudinal directions to be (1, 0). + # This is because the vehicle coordinate is used. + lng_directions[..., 0, 0] = np.where( + lng_magnitudes[..., 0] == 0, + 1, + lng_directions[..., 0, 0], + ) # x-axis + lng_directions[..., 0, 1] = np.where( + lng_magnitudes[..., 0] == 0, + 0, + lng_directions[..., 0, 1], + ) # y-axis + + # For the rest of the timesteps, we bring the longitudinal directions from the + # previous timestep. + for t in range(1, lng_directions.shape[2]): + lng_directions[..., t, 0] = np.where( + lng_magnitudes[..., t] == 0, + lng_directions[..., t - 1, 0], + lng_directions[..., t, 0], + ) # x-axis + lng_directions[..., t, 1] = np.where( + lng_magnitudes[..., t] == 0, + lng_directions[..., t - 1, 1], + lng_directions[..., t, 1], + ) # y-axis + + # Lateral directions are 90-degree counterclockwise rotation of longitudinal + # directions, i.e., (x_new, y_new) = (-y, x) + lat_directions = np.stack( + [lng_directions[..., 1] * -1, lng_directions[..., 0]], axis=-1 + ) # [B, P, T, 2] + + # Normalize directions + lng_directions = lng_directions / np.linalg.norm( + lng_directions, axis=-1, keepdims=True + ) # [B, P, T, 2] + lat_directions = lat_directions / np.linalg.norm( + lat_directions, axis=-1, keepdims=True + ) # [B, P, T, 2] + + # Get longitudinal and lateral distances from rater-specified trajectories + rater_specified_to_inference_vectors = ( + inference_trajectories[..., None, :, :, :] + - rater_specified_trajectories[..., None, :, :] + ) # [B, 1, I, T, 2] - [B, P, 1, T, 2] --> [B, P, I, T, 2] + lng_projections = np.sum( + lng_directions[..., None, :, :] * rater_specified_to_inference_vectors, + axis=-1, + ) # [B, P, I, T], directions are broadcasted to the inference trajectories + lat_projections = np.sum( + lat_directions[..., None, :, :] * rater_specified_to_inference_vectors, + axis=-1, + ) # [B, P, I, T], directions are broadcasted to the inference trajectories + lng_distances = np.abs(lng_projections) # [B, P, I, T] + lat_distances = np.abs(lat_projections) # [B, P, I, T] + + # Filter distances at 3 and 5 seconds + selected_indices = _THRESHOLD_TIME_SECONDS * frequency - 1 + lng_distances = lng_distances[..., selected_indices] # [B, P, I, 2] + lat_distances = lat_distances[..., selected_indices] # [B, P, I, 2] + + lat_thresholds, lng_thresholds = get_lat_lng_thresholds( + init_speed, lat_lng_threshold_multipliers + ) + + outputs = {} + + # --------------------------------------------------------------------------- + # Visualization + # --------------------------------------------------------------------------- + if output_trust_region_visualization: + center_x = rater_specified_trajectories[..., selected_indices, :][ + ..., 0 + ] # [B, P, T (=2)] + center_y = rater_specified_trajectories[..., selected_indices, :][ + ..., 1 + ] # [B, P, T (=2)] + width = 2 * lng_thresholds # [B, 2] + height = 2 * lat_thresholds # [B, 2] + angle = np.degrees( + np.arctan2( + displacement_vectors[..., selected_indices, :][..., 1], + displacement_vectors[..., selected_indices, :][..., 0], + ) + ) # [B, P, T (=2)] + + outputs.update({ + 'trust_region_center_x': center_x, # [B, I, 2] + 'trust_region_center_y': center_y, # [B, I, 2] + 'trust_region_width': width, # [B, 2] + 'trust_region_height': height, # [B, 2] + 'trust_region_angle': angle, # [B, I, 2] + }) + + # --------------------------------------------------------------------------- + # Hard matching with decaying + # --------------------------------------------------------------------------- + + # Normalize distances with thresholds + normalized_lng_distances = ( + lng_distances / lng_thresholds[..., None, None, :] + ) # [B, P, I, 2] + normalized_lat_distances = ( + lat_distances / lat_thresholds[..., None, None, :] + ) # [B, P, I, 2] + + # Pick the maximum of the two normalized distances + normalized_distances = np.maximum( + normalized_lng_distances, normalized_lat_distances + ) # [B, P, I, 2] + + # Mask to indicate if the inference trajectory is fully within the trust + # region, i.e., distance from trajectory i is near any rated trajectory p. + # For inferences not fully within the trust region, scores are clipped to + # `minimum_score_outside_trust_region` during score computation below. + # [B, P, I, 2] -> [B, I] + is_fully_within_trust_region = np.any( + np.all(normalized_distances <= 1.0, axis=3), axis=1) + outputs['is_fully_within_trust_region'] = ( + is_fully_within_trust_region # [B, I] + ) + + # Make scores flat within the trust region. + exponent = np.maximum(normalized_distances - 1.0, 0.0) + decay = decay_factor**exponent + # Scores between every inference i and rated trajectory p along x,y axes. + rater_feedback_scores_per_axis_pairwise = ( + rater_feedback_labels[..., None, None] * decay + ) # [B, P, I, 2] + + # Scores for each inference trajectory along x,y axes. + # Each inference trajectory is assigned a score based on its best match with + # a rated trajectory. + rater_feedback_score_per_axis_per_inference = np.amax( + rater_feedback_scores_per_axis_pairwise, axis=1 + ) # [B, I, 2] + + # Scores for each inference trajectory averaged over x,y axes. + rater_feedback_score_per_inference = np.mean( + rater_feedback_score_per_axis_per_inference, + axis=-1, + ) # [B, I] + + # Clip scores for inferences not fully within the trust region. + rater_feedback_score_per_inference[~is_fully_within_trust_region] = ( + np.maximum( + minimum_score_outside_trust_region, + rater_feedback_score_per_inference[~is_fully_within_trust_region], + ) + ) # [B, I] + + # Weighted sum over scores for each inference trajectory. + rater_feedback_score = np.sum( + rater_feedback_score_per_inference * inference_probs, axis=-1 + ) # [B] + outputs['rater_feedback_score'] = rater_feedback_score # [B] + # Updated the truncated or padded rater-specified trajectories. + # [B, P, T, 2] + outputs['rater_specified_trajectories'] = rater_specified_trajectories + # Updated the truncated or padded rater feedback labels. + # [B, P] + outputs['rater_feedback_labels'] = rater_feedback_labels + + return outputs \ No newline at end of file