from diffsynth.core.data.operators import * import torch, json, pandas, sys, os import numpy as np from pathlib import Path from PIL import Image import cv2 sys.path.append(os.getcwd()) sys.path.append("./DepthAnything3/src") from DepthAnything3.src.depth_anything_3.api import DepthAnything3 from openexr_numpy import imread, imwrite import torch.nn.functional as F import pandas as pd class MetaViewUnifiedDataset(torch.utils.data.Dataset): def __init__( self, base_path=None, metadata_path=None, repeat=1, data_file_keys=tuple(), main_data_operator=lambda x: x, special_operator_map=None, prope=False, debug=False, mode="train", norm_scale=1.0, path_3D=None, export_3D_feat_layers=None, anno_src=None, add_depth=False, subset=None, base_model = "qwen", ): self.base_model = base_model self.base_path = base_path self.metadata_path = metadata_path self.repeat = repeat self.data_file_keys = data_file_keys self.main_data_operator = main_data_operator self.cached_data_operator = LoadTorchPickle() self.special_operator_map = {} if special_operator_map is None else special_operator_map self.data = [] self.cached_data = [] paths = base_path.split(";") if subset is None: self.videos = os.listdir(base_path) elif len(paths) > 1: self.videos = [] for path in paths: dirs = os.listdir(path) for sub in subset: if sub in dirs: p = os.path.join(path, sub) folders = os.listdir(p) for folder in folders: self.videos.append(os.path.join(p, folder)) else: self.videos = [] for sub in subset: p = os.path.join(base_path, sub) dirs = os.listdir(p) for d in dirs: self.videos.append(os.path.join(sub, d)) self.total_length = len(self.videos) self.prope = prope self.mode = mode self.norm_scale = norm_scale self.anno_src = anno_src self.add_depth = add_depth self.subset = subset if prope: self.load_from_cache = False self.model_3D = None if path_3D is not None: device = torch.device("cuda") self.model_3D = DepthAnything3.from_pretrained(path_3D) self.model_3D = self.model_3D.to(device=device) export_3D_feat_layers = export_3D_feat_layers.split(",") self.export_3D_feat_layers = [int(s) for s in export_3D_feat_layers] if debug: self.max_dist() exit(0) @staticmethod def default_image_operator( base_path="", max_pixels=1920*1080, height=None, width=None, height_division_factor=16, width_division_factor=16, ): return RouteByType(operator_map=[ (str, ToAbsolutePath(base_path) >> LoadImage() >> ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor)), (list, SequencialProcess(ToAbsolutePath(base_path) >> LoadImage() >> ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor))), ]) @staticmethod def default_video_operator( base_path="", max_pixels=1920*1080, height=None, width=None, height_division_factor=16, width_division_factor=16, num_frames=81, time_division_factor=4, time_division_remainder=1, ): return RouteByType(operator_map=[ (str, ToAbsolutePath(base_path) >> RouteByExtensionName(operator_map=[ (("jpg", "jpeg", "png", "webp"), LoadImage() >> ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor) >> ToList()), (("gif",), LoadGIF( num_frames, time_division_factor, time_division_remainder, frame_processor=ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor), )), (("mp4", "avi", "mov", "wmv", "mkv", "flv", "webm"), LoadVideo( num_frames, time_division_factor, time_division_remainder, frame_processor=ImageCropAndResize(height, width, max_pixels, height_division_factor, width_division_factor), )), ])), ]) def __getitem__(self, data_id): if self.prope: if len(self.base_path.split(";")) > 1: return self.getitem_prope_all(data_id) elif "DL3DV" in self.base_path and self.metadata_path is None: return self.getitem_prope_DL3DV(data_id) elif "DL3DV" in self.base_path and self.metadata_path: return self.getitem_metadata_DL3DV(data_id) elif self.metadata_path: return self.getitem_metadata(data_id) return data def quick_check(self, extrinsics: torch.Tensor) -> bool: """check extrinsics""" if extrinsics.shape[-2:] != (4, 4): print("Extrinsics wrong shape!") return False t_norm = torch.norm(extrinsics[:, :3, 3], dim=1) if t_norm.max() > 100: # 阈值根据应用设定 print(f"Extrinsics shift too large! {t_norm.max()}") return False R = extrinsics[..., :3, :3] # 检查行列式接近1 det_R = torch.det(R) if not torch.allclose(det_R, torch.ones_like(det_R), atol=1e-4): print("Extrinsics not ortho!") return False # 检查最后一行 last_row = extrinsics[..., 3, :] expected = torch.tensor([0.0, 0.0, 0.0, 1.0], device=extrinsics.device) if not torch.allclose(last_row, expected.expand_as(last_row), atol=1e-4): print("Extrinsics wrong row!") return False return True def getitem_prope_DL3DV(self, index): video = self.videos[index % self.total_length] if not "Evaluation" in self.base_path: while not os.path.isdir(os.path.join(self.base_path, video)) or not os.path.exists(os.path.join(self.base_path, video, "transforms.json")): index += 1 video = self.videos[index % self.total_length] blender2opencv = np.array( [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]] ) data = {} # if "Evaluation" in self.base_path: # video = f"{video}/{video}/nerfstudio" with open(os.path.join(self.base_path, video, "transforms.json"), 'r', encoding='utf-8') as file: json_str = file.read() meta = json.loads(json_str) # frames = meta["frames"] # frames = sorted(frames, key=lambda x: x["colmap_im_id"]) frames = sorted(os.listdir(os.path.join(self.base_path, video, "images_4"))) import random interval = min(40, len(frames)) extrinsics_check = False edit_idx = None while not extrinsics_check: extrinsics_check = True edit_idx = random.randint(0, len(frames) - interval) if "val" in self.mode: edit_idx = 10 # edit_image = os.path.join(video, frames[edit_idx]["file_path"].replace("images", "images_4")) edit_image = os.path.join(video, "images_4", frames[edit_idx]) # print(edit_idx, frames[edit_idx]) # if "Evaluation" in self.base_path and self.anno_src is not None and "vipe" in self.anno_src: # vipe_pose = np.load(os.path.join(self.base_path, video.split('/')[0], "pose/video.npz")) # edit_viewmats = vipe_pose["data"][edit_idx] if self.anno_src is not None and "vipe-DA3" in self.anno_src: vipe_pose = np.load(os.path.join(self.base_path, video, "vipe-DA3/pose/video.npz")) edit_viewmats = vipe_pose["data"][edit_idx] elif self.anno_src is not None and self.anno_src == "vipe": vipe_pose = np.load(os.path.join(self.base_path, video, "pose/video.npz")) edit_viewmats = vipe_pose["data"][edit_idx] else: edit_viewmats = np.array(frames[edit_idx]["transform_matrix"], dtype=np.float32) @ blender2opencv # c2w! and invert y z axis! edit_viewmats = torch.Tensor(edit_viewmats).unsqueeze(0) max_idx = min(40, len(frames) - edit_idx) target_idx = random.randint(edit_idx + max_idx // 2, edit_idx + max_idx - 1) if "val" in self.mode: target_idx = 30 # target_image = os.path.join(video, frames[target_idx]["file_path"].replace("images", "images_4")) target_image = os.path.join(video, "images_4", frames[target_idx]) # if "Evaluation" in self.base_path and self.anno_src is not None and "vipe" in self.anno_src: # vipe_pose = np.load(os.path.join(self.base_path, video.split('/')[0], "pose/video.npz")) # target_viewmats = vipe_pose["data"][target_idx] if self.anno_src is not None and "vipe-DA3" in self.anno_src: vipe_pose = np.load(os.path.join(self.base_path, video, "vipe-DA3/pose/video.npz")) target_viewmats = vipe_pose["data"][target_idx] elif self.anno_src is not None and self.anno_src == "vipe": vipe_pose = np.load(os.path.join(self.base_path, video, "pose/video.npz")) target_viewmats = vipe_pose["data"][target_idx] else: target_viewmats = np.array(frames[target_idx]["transform_matrix"], dtype=np.float32) @ blender2opencv target_viewmats = torch.Tensor(target_viewmats).unsqueeze(0) edit_c2w = edit_viewmats target_c2w = target_viewmats in_c2ws = torch.cat([target_c2w, edit_c2w], dim=0) # normalize c2ws = torch.einsum("ij,njk->nik", torch.linalg.inv(edit_c2w[0]), in_c2ws) # shift to src coord(edit_image) c2ws[:, :3, 3] /= self.norm_scale #20.0 # 10.0 # translation normalized # print(c2ws) if not self.quick_check(c2ws[0:1, :, :]): extrinsics_check = False #transform c2w to w2c align with PRoPE implementation viewmats = torch.linalg.inv(c2ws) s = 0 ks = [ [meta["fl_x"], s, meta["cx"]], [ 0, meta["fl_y"], meta["cy"]], [ 0, 0, 1]] ks = torch.Tensor(ks).unsqueeze(0) image_height = meta["h"] image_width = meta["w"] ks[..., 0, 0] = ks[..., 0, 0] / image_width ks[..., 1, 1] = ks[..., 1, 1] / image_height ks[..., 0, 2] = ks[..., 0, 2] / image_width - 0.5 ks[..., 1, 2] = ks[..., 1, 2] / image_height - 0.5 ks[..., 2, 2] = 1.0 # ks has been normalized!! data["edit_image"] = self.main_data_operator(edit_image).resize((960, 528)) data["image"] = self.main_data_operator(target_image).resize((960, 528)) data["viewmats"] = viewmats data["Ks"] = torch.cat([ks, ks], dim=0) # print("original size : ",data["image"].size, data["edit_image"].size) if torch.isnan(data["viewmats"]).any() or torch.isnan(data["Ks"]).any(): print("!!!camera param has NaN!!!") print(target_image, edit_image) exit(0) if "qwen" in self.base_model: data["prompt"] = "镜头视角转到指定位置" elif "flux" in self.base_model: data["prompt"] = "Turn to the target view" data["name"] = video # if data["edit_image"].size[0] != 960: #7103edc158a862dbfa3c3454e4de584dad59c3c30055919f1dfa7fd7acfdd5c9 # print(f"{video} has different size!!") if self.model_3D is not None and "val" not in self.mode: feat_3D = self.model_3D.inference( [data["edit_image"].resize((960, 528))], # (1, 33, 60) export_feat_layers=self.export_3D_feat_layers, # (1, 20, 36, 1536) H, W, C (1, 20, 36, 1024) 960 528 process_res=840, ) feats = [] for layer in self.export_3D_feat_layers: feats.append(torch.from_numpy(feat_3D.aux[f"feat_layer_{layer}"])) data["feat_3D"] = torch.cat(feats, dim=-1)[0] # (20, 36, 1536) H, W, C if torch.isnan(data["feat_3D"]).any(): print("!!!feat 3D has NaN!!!") print(video, edit_idx) exit(0) # target_z = imread(os.path.join(self.base_path, video, f"vipe-DA3/depth/{(target_idx):05d}.exr"), "Z") # target_z[np.isnan(target_z)] = 1000 # target_z[(target_z > 1000) | np.isinf(target_z)] = 1000 # target_depth = torch.Tensor(target_z) # data["target_depth"] = target_depth if self.add_depth: if "vipe-DA3" in self.anno_src: z_channel = imread(os.path.join(self.base_path, video, f"vipe-DA3/depth/{(edit_idx):05d}.exr"), "Z") elif "vipe" == self.anno_src: z_channel = imread(os.path.join(self.base_path, video, f"depth/{(edit_idx):05d}.exr"), "Z") z_channel[np.isnan(z_channel)] = 0 z_channel[(z_channel > 1000) | np.isinf(z_channel)] = 1000 depth_edit = torch.Tensor(z_channel).unsqueeze(0).unsqueeze(0) depth_edit = F.interpolate(depth_edit, size=(528, 960), mode='bilinear', align_corners=False)[0] # print(torch.max(depth_edit), torch.min(depth_edit)) depth_latent = torch.zeros_like(depth_edit) depth = torch.cat([depth_latent, depth_edit], dim=0) # n, h, w if torch.isnan(depth).any(): print("!!!depth has NaN!!!") print(video, edit_idx) exit(0) # src_depth = np.array(depth_edit[0]) # mx = np.max(src_depth) # mn = np.min(src_depth) # K = [ [meta["fl_x"] / image_width * 960, 0, 960 // 2], # [ 0, meta["fl_y"] / image_height * 528, 528 // 2], # [ 0, 0, 1]] # T = np.array(viewmats[0]) # tgt_depth = self.transform_depth( # src_depth, # K=np.array(K), # T=T # ) # depth_latent = torch.Tensor(tgt_depth).unsqueeze(0) # depth = torch.cat([depth_latent, depth_edit], dim=0) # n, h, w data["depth"] = depth #if "val" in self.mode: # tgt_depth = (tgt_depth / np.max(tgt_depth) * 255).astype(np.uint8) # src_depth = (src_depth / np.max(src_depth) * 255).astype(np.uint8) # z_channel = imread(os.path.join(self.base_path, video.split('/')[0], f"depth/{target_idx:05d}.exr"), "Z") # depth_gt = torch.Tensor(z_channel).unsqueeze(0).unsqueeze(0) # depth_gt = F.interpolate(depth_gt, size=(528, 960), mode='bilinear', align_corners=False)[0] # depth_gt = np.array(depth_gt[0]) # depth_gt = (depth_gt / np.max(depth_gt) * 255).astype(np.uint8) # depth_vis = np.concatenate((src_depth, tgt_depth, depth_gt), axis=1) # im = Image.fromarray(depth_vis) # im.save(f"depth_vis/{index}_{mx:.2f}_{mn:.2f}.png") return data def getitem_metadata_DL3DV(self, index): csv = pd.read_csv(self.metadata_path) row = csv.iloc[index] video = str(row['video']) edit_idx = int(row['edit_idx']) target_idx = int(row['target_idx']) data = {} with open(os.path.join(self.base_path, video, "transforms.json"), 'r', encoding='utf-8') as file: json_str = file.read() meta = json.loads(json_str) frames = sorted(os.listdir(os.path.join(self.base_path, video, "images_4"))) import random edit_image = os.path.join(video, "images_4", frames[edit_idx]) if self.anno_src is not None and "vipe-DA3" in self.anno_src: vipe_pose = np.load(os.path.join(self.base_path, video, "vipe-DA3/pose/video.npz")) edit_viewmats = vipe_pose["data"][edit_idx] elif self.anno_src is not None and self.anno_src == "vipe": vipe_pose = np.load(os.path.join(self.base_path, video, "pose/video.npz")) edit_viewmats = vipe_pose["data"][edit_idx] else: edit_viewmats = np.array(frames[edit_idx]["transform_matrix"], dtype=np.float32) @ blender2opencv # c2w! and invert y z axis! edit_viewmats = torch.Tensor(edit_viewmats).unsqueeze(0) target_image = os.path.join(video, "images_4", frames[target_idx]) if self.anno_src is not None and "vipe-DA3" in self.anno_src: vipe_pose = np.load(os.path.join(self.base_path, video, "vipe-DA3/pose/video.npz")) target_viewmats = vipe_pose["data"][target_idx] elif self.anno_src is not None and self.anno_src == "vipe": vipe_pose = np.load(os.path.join(self.base_path, video, "pose/video.npz")) target_viewmats = vipe_pose["data"][target_idx] else: target_viewmats = np.array(frames[target_idx]["transform_matrix"], dtype=np.float32) @ blender2opencv target_viewmats = torch.Tensor(target_viewmats).unsqueeze(0) edit_c2w = edit_viewmats target_c2w = target_viewmats in_c2ws = torch.cat([target_c2w, edit_c2w], dim=0) # normalize c2ws = torch.einsum("ij,njk->nik", torch.linalg.inv(edit_c2w[0]), in_c2ws) # shift to src coord(edit_image) c2ws[:, :3, 3] /= self.norm_scale #20.0 # 10.0 # translation normalized #transform c2w to w2c align with PRoPE implementation viewmats = torch.linalg.inv(c2ws) s = 0 ks = [ [meta["fl_x"], s, meta["cx"]], [ 0, meta["fl_y"], meta["cy"]], [ 0, 0, 1]] ks = torch.Tensor(ks).unsqueeze(0) image_height = meta["h"] image_width = meta["w"] ks[..., 0, 0] = ks[..., 0, 0] / image_width ks[..., 1, 1] = ks[..., 1, 1] / image_height ks[..., 0, 2] = ks[..., 0, 2] / image_width - 0.5 ks[..., 1, 2] = ks[..., 1, 2] / image_height - 0.5 ks[..., 2, 2] = 1.0 # ks has been normalized!! data["edit_image"] = self.main_data_operator(edit_image).resize((960, 528)) data["image"] = self.main_data_operator(target_image).resize((960, 528)) data["viewmats"] = viewmats data["Ks"] = torch.cat([ks, ks], dim=0) # print("original size : ",data["image"].size, data["edit_image"].size) if torch.isnan(data["viewmats"]).any() or torch.isnan(data["Ks"]).any(): print("!!!camera param has NaN!!!") print(target_image, edit_image) exit(0) if "qwen" in self.base_model: data["prompt"] = "镜头视角转到指定位置" elif "flux" in self.base_model: data["prompt"] = "Turn to the target view" data["name"] = video if self.model_3D is not None and "val" not in self.mode: feat_3D = self.model_3D.inference( [data["edit_image"].resize((960, 528))], # (1, 33, 60) export_feat_layers=self.export_3D_feat_layers, # (1, 20, 36, 1536) H, W, C (1, 20, 36, 1024) 960 528 process_res=840, ) feats = [] for layer in self.export_3D_feat_layers: feats.append(torch.from_numpy(feat_3D.aux[f"feat_layer_{layer}"])) data["feat_3D"] = torch.cat(feats, dim=-1) # (1, 20, 36, 1536) B, H, W, C if torch.isnan(data["feat_3D"]).any(): print("!!!feat 3D has NaN!!!") print(video, edit_idx) exit(0) target_z = imread(os.path.join(self.base_path, video, f"vipe-DA3/depth/{(target_idx):05d}.exr"), "Z") target_z[np.isnan(target_z)] = 1000 target_z[(target_z > 1000) | np.isinf(target_z)] = 1000 target_depth = torch.Tensor(target_z) data["target_depth"] = target_depth if self.add_depth: if "vipe-DA3" in self.anno_src: z_channel = imread(os.path.join(self.base_path, video, f"vipe-DA3/depth/{(edit_idx):05d}.exr"), "Z") elif "vipe" == self.anno_src: z_channel = imread(os.path.join(self.base_path, video, f"depth/{(edit_idx):05d}.exr"), "Z") z_channel[np.isnan(z_channel)] = 0 z_channel[(z_channel > 1000) | np.isinf(z_channel)] = 1000 depth_edit = torch.Tensor(z_channel).unsqueeze(0).unsqueeze(0) depth_edit = F.interpolate(depth_edit, size=(528, 960), mode='bilinear', align_corners=False)[0] # print(torch.max(depth_edit), torch.min(depth_edit)) depth_latent = torch.zeros_like(depth_edit) depth = torch.cat([depth_latent, depth_edit], dim=0) # n, h, w if torch.isnan(depth).any(): print("!!!depth has NaN!!!") print(video, edit_idx) exit(0) data["depth"] = depth return data def getitem_metadata(self, index): csv = pd.read_csv(self.metadata_path) row = csv.iloc[index] video = str(row['video']) edit_idx = int(row['edit_idx']) target_idx = int(row['target_idx']) vipe_pose = np.load(os.path.join(self.base_path, video, "vipe-DA3/pose/video.npz")) vipe_intr = np.load(os.path.join(self.base_path, video, "vipe-DA3/intrinsics/video.npz")) data = {} cap = cv2.VideoCapture(os.path.join(self.base_path, video, "video.mp4")) len_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - 12 # Re10K shot change import random interval = min(40, len_frames - 1) cap.set(cv2.CAP_PROP_POS_FRAMES, edit_idx) ret, f = cap.read() edit_image = Image.fromarray(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) edit_viewmats = vipe_pose["data"][edit_idx] edit_viewmats = torch.Tensor(edit_viewmats).unsqueeze(0) cap.set(cv2.CAP_PROP_POS_FRAMES, target_idx) ret, f = cap.read() target_image = Image.fromarray(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) target_viewmats = vipe_pose["data"][target_idx] target_viewmats = torch.Tensor(target_viewmats).unsqueeze(0) edit_c2w = edit_viewmats target_c2w = target_viewmats print(edit_idx, target_idx, edit_c2w.shape, target_c2w.shape) in_c2ws = torch.cat([target_c2w, edit_c2w], dim=0) # normalize c2ws = torch.einsum("ij,njk->nik", torch.linalg.inv(edit_c2w[0]), in_c2ws) # shift to src coord(edit_image) c2ws[:, :3, 3] /= self.norm_scale # translation normalized cap.release() # release video #transform c2w to w2c align with PRoPE implementation viewmats = torch.linalg.inv(c2ws) intri = vipe_intr["data"][edit_idx] s = 0 ks = [ [intri[0], s, intri[2]], [ 0, intri[1], intri[3]], [ 0, 0, 1.]] image_width = torch.tensor(intri[2]) * 2 image_height = torch.tensor(intri[3]) * 2 ks = torch.Tensor(ks).unsqueeze(0) ks[..., 0, 0] = ks[..., 0, 0] / image_width ks[..., 1, 1] = ks[..., 1, 1] / image_height ks[..., 0, 2] = ks[..., 0, 2] / image_width - 0.5 ks[..., 1, 2] = ks[..., 1, 2] / image_height - 0.5 data["edit_image"] = edit_image.resize((960, 528)) data["image"] = target_image.resize((960, 528)) data["viewmats"] = viewmats data["Ks"] = torch.cat([ks, ks], dim=0) # print(data["Ks"]) if torch.isnan(data["viewmats"]).any() or torch.isnan(data["Ks"]).any(): print("!!!camera param has NaN!!!") print(video, edit_idx, target_idx) exit(0) if "qwen" in self.base_model: data["prompt"] = "镜头视角转到指定位置" elif "flux" in self.base_model: data["prompt"] = "Turn to the target view" data["name"] = video if self.model_3D is not None and "val" not in self.mode: feat_3D = self.model_3D.inference( [data["edit_image"].resize((960, 528))], # (1, 33, 60) export_feat_layers=self.export_3D_feat_layers, # (1, 20, 36, 1536) H, W, C (1, 20, 36, 1024) 960 528 process_res=840, ) feats = [] for layer in self.export_3D_feat_layers: feats.append(torch.from_numpy(feat_3D.aux[f"feat_layer_{layer}"])) data["feat_3D"] = torch.cat(feats, dim=-1)[0] # (20, 36, 1536) H, W, C if torch.isnan(data["feat_3D"]).any(): print("!!!feat 3D has NaN!!!") print(video, edit_idx) exit(0) if self.add_depth: z_channel = imread(os.path.join(self.base_path, video, f"vipe-DA3/depth/{(edit_idx):05d}.exr"), "Z") z_channel[np.isnan(z_channel)] = 0 z_channel[(z_channel > 1000) | np.isinf(z_channel)] = 1000 depth_edit = torch.Tensor(z_channel).unsqueeze(0).unsqueeze(0) depth_edit = F.interpolate(depth_edit, size=(528, 960), mode='bilinear', align_corners=False)[0] # print(torch.max(depth_edit), torch.min(depth_edit)) depth_latent = torch.zeros_like(depth_edit) depth = torch.cat([depth_latent, depth_edit], dim=0) # n, h, w if torch.isnan(depth).any(): print("!!!depth has NaN!!!") print(video, edit_idx) exit(0) data["depth"] = depth return data def get_frame(self, video_path, frame_num): cap = cv2.VideoCapture(video_path) cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num) ret, frame = cap.read() cap.release() if ret: return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) else: print(f'error video {video_path} {frame_num}') exit(0) def getitem_prope_all(self, index): video = self.videos[index % self.total_length] if not "Evaluation" in self.base_path: while not os.path.isdir(video): index += 1 video = self.videos[index % self.total_length] vipe_pose = np.load(os.path.join(video, "vipe-DA3/pose/video.npz")) vipe_intr = np.load(os.path.join(video, "vipe-DA3/intrinsics/video.npz")) data = {} if 'DL3DV' in video: frames = sorted(os.listdir(os.path.join(video, "images_4"))) len_frames = len(frames) else: cap = cv2.VideoCapture(os.path.join(video, "video.mp4")) len_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - 12 # Re10K shot change import random interval = min(40, len_frames - 1) extrinsics_check = False edit_idx = None while not extrinsics_check: if "val" in self.mode: edit_idx = 10 target_idx = 30 else: edit_idx = random.randint(0, len_frames - interval) max_idx = min(40, len_frames - edit_idx) # 0 -12 target_idx = random.randint(edit_idx + max_idx // 2, edit_idx + max_idx - 1) # print(len_frames, edit_idx, target_idx) if 'DL3DV' in video: edit_image = os.path.join(video, "images_4", frames[edit_idx]) edit_image = Image.open(edit_image) else: cap.set(cv2.CAP_PROP_POS_FRAMES, edit_idx) ret, f = cap.read() edit_image = Image.fromarray(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) edit_viewmats = vipe_pose["data"][edit_idx] edit_viewmats = torch.Tensor(edit_viewmats).unsqueeze(0) if 'DL3DV' in video: target_image = os.path.join(video, "images_4", frames[target_idx]) target_image = Image.open(target_image) else: cap.set(cv2.CAP_PROP_POS_FRAMES, target_idx) ret, f = cap.read() target_image = Image.fromarray(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) target_viewmats = vipe_pose["data"][target_idx] target_viewmats = torch.Tensor(target_viewmats).unsqueeze(0) edit_c2w = edit_viewmats target_c2w = target_viewmats in_c2ws = torch.cat([target_c2w, edit_c2w], dim=0) # normalize c2ws = torch.einsum("ij,njk->nik", torch.linalg.inv(edit_c2w[0]), in_c2ws) # shift to src coord(edit_image) c2ws[:, :3, 3] /= self.norm_scale # translation normalized if self.quick_check(c2ws[0:1, :, :]): extrinsics_check = True if 'DL3DV' not in video: cap.release() # release video break #transform c2w to w2c align with PRoPE implementation viewmats = torch.linalg.inv(c2ws) intri = vipe_intr["data"][edit_idx] s = 0 ks = [ [intri[0], s, intri[2]], [ 0, intri[1], intri[3]], [ 0, 0, 1.]] image_width = torch.tensor(intri[2]) * 2 image_height = torch.tensor(intri[3]) * 2 ks = torch.Tensor(ks).unsqueeze(0) ks[..., 0, 0] = ks[..., 0, 0] / image_width ks[..., 1, 1] = ks[..., 1, 1] / image_height ks[..., 0, 2] = ks[..., 0, 2] / image_width - 0.5 ks[..., 1, 2] = ks[..., 1, 2] / image_height - 0.5 # ks = [ [meta["fl_x"], s, meta["cx"]], # [ 0, meta["fl_y"], meta["cy"]], # [ 0, 0, 1]] # ks = torch.Tensor(ks).unsqueeze(0) # image_height = meta["h"] # image_width = meta["w"] # ks[..., 0, 0] = ks[..., 0, 0] / image_width # ks[..., 1, 1] = ks[..., 1, 1] / image_height # ks[..., 0, 2] = ks[..., 0, 2] / image_width - 0.5 # ks[..., 1, 2] = ks[..., 1, 2] / image_height - 0.5 # ks[..., 2, 2] = 1.0 # ks has been normalized!! data["edit_image"] = edit_image.resize((960, 528)) data["image"] = target_image.resize((960, 528)) data["viewmats"] = viewmats data["Ks"] = torch.cat([ks, ks], dim=0) # print(data["Ks"]) if torch.isnan(data["viewmats"]).any() or torch.isnan(data["Ks"]).any(): print("!!!camera param has NaN!!!") print(video, edit_idx, target_idx) exit(0) if "qwen" in self.base_model: data["prompt"] = "镜头视角转到指定位置" elif "flux" in self.base_model: data["prompt"] = "Turn to the target view" data["name"] = video if self.model_3D is not None and "val" not in self.mode: feat_3D = self.model_3D.inference( [data["edit_image"].resize((960, 528))], # (1, 33, 60) export_feat_layers=self.export_3D_feat_layers, # (1, 20, 36, 1536) H, W, C (1, 20, 36, 1024) 960 528 process_res=840, ) feats = [] for layer in self.export_3D_feat_layers: feats.append(torch.from_numpy(feat_3D.aux[f"feat_layer_{layer}"])) data["feat_3D"] = torch.cat(feats, dim=-1)[0] # (20, 36, 1536) H, W, C if torch.isnan(data["feat_3D"]).any(): print("!!!feat 3D has NaN!!!") print(video, edit_idx) exit(0) # target_z = imread(os.path.join(self.base_path, video, f"vipe-DA3/depth/{(target_idx):05d}.exr"), "Z") # target_z[np.isnan(target_z)] = 1000 # target_z[(target_z > 1000) | np.isinf(target_z)] = 1000 # target_depth = torch.Tensor(target_z) # data["target_depth"] = target_depth if self.add_depth: z_channel = imread(os.path.join(video, f"vipe-DA3/depth/{(edit_idx):05d}.exr"), "Z") z_channel[np.isnan(z_channel)] = 0 z_channel[(z_channel > 1000) | np.isinf(z_channel)] = 1000 depth_edit = torch.Tensor(z_channel).unsqueeze(0).unsqueeze(0) depth_edit = F.interpolate(depth_edit, size=(528, 960), mode='bilinear', align_corners=False)[0] # print(torch.max(depth_edit), torch.min(depth_edit)) depth_latent = torch.zeros_like(depth_edit) depth = torch.cat([depth_latent, depth_edit], dim=0) # n, h, w if torch.isnan(depth).any(): print("!!!depth has NaN!!!") print(video, edit_idx) exit(0) data["depth"] = depth return data def __len__(self): if self.prope: return self.total_length * self.repeat if __name__ == '__main__': metadata_path = '../RealEstate10K/meta_view/hard.csv' base_path = "../RealEstate10K/JiaHWang/Re10K/test" # 48aaed5a44005bccd51d529ab90335b144fe5e7f3c8a22ba399f4ee3b3fb6728 dataset = MetaViewUnifiedDataset( base_path=base_path, metadata_path=metadata_path, # subset=['1K', '2K', '3K', '4K', '5K', '6K', '7K', '9K', '10K', '11K', 'train', 'Sekai-Real-Walking-HQ-split'], repeat=1, data_file_keys="image,edit_image".split(","), prope=True, debug=False, mode="train", norm_scale=1.0, # path_3D="../Depth-Anything-3/model/DA3-GIANT-1.1", # export_3D_feat_layers="19,39", anno_src="vipe-DA3", add_depth=True, main_data_operator=UnifiedDataset.default_image_operator( base_path=base_path, max_pixels=1048576, height=None, width=None, height_division_factor=16, width_division_factor=16, ) ) dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, collate_fn=lambda x: x[0], num_workers=1) print(len(dataset)) cnt = 0 for data in dataloader: # print(data["name"]) cnt += 1 print(data['name']) #if cnt > 11: # exit(0) print(cnt)