Ben dataloader v0

.gitignore

@@ -3,6 +3,9 @@ __pycache__/
 *.py[codz]
 *$py.class
 
+# env setup script
+setup.sh
+
 # C extensions
 *.so
 

README.md

@@ -1,2 +1,2 @@
 # robotvision
-Machine Learning @ Purdue - Robot Vision Team (2025-2026
+Machine Learning @ Purdue - Robot Vision Team (2025-2026)

copy_lol.py

@@ -0,0 +1,13 @@
+import shutil
+import time
+import os
+DATA_DIR = '/scratch/gilbreth/bnamikas/data/waymo_open_dataset_end_to_end_camera_v_1_0_0/training_202504031202_202504151040.tfrecord-00111-of-00263'
+DEST = '/tmp/cache/'
+if not os.path.exists(DEST):
+    os.mkdir(DEST)
+start=time.time()
+try:
+    shutil.copy(DATA_DIR, DEST)
+except FileExistsError as e:
+    pass
+print('Total time', time.time()-start)

indexer.py

@@ -0,0 +1,30 @@
+should_index = True
+if should_index: # Takes approx. 9 mins to index
+    import os
+    import mmap
+    import struct
+    import time
+    import pickle
+    from tqdm import tqdm
+
+    DATA_DIR = '/scratch/gilbreth/bnamikas/data/waymo_open_dataset_end_to_end_camera_v_1_0_0'
+    # DATA_DIR = './data'
+
+    indexes = []
+
+    start = time.time()
+    for i, fn in enumerate(tqdm([file for file in os.listdir(DATA_DIR) if '.tfrecord' in file and file.startswith('train')])):
+        with open(os.path.join(DATA_DIR, fn), 'rb') as file:
+            while True:
+                blenth = file.read(8)
+                if len(blenth) == 0:
+                    break
+                proto_len = struct.unpack('q', blenth)[0]
+                indexes.append((fn, file.tell()+4, proto_len)) #file.tell()+4 so that I can make sure we skip checksum
+                file.seek(file.tell() + proto_len+8)
+
+
+    with open('index.pkl', 'wb') as f:
+        pickle.dump(indexes, f)
+
+    print("Done indexing")

loader.py

@@ -0,0 +1,70 @@
+# HISTORICAL, loader1-3 have memory leaks and issues
+import torch
+from torch.utils.data import Dataset
+from protos import e2e_pb2
+import pickle
+import struct
+import os
+import numpy as np
+from PIL import Image
+from io import BytesIO 
+import cv2 
+
+class WaymoE2E(Dataset): 
+    def __init__(self, indexFile = 'index_local.pkl', data_dir='./dataset'):
+        self.DATA_DIR = data_dir
+
+        with open(indexFile, 'rb') as f:
+            self.indexes = pickle.load(f)
+
+    def decode_img(self, img):
+        img_array = np.frombuffer(img, np.uint8)
+        return cv2.imdecode(img_array, cv2.IMREAD_COLOR)
+
+
+    def __len__(self):
+        return len(self.indexes)
+
+    def __getitem__(self, idx):
+        filename, offset, size = self.indexes[idx]
+        full_filepath = os.path.join(self.DATA_DIR, filename)
+
+
+        with open(full_filepath, 'rb') as f:
+            # pass
+            f.seek(offset)
+            # blenth = f.read(8)
+            # proto_len = struct.unpack('q', blenth)[0]
+            # f.read(4)
+            protobuff = f.read(size)
+
+        frame = e2e_pb2.E2EDFrame()
+        frame.ParseFromString(protobuff)
+        return np.vstack([self.decode_img(images.image) for images in frame.frame.images]), np.array(list(zip(frame.future_states.pos_x, frame.future_states.pos_y, frame.future_states.pos_z)))
+
+from torch.utils.data import DataLoader
+import time
+from tqdm import tqdm
+# DATA_DIR = '/scratch/gilbreth/bnamikas/data/waymo_open_dataset_end_to_end_camera_v_1_0_0'
+DATA_DIR = './data'
+# DATA_DIR = '/tmp/'
+
+
+def main():
+    dataset = WaymoE2E(data_dir = DATA_DIR)
+
+    loader = DataLoader(
+        dataset, 
+        batch_size=32,
+        num_workers=16,
+        shuffle=False
+    )
+
+    start = time.time()
+    for batch_of_frames in tqdm(loader):
+        pass
+    print("Total Time:", time.time()-start)
+
+import cProfile
+cProfile.run('main()')
+# main()

loader1.py

@@ -0,0 +1,61 @@
+# HISTORICAL, loader1-3 have memory leaks and issues
+import torch
+from torch.utils.data import Dataset
+from protos import e2e_pb2
+import pickle
+import struct
+import os
+import numpy as np
+from PIL import Image
+from io import BytesIO 
+
+class WaymoE2E(Dataset): 
+    def __init__(self, indexFile = 'index.pkl', data_dir='./dataset'):
+        self.DATA_DIR = data_dir
+
+        with open(indexFile, 'rb') as f:
+            self.indexes = pickle.load(f)
+
+        self.file = open(os.path.join(self.DATA_DIR, self.indexes[0][0]), 'rb')
+        self.filename=self.indexes[0][0]
+
+    def decode_img(self, img):
+        return np.array(Image.open(BytesIO(img)))
+
+
+    def __len__(self):
+        return len(self.indexes)
+
+    def __getitem__(self, idx):
+        frame = e2e_pb2.E2EDFrame()
+
+        if self.filename != self.indexes[idx][0]:
+            self.file.close()
+            self.file = open(os.path.join(self.DATA_DIR, self.indexes[idx][0]), 'rb')
+            self.filename = self.indexes[idx][0]
+
+            # pass
+        self.file.seek(self.indexes[idx][1])
+        protobuff = self.file.read(self.indexes[idx][2])
+        frame.ParseFromString(protobuff)
+
+        return 1 # np.vstack([self.decode_img(images.image) for images in frame.frame.images]), np.array(list(zip(frame.future_states.pos_x, frame.future_states.pos_y, frame.future_states.pos_z)))
+
+from torch.utils.data import DataLoader
+import time
+from tqdm import tqdm
+DATA_DIR = '/scratch/gilbreth/bnamikas/data/waymo_open_dataset_end_to_end_camera_v_1_0_0'
+
+dataset = WaymoE2E(data_dir = DATA_DIR)
+
+loader = DataLoader(
+    dataset, 
+    batch_size=32,
+    num_workers=0,
+    shuffle=False
+)
+
+start = time.time()
+for batch_of_frames in tqdm(loader):
+    pass
+print("Total Time:", time.time()-start)

loader1m.py

@@ -0,0 +1,56 @@
+# HISTORICAL, loader1-3 have memory leaks and issues
+import torch
+from torch.utils.data import Dataset
+from protos import e2e_pb2
+import pickle
+import struct
+import os
+import numpy as np
+from PIL import Image
+from io import BytesIO 
+
+class WaymoE2E(Dataset): 
+    def __init__(self, indexFile='index.pkl', data_dir='./dataset'):
+        self.DATA_DIR = data_dir
+
+        with open(indexFile, 'rb') as f:
+            self.indexes = pickle.load(f)
+
+
+    def decode_img(self, img):
+        return np.array(Image.open(BytesIO(img)))
+
+    def __len__(self):
+        return len(self.indexes)
+
+    def __getitem__(self, idx):
+        filename, offset, size = self.indexes[idx]
+        full_filepath = os.path.join(self.DATA_DIR, filename)
+
+        with open(full_filepath, 'rb') as f:
+            f.seek(offset)
+            protobuff = f.read(size)
+
+        frame = e2e_pb2.E2EDFrame()
+        frame.ParseFromString(protobuff)
+
+        return 1 # np.vstack([self.decode_img(images.image) for images in frame.frame.images]), np.array(list(zip(frame.future_states.pos_x, frame.future_states.pos_y, frame.future_states.pos_z)))
+
+from torch.utils.data import DataLoader
+import time
+from tqdm import tqdm
+DATA_DIR = '/scratch/gilbreth/bnamikas/data/waymo_open_dataset_end_to_end_camera_v_1_0_0'
+
+dataset = WaymoE2E(data_dir = DATA_DIR)
+
+loader = DataLoader(
+    dataset, 
+    batch_size=32,
+    num_workers=8,
+    shuffle=False
+)
+
+start = time.time()
+for batch_of_frames in tqdm(loader):
+    pass
+print("Total Time:", time.time()-start)

loader2.py

@@ -0,0 +1,72 @@
+# HISTORICAL, loader1-3 have memory leaks and issues
+import torch
+from torch.utils.data import Dataset
+from protos import e2e_pb2
+import pickle
+import struct
+import os
+import numpy as np
+from PIL import Image
+from io import BytesIO 
+import mmap
+import cv2
+
+class WaymoE2E(Dataset): 
+    def __init__(self, indexFile = 'index.pkl', data_dir='./dataset'):
+        self.DATA_DIR = data_dir
+
+        with open(indexFile, 'rb') as f:
+            self.indexes = pickle.load(f)
+
+        all_paths = list(set([index[0] for index in self.indexes]))
+        self.fds = [(os.open(os.path.join(self.DATA_DIR, fp), os.O_RDONLY), fp) for fp in all_paths]
+        self.mmaps = {fd[1]: mmap.mmap(fd[0], length=0, access=mmap.ACCESS_READ) for fd in self.fds}
+
+    def decode_img(self, img):
+        img_array = np.frombuffer(img, np.uint8)
+        return cv2.imdecode(img_array, cv2.IMREAD_COLOR)
+
+
+    def __len__(self):
+        return len(self.indexes)
+
+    def __getitem__(self, idx):
+        frame = e2e_pb2.E2EDFrame()
+
+        # self.mmaps[self.indexes[idx][0]].seek(self.indexes[idx][1])
+        # frame.ParseFromString(self.mmaps[self.indexes[idx][0]].read(self.indexes[idx][2]))
+        frame.ParseFromString(self.mmaps[self.indexes[idx][0]][self.indexes[idx][1]:self.indexes[idx][1]+self.indexes[idx][2]])
+
+        return 1#np.vstack([self.decode_img(images.image) for images in frame.frame.images]), np.array(list(zip(frame.future_states.pos_x, frame.future_states.pos_y, frame.future_states.pos_z)))
+
+from torch.utils.data import DataLoader
+import time
+from tqdm import tqdm
+DATA_DIR = '/scratch/gilbreth/bnamikas/data/waymo_open_dataset_end_to_end_camera_v_1_0_0'
+
+
+
+def main():
+    try:
+        dataset = WaymoE2E(data_dir = DATA_DIR)
+
+        loader = DataLoader(
+            dataset, 
+            batch_size=32,
+            num_workers=8,
+            shuffle=True
+        )
+
+        start = time.time()
+        for batch_of_frames in tqdm(loader):
+            pass
+        print("Total Time:", time.time()-start)
+    except Exception as e:
+        print(e)
+    finally:
+        [mmap.close() for mmap in dataset.mmaps.values()]
+        [os.close(fd[0]) for fd in dataset.fds]
+
+import cProfile
+# cProfile.run('main()')
+main()