import os.path as osp import os import sys import json import itertools import time from collections import deque import torch import tqdm import concurrent.futures import psutil import io from models.SpaTrackV2.models.utils import matrix_to_quaternion import cv2 from PIL import Image import numpy as np from models.SpaTrackV2.datasets.base_sfm_dataset import BaseSfMViewDataset from models.SpaTrackV2.models.utils import ( camera_to_pose_encoding, pose_encoding_to_camera ) from models.SpaTrackV2.models.camera_transform import normalize_cameras from models.SpaTrackV2.datasets.dataset_util import imread_cv2, npz_loader,resize_crop_video import glob def bytes_to_gb(bytes): return bytes / (1024 ** 3) def get_total_size(obj, seen=None): size = sys.getsizeof(obj) if seen is None: seen = set() obj_id = id(obj) if obj_id in seen: return 0 seen.add(obj_id) if isinstance(obj, dict): size += sum([get_total_size(v, seen) for v in obj.values()]) size += sum([get_total_size(k, seen) for k in obj.keys()]) elif hasattr(obj, '__dict__'): size += get_total_size(obj.__dict__, seen) elif hasattr(obj, '__iter__') and not isinstance(obj, (str, bytes, bytearray)): size += sum([get_total_size(i, seen) for i in obj]) return size class Waymo(BaseSfMViewDataset): def __init__(self, mask_bg=False, scene_st=None, scene_end=None, debug=False, *args, ROOT, **kwargs): self.ROOT = ROOT super().__init__(*args, **kwargs) assert mask_bg in (True, False, 'rand') self.mask_bg = mask_bg self.dataset_label = 'Waymo' # load all scenes self.scene_list = os.listdir(self.ROOT) # get the scene info self.scene_list = [os.path.join(self.ROOT,dir_i) for dir_i in self.scene_list if os.path.isdir(os.path.join(self.ROOT,dir_i))] def __len__(self): return len(self.scene_list) def _get_metadatapath(self, obj, instance, view_idx): return osp.join(self.ROOT, obj, instance, 'images', f'frame{view_idx:06n}.npz') def _get_impath(self, obj, instance, view_idx): return osp.join(self.ROOT, obj, instance, 'images', f'frame{view_idx:06n}.jpg') def _get_depthpath(self, obj, instance, view_idx): return osp.join(self.ROOT, obj, instance, 'depths', f'frame{view_idx:06n}.jpg.geometric.png') def _get_maskpath(self, obj, instance, view_idx): return osp.join(self.ROOT, obj, instance, 'masks', f'frame{view_idx:06n}.png') def _read_depthmap(self, depthpath, input_metadata=None): depthmap = imread_cv2(depthpath, cv2.IMREAD_UNCHANGED) depthmap = depthmap.astype(np.float32) return depthmap def _get_views(self, idx, resolution, rng): # choose a scene scene_root = self.scene_list[idx] # random choose view view_pick = np.random.randint(1, 6) cam_root = osp.join(scene_root, f"cam_{view_pick}") # image list imgs_pool = sorted(glob.glob(osp.join(cam_root, "*_rgb.jpg"))) T = len(imgs_pool) # randomly choose a scene sclae_num = int(np.random.uniform(1, 2)) start = np.random.choice(np.arange(0, max(T - sclae_num*self.num_views, 1))) idxs = np.arange(start, start+sclae_num*self.num_views, sclae_num).clip(0, T-1) images_pick = np.array(imgs_pool)[idxs]# choose the image extrs = [] intrs = [] rgbs = [] depths = [] import pdb; pdb.set_trace() for i, img_dir_i in enumerate(images_pick): img_dir = img_dir_i depth_dir = img_dir.replace("_rgb.jpg", "_depth.npz") cam_dir = img_dir.replace("_rgb.jpg", "_pose.npy") # load rgb and depth rgb = imread_cv2(img_dir) depth = np.load(depth_dir, allow_pickle=True)['data'] cam = np.load(cam_dir, allow_pickle=True).item() rgbs.append(rgb) depths.append(depth) # load pose extrs.append(cam['extr']) intrs.append(cam['intr']) extrs = np.stack(extrs, axis=0) intrs = np.stack(intrs, axis=0) rgbs = np.stack(rgbs, axis=0) depths = np.stack(depths, axis=0) # crop and resize rgbs, depths, Intrs = resize_crop_video(rgbs, depths, intrs, resolution[0]) # encode the camera poses Extrs = torch.from_numpy(extrs) camera_poses = torch.inverse(Extrs) #NOTE: C2W focal0 = Intrs[:, 0, 0] / resolution[0] focal1 = Intrs[:, 1, 1] / resolution[0] focal = (focal0.unsqueeze(1)+focal1.unsqueeze(1))/2 # first frame normalize camera_poses = torch.inverse(camera_poses[:1]) @ camera_poses T_center = camera_poses[:, :3, 3].mean(dim=0) Radius = (camera_poses[:, :3, 3].norm(dim=1).max()) # if Radius < 1e-2: Radius = 1 camera_poses[:, :3, 3] = (camera_poses[:, :3, 3])/Radius R = camera_poses[:, :3, :3] t = camera_poses[:, :3, 3] rot_vec = matrix_to_quaternion(R) pose_enc = torch.cat([t, rot_vec, focal], dim=1) # depth_cano = Radius*focal[:,:,None,None] / depths.clamp(min=1e-6) depth_cano = depths / Radius traj_3d = torch.zeros(self.num_views, self.track_num, 3) vis = torch.zeros(self.num_views, self.track_num) syn_real = torch.tensor([1]) metric_rel = torch.tensor([1]) static = torch.tensor([0]) data_dir = cam_root views = dict( rgbs=rgbs, depths=depth_cano, pose_enc=pose_enc, traj_mat=camera_poses, intrs=Intrs, traj_3d=traj_3d, vis=vis, syn_real=syn_real, metric_rel=metric_rel, data_dir=data_dir ) return views if __name__ == "__main__": from models.SpaTrackV2.datasets.base_sfm_dataset import view_name from functools import partial # from dust3r.viz import SceneViz, auto_cam_size DATA_DIR = "/mnt/bn/haotongdata3/Datasets/WAYMO_FULL/training/data/" dataset = Waymo(split='train', ROOT=DATA_DIR, resolution=518, aug_crop=16, num_views=32) rng = np.random.default_rng(seed=0) data_ret = dataset._get_views(0,(518,518),rng) from models.SpaTrackV2.datasets.vis3d_check import vis4d vis4d(data_ret["rgbs"], data_ret["depths"], data_ret["traj_mat"], data_ret["intrs"], track3d=data_ret["traj_3d"], workspace="/mnt/bn/xyxdata/home/codes/my_projs/SpaTrack2/viser_result/test") import pdb; pdb.set_trace()