File size: 6,083 Bytes
9205b56 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | #
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use
# under the terms of the LICENSE.md file.
#
# For inquiries contact george.drettakis@inria.fr
#
import numpy as np
import torch
from scene import Scene
import os
from tqdm import tqdm
from os import makedirs
from gaussian_renderer import render
import torchvision
from utils.general_utils import safe_state
from argparse import ArgumentParser
from arguments import ModelParams, PipelineParams, get_combined_args
from gaussian_renderer import GaussianModel
from autoencoder.model import Autoencoder
from eval.openclip_encoder import OpenCLIPNetwork
import cv2
def render_set(model_path, source_path, name, iteration, views, gaussians, pipeline, background, args, ae_model, clip_model):
render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders")
gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt")
render_npy_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders_npy")
gts_npy_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt_npy")
makedirs(render_npy_path, exist_ok=True)
makedirs(gts_npy_path, exist_ok=True)
makedirs(render_path, exist_ok=True)
makedirs(gts_path, exist_ok=True)
for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
output = render(view, gaussians, pipeline, background, args, ae_model=ae_model, clip_model=clip_model)
if idx == 0: continue
if not args.include_feature:
rendering = output["render"]
else:
rendering = output["language_feature_image"]
if not args.include_feature:
gt = view.original_image[0:3, :, :]
else:
gt, mask = view.get_language_feature(os.path.join(source_path, args.language_features_name), feature_level=args.feature_level)
# if clip_model is not None and ae_model is not None:
# language_image = rendering.clone().permute(1, 2, 0).unsqueeze(0)
# lvl, h, w, _ = language_image.shape
# restored_feat = ae_model.decode(language_image.flatten(0, 2))
# restored_feat = restored_feat.view(lvl, h, w, -1)
# relevancy_map = clip_model.get_max_across(restored_feat) #.sum(dim=0)
# n_head, n_prompt, h, w = relevancy_map.shape
# for i in range(relevancy_map.shape[1]):
# scale = 30
# kernel = np.ones((scale,scale)) / (scale**2)
# np_relev = relevancy_map[:, i].cpu().numpy()
# avg_filtered = cv2.filter2D(np_relev.transpose(1,2,0), -1, kernel)
# h, w = avg_filtered.shape
# avg_filtered = avg_filtered.reshape(h, w, 1)
# score = avg_filtered[..., 0].max()
# coord = np.nonzero(avg_filtered[..., 0] == score)
# cx, cy = round(coord[1][0]), round(coord[0][0])
# # breakpoint()
# # score_lvl[0] = score
# width = 20
# x1 , y1 , x2, y2 = round(max(cx - width, 0)), round(max(cy - width, 0)), round(min(w, cx + width)), round(min(h, cy + width))
# rendering[:, y1:y2, x1:x2] = 1.0
# break
np.save(os.path.join(render_npy_path, '{0:05d}'.format(idx) + ".npy"),rendering.permute(1,2,0).cpu().numpy())
np.save(os.path.join(gts_npy_path, '{0:05d}'.format(idx) + ".npy"),gt.permute(1,2,0).cpu().numpy())
torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(idx) + ".png"))
torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ".png"))
breakpoint()
break
def render_sets(dataset : ModelParams, iteration : int, pipeline : PipelineParams, skip_train : bool, skip_test : bool, args):
with torch.no_grad():
gaussians = GaussianModel(dataset.sh_degree)
scene = Scene(dataset, gaussians, shuffle=False)
checkpoint = os.path.join(args.model_path, 'chkpnt30000.pth')
(model_params, first_iter) = torch.load(checkpoint)
gaussians.restore(model_params, args, mode='test')
ae_model = Autoencoder([256, 128, 64, 32, 3], [16, 32, 64, 128, 256, 256, 512]).to("cuda")
ae_model.load_state_dict(torch.load('autoencoder/ckpt/office_scene_50/best_ckpt.pth', map_location='cuda'))
ae_model.eval()
clip_model = OpenCLIPNetwork("cuda")
clip_model.set_positives(["bottle", "sanitizer", "tv", "screen", "television", "chair"])
bg_color = [1,1,1] if dataset.white_background else [0, 0, 0]
background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
if not skip_train:
render_set(dataset.model_path, dataset.source_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipeline, background, args, ae_model, clip_model)
if not skip_test:
render_set(dataset.model_path, dataset.source_path, "test", scene.loaded_iter, scene.getTestCameras(), gaussians, pipeline, background, args, ae_model, clip_model)
if __name__ == "__main__":
# Set up command line argument parser
parser = ArgumentParser(description="Testing script parameters")
model = ModelParams(parser, sentinel=True)
pipeline = PipelineParams(parser)
parser.add_argument("--iteration", default=-1, type=int)
parser.add_argument("--skip_train", action="store_true")
parser.add_argument("--skip_test", action="store_true")
parser.add_argument("--quiet", action="store_true")
parser.add_argument("--include_feature", action="store_true")
args = get_combined_args(parser)
print("Rendering " + args.model_path)
safe_state(args.quiet)
render_sets(model.extract(args), args.iteration, pipeline.extract(args), args.skip_train, args.skip_test, args) |