import os import glob import typing import imageio import numpy as np import json import torch import torch.nn.functional as F from torch.utils.data import Dataset from body_model import MANO_JOINTS assert len(MANO_JOINTS) == 16 import time from util.logger import Logger from geometry.camera import invert_camera from .tools import read_keypoints, read_mask_path, load_mano_preds, load_keypoints_with_interp from .vidproc import preprocess_cameras, preprocess_frames, preprocess_tracks """ Define data-related constants """ # DEFAULT_GROUND = np.array([0.0, -1.0, 0.0, -0.5]) # XXX: TEMPORARY CONSTANTS SHOT_PAD = 0 MIN_SEQ_LEN = 10 MAX_NUM_TRACKS = 12 MIN_TRACK_LEN = 60 MIN_KEYP_CONF = 0.4 def get_dataset_from_cfg(cfg): args = cfg.data if not args.use_cams: args.sources.cameras = "" args.sources = expand_source_paths(args.sources) print("DATA SOURCES", args.sources) a = time.time() check_data_sources(args, cfg) print('end of check_data_sources, time: ', time.time() - a) return MultiPeopleDataset( args.sources, args.seq, tid_spec=args.track_ids, shot_idx=args.shot_idx, start_idx=int(args.start_idx), end_idx=int(args.end_idx), is_static=cfg.is_static, split_cameras=args.get("split_cameras", True), ) def expand_source_paths(data_sources): return {k: get_data_source(v) for k, v in data_sources.items()} def get_data_source(source): matches = glob.glob(source) if len(matches) < 1: print(f"{source} does not exist") return source # return anyway for default values if len(matches) > 1: raise ValueError(f"{source} is not unique") return matches[0] def check_data_sources(args, cfg): print() a = time.time() if args.type == "video": preprocess_frames(args.sources.images, args.src_path, **args.frame_opts) b = time.time() preprocess_tracks(cfg.datatype, args.sources.images, args.sources.tracks, args.sources.shots, gpu=cfg.gpu) c = time.time() preprocess_cameras(args, overwrite=args.get("overwrite_cams", False)) d = time.time() print('frame, hand track, slam camera: ', b-a, c-b, d-c) class MultiPeopleDataset(Dataset): def __init__( self, data_sources: typing.Dict, seq_name, tid_spec="all", shot_idx=0, start_idx=0, end_idx=-1, is_static=False, pad_shot=False, split_cameras=True, ): self.is_static = is_static self.seq_name = seq_name self.data_sources = data_sources self.split_cameras = split_cameras # select only images in the desired shot img_files, _ = get_shot_img_files( self.data_sources["shots"], shot_idx, pad_shot ) end_idx = end_idx if end_idx > 0 else len(img_files) self.data_start, self.data_end = start_idx, end_idx img_files = img_files[start_idx:end_idx] self.img_names = [get_name(f) for f in img_files] self.num_imgs = len(self.img_names) img_dir = self.data_sources["images"] assert os.path.isdir(img_dir) self.img_paths = [os.path.join(img_dir, f) for f in img_files] img_h, img_w = imageio.imread(self.img_paths[0]).shape[:2] self.img_size = img_w, img_h print(f"USING TOTAL {self.num_imgs} {img_w}x{img_h} IMGS") # find the tracks in the video track_root = self.data_sources["tracks"] if tid_spec == "all" or tid_spec.startswith("longest"): n_tracks = MAX_NUM_TRACKS if tid_spec.startswith("longest"): n_tracks = int(tid_spec.split("-")[1]) # get the longest tracks in the selected shot track_ids = sorted(os.listdir(track_root)) track_paths = [ [f"{track_root}/{tid}/{name}_keypoints.json" for name in self.img_names] for tid in track_ids ] track_lens = [ len(list(filter(os.path.isfile, paths))) for paths in track_paths ] print("raw TRACK IDS and LENGTHS", track_ids, track_lens) track_ids = [ track_ids[i] for i in range(len(track_lens))# np.argsort(track_lens)[::-1] if track_lens[i] > MIN_TRACK_LEN ] print("TRACK LENGTHS", track_ids, track_lens) track_ids = track_ids[:n_tracks] else: track_ids = [f"{int(tid):03d}" for tid in tid_spec.split("-")] print("TRACK IDS", track_ids) self.track_ids = track_ids self.n_tracks = len(track_ids) self.track_dirs = [os.path.join(track_root, tid) for tid in track_ids] # keep a list of frame index masks of whether a track is available in a frame sidx = np.inf eidx = -1 self.track_vis_masks = [] for pred_dir in self.track_dirs: kp_paths = [f"{pred_dir}/{x}_keypoints.json" for x in self.img_names] has_kp = [os.path.isfile(x) for x in kp_paths] # keep track of which frames this track is visible in vis_mask = np.array(has_kp) idcs = np.where(vis_mask)[0] if len(idcs) > 0: si, ei = min(idcs), max(idcs) sidx = min(sidx, si) eidx = max(eidx, ei) self.track_vis_masks.append(vis_mask) eidx = max(eidx + 1, 0) sidx = min(sidx, eidx) print("START", sidx, "END", eidx) self.start_idx = sidx self.end_idx = eidx self.seq_len = eidx - sidx self.seq_intervals = [(sidx, eidx) for _ in track_ids] self.sel_img_paths = self.img_paths[sidx:eidx] self.sel_img_names = self.img_names[sidx:eidx] # used to cache data self.data_dict = {} self.cam_data = None def __len__(self): return self.n_tracks def load_data(self, interp_input=True): if len(self.data_dict) > 0: return # load camera data self.load_camera_data() # get data for each track data_out = { "mask_paths": [], "floor_plane": [], "joints2d": [], "vis_mask": [], "track_interval": [], "init_body_pose": [], "init_body_shape": [], "init_root_orient": [], "init_trans": [], "is_right": [] } # create batches of sequences # each batch is a track for a person T = self.seq_len sidx, eidx = self.start_idx, self.end_idx for i, tid in enumerate(self.track_ids): if len(self.track_ids) == 2: assert int(i) == int(tid), f'{i}/{tid}/{self.track_ids}' # load mask of visible frames for this track vis_mask = self.track_vis_masks[i][sidx:eidx] # (T) vis_idcs = np.where(vis_mask)[0] track_s, track_e = min(vis_idcs), max(vis_idcs) + 1 data_out["track_interval"].append([track_s, track_e]) vis_mask = get_ternary_mask(vis_mask) data_out["vis_mask"].append(vis_mask) # load 2d keypoints for visible frames with interpolation kp_paths = [ f"{self.track_dirs[i]}/{x}_keypoints.json" for x in self.sel_img_names ] # (T, J, 3) (x, y, conf) - with interpolation for missing frames joints2d_data = load_keypoints_with_interp(kp_paths, interp=interp_input) # print('joints2d_data: ', joints2d_data.shape) # print('joints2d_data[:, :, [2]]: ', np.unique(joints2d_data[:, :, [2]])) # raise ValueError # Discard bad ViTPose detections joints2d_data[:, :, 2] = 1.0 joints2d_data[ np.repeat(joints2d_data[:, :, [2]] < MIN_KEYP_CONF, 3, axis=2) ] = 0 # Set all confidence values to 1.0 for optimization data_out["joints2d"].append(joints2d_data) # load single image mano predictions pred_paths = [ f"{self.track_dirs[i]}/{x}_mano.json" for x in self.sel_img_names ] pose_init, orient_init, trans_init, betas_init, is_right = load_mano_preds( pred_paths, tid=tid, interp=interp_input ) n_joints = len(MANO_JOINTS) - 1 # 15 data_out["init_body_pose"].append(pose_init) data_out["init_body_shape"].append(betas_init) data_out["init_root_orient"].append(orient_init) data_out["init_trans"].append(trans_init) data_out['is_right'].append(is_right) # DEBUG: Print hand type for each track print(f"DEBUG: Track {i} (tid={tid}): is_right={is_right[0] if len(is_right) > 0 else 'empty'}, {len(pred_paths)} frames") # data_out["floor_plane"].append(DEFAULT_GROUND[:3] * DEFAULT_GROUND[3:]) self.data_dict = data_out def __getitem__(self, idx): if len(self.data_dict) < 1: self.load_data() obs_data = dict() # 2D keypoints joint2d_data = self.data_dict["joints2d"][idx] obs_data["joints2d"] = torch.Tensor(joint2d_data) # single frame predictions obs_data["init_body_pose"] = torch.Tensor(self.data_dict["init_body_pose"][idx]) obs_data["init_body_shape"] = torch.Tensor(self.data_dict["init_body_shape"][idx]) obs_data["init_root_orient"] = torch.Tensor( self.data_dict["init_root_orient"][idx] ) obs_data["init_trans"] = torch.Tensor(self.data_dict["init_trans"][idx]) obs_data["is_right"] = torch.Tensor(self.data_dict["is_right"][idx]) # floor plane # obs_data["floor_plane"] = torch.Tensor(self.data_dict["floor_plane"][idx]) # the frames the track is visible in obs_data["vis_mask"] = torch.Tensor(self.data_dict["vis_mask"][idx]) # the frames used in this subsequence obs_data["seq_interval"] = torch.Tensor(list(self.seq_intervals[idx])).to( torch.int ) # the start and end interval of available keypoints obs_data["track_interval"] = torch.Tensor( self.data_dict["track_interval"][idx] ).int() print(len(self.track_ids), self.track_ids[idx], obs_data["is_right"].shape, obs_data["seq_interval"].shape) obs_data["track_id"] = int(self.track_ids[idx]) # print(self.track_ids[idx], obs_data["is_right"]) # print(self.track_ids[idx] == obs_data["is_right"]) # is_right should be constant across time for each track assert torch.all(obs_data["is_right"] == obs_data["is_right"][0]), "is_right should be constant" assert int(self.track_ids[idx]) == int(obs_data["is_right"][0]), f"Track ID {self.track_ids[idx]} != is_right {obs_data['is_right'][0]}" obs_data["seq_name"] = self.seq_name return obs_data def load_camera_data(self): cam_dir = self.data_sources["cameras"] data_interval = 0, -1 if self.split_cameras: data_interval = self.data_start, self.data_end track_interval = self.start_idx, self.end_idx self.cam_data = CameraData( cam_dir, self.seq_len, self.img_size, self.is_static, data_interval, track_interval ) def get_camera_data(self): if self.cam_data is None: raise ValueError return self.cam_data.as_dict() class CameraData(object): def __init__( self, cam_dir, seq_len, img_size, is_static, data_interval=[0, -1], track_interval=[0, -1] ): self.img_size = img_size self.cam_dir = cam_dir self.is_static = is_static # inclusive exclusive data_start, data_end = data_interval if data_end < 0: data_end += seq_len + 1 data_len = data_end - data_start # start and end indices are with respect to the data interval sidx, eidx = track_interval if eidx < 0: eidx += data_len + 1 self.sidx, self.eidx = sidx + data_start, eidx + data_start print(self.sidx, self.eidx, data_start, data_end, data_interval, track_interval) # self.seq_len = self.eidx - self.sidx self.load_data() def load_data(self): # camera info sidx, eidx = self.sidx, self.eidx img_w, img_h = self.img_size fpath = os.path.join(self.cam_dir, "cameras.npz") if os.path.isfile(fpath) and not self.is_static: Logger.log(f"Loading cameras from {fpath}...") cam_R, cam_t, intrins, width, height = load_cameras_npz(fpath, self.seq_len) scale = img_w / width self.intrins = scale * intrins[sidx:eidx] # move first camera to origin # R0, t0 = invert_camera(cam_R[sidx], cam_t[sidx]) # self.cam_R = torch.einsum("ij,...jk->...ik", R0, cam_R[sidx:eidx]) # self.cam_t = t0 + torch.einsum("ij,...j->...i", R0, cam_t[sidx:eidx]) # t0 = -cam_t[sidx:eidx].mean(dim=0) + torch.randn(3) * 0.1 t0 = -cam_t[sidx:sidx+1] + torch.randn(3) * 0.1 self.cam_R = cam_R[sidx:eidx] self.cam_t = cam_t[sidx:eidx] - t0 else: # raise ValueError Logger.log(f"WARNING: {fpath} does not exist, using static cameras...") default_focal = 0.5 * (img_h + img_w) self.intrins = torch.tensor( [default_focal, default_focal, img_w / 2, img_h / 2] )[None].repeat(self.seq_len, 1) self.cam_R = torch.eye(3)[None].repeat(self.seq_len, 1, 1) self.cam_t = torch.zeros(self.seq_len, 3) Logger.log(f"Images have {img_w}x{img_h}, intrins {self.intrins[0]}") print("CAMERA DATA", self.cam_R.shape, self.cam_t.shape, self.intrins[0]) def world2cam(self): return self.cam_R, self.cam_t def cam2world(self): print('dataset.py, cam2world: ', self.cam_R) R = self.cam_R.transpose(-1, -2) t = -torch.einsum("bij,bj->bi", R, self.cam_t) print('dataset.py, cam2world: ', R) return R, t def as_dict(self): return { "cam_R": self.cam_R, # (T, 3, 3) "cam_t": self.cam_t, # (T, 3) "intrins": self.intrins, # (T, 4) "static": self.is_static, # bool } def get_ternary_mask(vis_mask): # get the track start and end idcs relative to the filtered interval vis_mask = torch.as_tensor(vis_mask) vis_idcs = torch.where(vis_mask)[0] track_s, track_e = min(vis_idcs), max(vis_idcs) + 1 # -1 = track out of scene, 0 = occlusion, 1 = visible vis_mask = vis_mask.float() vis_mask[:track_s] = -1 vis_mask[track_e:] = -1 return vis_mask def get_shot_img_files(shots_path, shot_idx, shot_pad=SHOT_PAD): assert os.path.isfile(shots_path) with open(shots_path, "r") as f: shots_dict = json.load(f) img_names = sorted(shots_dict.keys()) N = len(img_names) shot_mask = np.array([shots_dict[x] == shot_idx for x in img_names]) idcs = np.where(shot_mask)[0] if shot_pad > 0: # drop the frames before/after shot change raise ValueError if min(idcs) > 0: idcs = idcs[shot_pad:] if len(idcs) > 0 and max(idcs) < N - 1: idcs = idcs[:-shot_pad] if len(idcs) < MIN_SEQ_LEN: raise ValueError("shot is too short for optimization") shot_mask = np.zeros(N, dtype=bool) shot_mask[idcs] = 1 sel_paths = [img_names[i] for i in idcs] print(f"FOUND {len(idcs)}/{len(shots_dict)} FRAMES FOR SHOT {shot_idx}") return sel_paths, idcs def load_cameras_npz(camera_path, v): assert os.path.splitext(camera_path)[-1] == ".npz" cam_data = np.load(camera_path) height, width, focal = ( int(cam_data["height"]), int(cam_data["width"]), float(cam_data["focal"]), ) w2c = torch.from_numpy(cam_data["w2c"]) # (N, 4, 4) cam_R = w2c[:, :3, :3] # (N, 3, 3) cam_t = w2c[:, :3, 3] # (N, 3) N = len(w2c) if "intrins" in cam_data: intrins = torch.from_numpy(cam_data["intrins"].astype(np.float32)) else: intrins = torch.tensor([focal, focal, width / 2, height / 2])[None].repeat(N, 1) print(f"Loaded {N} cameras") return cam_R, cam_t, intrins, width, height def is_image(x): return (x.endswith(".png") or x.endswith(".jpg")) and not x.startswith(".") def get_name(x): return os.path.splitext(os.path.basename(x))[0] def split_name(x, suffix): return os.path.basename(x).split(suffix)[0] def get_names_in_dir(d, suffix): files = [split_name(x, suffix) for x in glob.glob(f"{d}/*{suffix}")] return sorted(files) def batch_join(parent, names, suffix=""): return [os.path.join(parent, f"{n}{suffix}") for n in names]