|
|
| import os |
| import cv2 |
| import random |
| import numpy as np |
| from PIL import Image |
| |
| import torch |
| from torch.utils.data import Dataset, DataLoader |
| from torch.utils.data.sampler import Sampler |
| from torchvision import transforms, utils |
| import random |
| def get_stamp_list(dataset, timestamp): |
| base_dataset = getattr(dataset, "dataset", None) |
| if base_dataset is None: |
| raise AttributeError("Expected dataset to expose the underlying data via `.dataset`.") |
|
|
| if hasattr(base_dataset, "poses"): |
| pose_count = len(base_dataset.poses) |
| else: |
| |
| if len(base_dataset) == 0: |
| raise ValueError("Cannot build timestamp list from an empty dataset.") |
| camera_ids = [] |
| for item in base_dataset: |
| image_name = getattr(item, "image_name", None) |
| if not isinstance(image_name, str): |
| continue |
| camera_ids.append(image_name.split("/")[0]) |
| if not camera_ids: |
| raise AttributeError("Unable to infer unique camera ids from dataset.") |
| pose_count = len(set(camera_ids)) |
|
|
| if pose_count == 0: |
| raise ValueError("Pose count inferred as zero; cannot compute timestamp slices.") |
|
|
| frame_length = int(len(dataset) / pose_count) |
| if frame_length == 0: |
| raise ValueError("Frame length inferred as zero; check dataset ordering.") |
|
|
| if len(dataset) % pose_count != 0: |
| print(f"[WARN] Dataset length {len(dataset)} not divisible by camera count {pose_count}.") |
|
|
| if timestamp >= frame_length: |
| raise IndexError("input timestamp bigger than total timestamp.") |
|
|
| indices = [i * frame_length + timestamp for i in range(pose_count)] |
| return [dataset[idx] for idx in indices] |
|
|
|
|
| class FineSampler(Sampler): |
| def __init__(self, dataset): |
| self.len_dataset = len(dataset) |
| self.len_pose = len(dataset.dataset.poses) |
| self.frame_length = int(self.len_dataset/ self.len_pose) |
|
|
| sample_list = [] |
| for i in range(self.frame_length): |
| for j in range(4): |
| idx = torch.randperm(self.len_pose) *self.frame_length + i |
| now_list = [] |
| cnt = 0 |
| for item in idx.tolist(): |
| now_list.append(item) |
| cnt+=1 |
| if cnt % 2 == 0 and len(sample_list)>2: |
| select_element = [x for x in random.sample(sample_list,2)] |
| now_list += select_element |
| |
| sample_list += now_list |
| |
| self.sample_list = sample_list |
| def __iter__(self): |
|
|
| return iter(self.sample_list) |
| |
| def __len__(self): |
| return len(self.sample_list) |
|
|
