| |
| |
| |
| import torch |
| from torch.utils.data import Dataset |
| import cv2 |
| import json |
| from pathlib import Path |
| from torchvision import transforms |
| import os |
|
|
| class LMDriveDataset(Dataset): |
| def __init__(self, data_dir, transform=None, lidar_transform=None): |
| self.data_dir = Path(data_dir) |
| self.transform = transform or self.get_default_transform() |
| self.lidar_transform = lidar_transform or self.get_default_lidar_transform() |
| self.samples = [] |
| measurement_dir = self.data_dir / "measurements" |
| image_dir = self.data_dir / "rgb_full" |
| if not measurement_dir.exists() or not image_dir.exists(): |
| print(f"Warning: Data directory {data_dir} not found or incomplete.") |
| return |
| measurement_files = sorted(measurement_dir.glob("*.json")) |
| for measurement_path in measurement_files: |
| frame_id = int(measurement_path.stem) |
| image_path = image_dir / f"{frame_id:04d}.jpg" |
| if not image_path.exists(): continue |
| with open(measurement_path, "r") as f: measurements_data = json.load(f) |
| self.samples.append({"image_path": str(image_path),"measurement_path": str(measurement_path),"frame_id": frame_id,"measurements": measurements_data}) |
| def __len__(self): return len(self.samples) |
| def __getitem__(self, idx): |
| sample = self.samples[idx]; full_image = cv2.imread(sample["image_path"]) |
| if full_image is None: raise ValueError(f"Failed to load image: {sample['image_path']}") |
| full_image = cv2.cvtColor(full_image, cv2.COLOR_BGR2RGB) |
| front_image = full_image[:600, :800]; left_image = full_image[600:1200, :800] |
| right_image = full_image[1200:1800, :800]; center_image = full_image[1800:2400, :800] |
| front_tensor=self.transform(front_image);left_tensor=self.transform(left_image);right_tensor=self.transform(right_image);center_tensor=self.transform(center_image) |
| lidar_path = str(self.data_dir / "lidar" / f"{sample['frame_id']:04d}.png") |
| if os.path.exists(lidar_path): |
| lidar = cv2.imread(lidar_path) |
| if len(lidar.shape) == 2: lidar = cv2.cvtColor(lidar, cv2.COLOR_GRAY2BGR) |
| lidar = cv2.cvtColor(lidar, cv2.COLOR_BGR2RGB) |
| else: lidar = np.zeros((112, 112, 3), dtype=np.uint8) |
| lidar_tensor = self.lidar_transform(lidar) |
| m = sample["measurements"] |
| target_point = torch.tensor([m.get("x_command", 0.0), m.get("y_command", 0.0)], dtype=torch.float32) |
| measurements = torch.tensor([m.get("x",0.0),m.get("y",0.0),m.get("theta",0.0),m.get("speed",0.0),m.get("steer",0.0),m.get("throttle",0.0),int(m.get("brake",False)),m.get("command",0),int(m.get("is_junction",False)),int(m.get("should_brake",0))], dtype=torch.float32) |
| return {"rgb":front_tensor,"rgb_left":left_tensor,"rgb_right":right_tensor,"rgb_center":center_tensor,"lidar":lidar_tensor,"measurements":measurements,"target_point":target_point} |
| def get_default_transform(self): |
| return transforms.Compose([transforms.ToPILImage(),transforms.Resize((224,224)),transforms.ToTensor(),transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])]) |
| def get_default_lidar_transform(self): |
| return transforms.Compose([transforms.ToPILImage(),transforms.Resize((112,112)),transforms.ToTensor(),transforms.Normalize(mean=[0.5],std=[0.5])]) |
|
|
