from dataclasses import dataclass, field import numpy as np import json import copy import torch import torch.nn.functional as F from skimage import measure from einops import repeat from tqdm import tqdm from PIL import Image import os import cv2 import diffusionGS from easydict import EasyDict as edict from einops import repeat, rearrange from diffusionGS.utils.misc import get_rank from diffusionGS.systems.base import BaseSystem from diffusionGS.models.diffusion import create_diffusion from diffusionGS.models.gsrenderer.gs_core import render_opencv_cam from diffusionGS.utils.typing import * from diffusionGS.utils.losses import LossComputer,MetricComputer from diffusionGS.systems.utils import * import random # DEBUG = True @diffusionGS.register("diffusion-gs-scene-system") class PointDiffusionSystem(BaseSystem): @dataclass class Config(BaseSystem.Config): num_inference_steps: int = 30 snr: bool = True #validate_metric: bool = False save_intermediate_video: bool = True save_result_for_eval: bool = False # shape vae model shape_model_type: str = None shape_model: dict = field(default_factory=dict) # noise scheduler noise_scheduler_type: str = None noise_scheduler: dict = field(default_factory=dict) # loss # denoise_scheduler_type: str = None # denoise_scheduler: dict = field(default_factory=dict) cfg: Config def configure(self): super().configure() self.shape_model = diffusionGS.find(self.cfg.shape_model_type)(self.cfg.shape_model) self.noise_scheduler = diffusionGS.find(self.cfg.noise_scheduler_type)(**self.cfg.noise_scheduler) self.diffusion_training = create_diffusion(str(1000), predict_xstart=True) self.diffusion_inference = create_diffusion(str(self.cfg.num_inference_steps), predict_xstart=True) self.loss_computer = LossComputer() def get_example_data(self): example_data_path = "examp_data/example/debug_realestate_10k/data_examples/batch_realestate_example.json" with open(example_data_path, 'r') as file: batch = json.load(file) batch['c2w'] = torch.tensor(np.stack(batch['c2w']))[:4].float().unsqueeze(0) batch['fxfycxcy'] = torch.tensor(np.stack(batch['fxfycxcy']))[:4].float().unsqueeze(0) batch['fxfycxcy'] = batch['fxfycxcy'] * 2 batch['image'] = F.interpolate(torch.tensor(np.stack(batch['image']))[:4].float(),scale_factor=2).unsqueeze(0) batch['index'] = torch.tensor(np.stack(batch['index']))[:4].float().unsqueeze(0) batch = {k: v.cuda() for k, v in batch.items() if k!="camera_json_path" and k!="view_choices"} return batch def forward(self, batch: Dict[str, Any], skip_noise=False) -> Dict[str, Any]: #get training input sel_images = batch['rgbs_input'] #(batch['rgbs_input']*2.) - 1. ##归一化 ray_o, ray_d = TransformInput(sel_images, batch['c2ws_input'], batch['fxfycxcys_input']) bs, v, c, h, w = sel_images.shape # 3. sample noise that we"ll add to the latents noise = torch.randn_like(sel_images)#.to(sel_images) # 4. Sample a random timestep for each motion timesteps = torch.randint( 0, self.cfg.noise_scheduler.num_train_timesteps, (bs,), device=sel_images.device, ) timesteps = timesteps.long() # breakpoint() sel_images[:,1:] = self.diffusion_training.q_sample(sel_images[:,1:], timesteps, noise=noise[:,1:]) #sel_images[:,1:] = self.noise_scheduler.add_noise(sel_images[:,1:],noise[:,1:],timesteps) # #breakpoint() guassians_parameters, img_aligned_xyz = self.shape_model.image_to_gaussians(sel_images, ray_o, ray_d, timesteps) ### 5. render image from gen gaussians rendered_images = self.shape_model.render_gaussians(guassians_parameters, batch['c2ws'], batch['fxfycxcys'], sel_images.shape[3],sel_images.shape[4]) #rendered_images,rendered_depths = rendered_images[:,:,:3],rendered_images[:,:,3:] # 6. compute lossc # breakpoint() l2_loss, lpips_loss, ssim_loss, pointsdist_loss, l2_loss_xyz = self.loss_computer( rendered_images, batch['rgbs'], batch['masks'], batch['masks_input'], ray_o, img_aligned_xyz=img_aligned_xyz, ) # breakpoint() return { "loss_diffusion": l2_loss.mean(), "loss_lpips": lpips_loss.mean(), "loss_ssim": ssim_loss.mean(), "loss_xyz": l2_loss_xyz.mean(), "loss_pointsdist": pointsdist_loss.mean(), "latents": sel_images, "x_t": sel_images, "noise": noise, "noise_pred": rendered_images, "timesteps": timesteps, } def training_step(self, batch, batch_idx): out = self(batch) loss = 0. for name, value in out.items(): if name.startswith("loss_"): self.log(f"train/{name}", value) loss += value * self.C(self.cfg.loss[name.replace("loss_", "lambda_")]) for name, value in self.cfg.loss.items(): if name.startswith("lambda_"): self.log(f"train_params/{name}", self.C(value)) return {"loss": loss} @torch.no_grad() def validation_step(self, batch, batch_idx): #self.eval() # TODO: write a save and denoise process # out = self(batch) # if get_rank() == 0: ######################## for debug ######################## # input_batch_test = self.get_example_data() # ray_o, ray_d = TransformInput(input_batch_test['image'], input_batch_test['c2w'], input_batch_test['fxfycxcy']) # input_batch = edict( # image = input_batch_test['image'][:,:1], # c2w = input_batch_test['c2w'], # fxfycxcy = input_batch_test['fxfycxcy'], # ray_o=ray_o, # ray_d=ray_d, # ) # input_images = input_batch_test['image'] #batch['rgbs_input'] # image_condition = input_batch_test['image'][0,:1] # b, v, c, h, w = input_images.shape # b = n # sample_noise = torch.randn(b, v-1, c, h, w, device=self.device) # 在 timestep 为 T 的时候采的 noise # input_batch["image_noisy"] = sample_noise ####################### for debug ######################## ######################## for real run ######################## ray_o, ray_d = TransformInput(batch['rgbs_input'], batch['c2ws_input'], batch['fxfycxcys_input']) input_batch = edict( image = batch['rgbs_input'][:,:1], c2w = batch['c2ws_input'], fxfycxcy = batch['fxfycxcys_input'], ray_o=ray_o, ray_d=ray_d, ) input_images = batch['rgbs_input'] #[:,:1] b, v, c, h, w = input_images.shape # b = n sample_noise = torch.randn(b, v-1, c, h, w, device=self.device) # 在 timestep 为 T 的时候采的 noise input_batch["image_noisy"] = sample_noise ######################## for real run ######################## image_condition = input_images[:, 0, ...].unsqueeze(1) traj_timesteps = list(self.diffusion_inference.timestep_map)[::-1] # 提前设计好的 timestep 路线 traj_len = len(traj_timesteps) traj_samples = [] traj_pred_xstart = [] final_out = None for out in self.diffusion_inference.p_sample_loop_progressive( self.shape_model, sample_noise.shape, input_batch, clip_denoised=False, progress=True, device=self.device, ): samples = out["sample"] pred_xstart = out["pred_xstart"] # 把 image condition 和 samples concate 起来 samples_present = torch.cat((image_condition, samples), dim=1) pred_xstart_present = torch.cat((image_condition, pred_xstart), dim=1) traj_samples.append(samples_present) traj_pred_xstart.append(pred_xstart_present) final_out = out traj_samples = torch.stack(traj_samples, 0) # [T, B, V, C, H, W] [250, 8, 4, 3, 256, 256] traj_pred_xstart = torch.stack(traj_pred_xstart, 0) # [T, B, V, C, H, W] [250, 8, 4, 3, 256, 256] traj_samples = rearrange( traj_samples, "t n v c h w -> n t h (v w) c" ).contiguous().cpu().numpy() traj_pred_xstart = rearrange( traj_pred_xstart, "t n v c h w -> n t h (v w) c" ).contiguous().cpu().numpy() # save multiple videos if self.cfg.save_intermediate_video: for i in range(traj_samples.shape[0]): vis = traj_samples[i] vis = display_timestep_on_video(vis, traj_timesteps) self.save_videos(f"it{self.true_global_step}/{batch['uid'][i]}_{batch['sel_idx'][i] if 'sel_idx' in batch.keys() else 0}_{i:04d}_traj_xt.mp4", vis, fps=24, quality=8) # 将输出的图片保存为视频 image_grid = [] for i in range(traj_pred_xstart.shape[0]): vis = traj_pred_xstart[i] vis = display_timestep_on_video(vis, traj_timesteps) self.save_videos(f"it{self.true_global_step}/{batch['uid'][i]}_{batch['sel_idx'][i] if 'sel_idx' in batch.keys() else 0}_{i:04d}_traj_xstart.mp4", vis, fps=24, quality=8) # breakpoint() for i in range(traj_pred_xstart.shape[0]): self.save_guassians_ply_scene(f"it{self.true_global_step}/{batch['uid'][i]}_{batch['sel_idx'][i] if 'sel_idx' in batch.keys() else 0}.ply", \ final_out['denoiser_output_dict']['pred_gaussians'][i], render_keyframe_c2ws = batch['c2ws_input'][i], \ render_intrinsics = batch['fxfycxcys_input'][i], render_video=True, h=batch['rgbs_input'].shape[-2], w=batch['rgbs_input'].shape[-1]) self.save_torch_images(f"it{self.true_global_step}/{batch['uid'][i]}_{batch['sel_idx'][i] if 'sel_idx' in batch.keys() else 0}.png", torch.cat([image_condition[i],final_out['denoiser_output_dict']['render_images'][i]], dim=0)) if self.cfg.save_result_for_eval: for i in range(len(batch['uid'])): # Score every specified target view, excluding the conditioning image. targets = batch['rgbs'][i, 1:] cameras = batch['c2ws'][i, 1:] intrinsics = batch['fxfycxcys'][i, 1:] gaussians = final_out['denoiser_output_dict']['pred_gaussians'][i] # The native rasterizer requires float32 camera matrices, even # while Lightning runs the denoiser with mixed precision. with torch.autocast(device_type="cuda", enabled=False): predictions = torch.stack([ render_opencv_cam(gaussians, targets.shape[-2], targets.shape[-1], camera.float(), intrinsic.float())['render'] for camera, intrinsic in zip(cameras, intrinsics) ]) result_pkg = { 'render_images': predictions.detach().cpu(), 'image': targets.detach().cpu(), 'target_indices': batch['image_indices'][i, 1:].detach().cpu(), } self.save_torch(f"it{self.true_global_step}/{batch['uid'][i]}.pt", result_pkg) torch.cuda.empty_cache() # Metrics are computed from the saved target views by eval_scene_result.py. return {} def on_validation_epoch_end(self): pass