-
Notifications
You must be signed in to change notification settings - Fork 1
/
findpair.py
151 lines (129 loc) · 5.34 KB
/
findpair.py
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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
#function takes in two GS and camera infos, output best image pair for dust3r alignment
#
# 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 [email protected]
#
import cv2
from scene.colmap_loader import qvec2rotmat, read_extrinsics_binary
from dinov2_utils import classify_image
import numpy as np
import torch.nn.functional as F
import os
from scene.cameras import myMiniCam3
from myUtils import gaus_copy
import torch
from gaussian_renderer import render
from scene import Scene, GaussianModel
#pick suitable campos and images for dust3r input by dinov2 similarity
def campick(g1, g2, extrinsics1, extrinsics2, dataset, pipe):
#canvas set up
bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
fovx = 0.812657831303291
fovy = 0.5579197285849142
znear = 0.01
zfar = 100.0
scaling_modifer = 1.0
pipe.convert_SHs_python = False
pipe.compute_cov3D_python = False
d2list1 = []
# cnt = 0
for key in list(extrinsics1.keys()):
img = extrinsics1[key]
qvec, tvec = img.qvec, img.tvec
rotmat = qvec2rotmat(qvec)
rotmat = torch.tensor(rotmat, dtype=torch.float32, device="cuda")
tvec = torch.tensor(tvec, dtype=torch.float32, device="cuda")
rotmat = rotmat.T
cam_pos = myMiniCam3(512, 336, rotmat, tvec, fovx, fovy, znear, zfar)
img1 = render(cam_pos, g1, pipe, background, scaling_modifer)["render"]
dv1 = classify_image(img1)
d2list1.append((dv1, rotmat, tvec))
d2list2 = []
for key in list(extrinsics2.keys()):
img = extrinsics2[key]
qvec, tvec = img.qvec, img.tvec
rotmat = qvec2rotmat(qvec)
rotmat = torch.tensor(rotmat, dtype=torch.float32, device="cuda")
tvec = torch.tensor(tvec, dtype=torch.float32, device="cuda")
rotmat = rotmat.T
cam_pos = myMiniCam3(512, 336, rotmat, tvec, fovx, fovy, znear, zfar)
img2 = render(cam_pos, g2, pipe, background, scaling_modifer)["render"]
dv2 = classify_image(img2)
d2list2.append((dv2, rotmat, tvec))
max_sim = 0
res = ()
for v1, r1, t1 in d2list1:
for v2, r2, t2 in d2list2:
cos_sim = F.cosine_similarity(v1, v2, dim = 0)
if cos_sim > max_sim and cos_sim < 0.91:
max_sim = cos_sim
res = (r1, t1, r2, t2)
r1, t1, r2, t2 = res
cam_pos1 = myMiniCam3(512, 336, r1, t1, fovx, fovy, znear, zfar)
I1 = render(cam_pos1, g1, pipe, background, scaling_modifer)
img1 = I1["render"]
dep1 = I1["depth"]
wei1 = I1["weight"]
cam_pos2 = myMiniCam3(512, 336, r2, t2, fovx, fovy, znear, zfar)
I2 = render(cam_pos2, g2, pipe, background, scaling_modifer)
img2 = I2["render"]
dep2 = I2["depth"]
wei2 = I2["weight"]
print("found img pair with sim:", max_sim)
return img1, img2, r1, t1, r2, t2, dep1, dep2, wei1, wei2
def tensor2np(img):
# Convert from [C, H, W] to [H, W, C] and move to CPU if needed
img = img.permute(1, 2, 0).cpu().detach().numpy()
# Ensure pixel values are in [0, 1] if float32 or scale if necessary
if img.dtype == np.float32 and img.max() > 1.0:
img = np.clip(img, 0.0, 1.0) # Clip to [0, 1]
# Convert from RGB to BGR for OpenCV display
img_bgr = cv2.cvtColor((img * 255).astype('uint8'), cv2.COLOR_RGB2BGR)
return img_bgr
def findpair(dataset, opt, pipe, checkpoint: str, checkpoint2: str, campose1, campose2):
g1 = GaussianModel(dataset.sh_degree)
g1.training_setup(opt)
g2 = GaussianModel(dataset.sh_degree)
g2.training_setup(opt)
scene = Scene(dataset, g1)
if checkpoint:
(model_params, first_iter) = torch.load(checkpoint)
g1.restore(model_params, opt)
if checkpoint2:
(model_params, first_iter) = torch.load(checkpoint2)
g2.restore(model_params, opt)
print("Second model loaded from checkpoint2.")
print("Both GS models loaded with following shape:", g1._xyz.shape, g2._xyz.shape)
bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
g1copy = GaussianModel(dataset.sh_degree)
g1copy.training_setup(opt)
g2copy = GaussianModel(dataset.sh_degree)
g2copy.training_setup(opt)
gaus_copy(g1, g1copy)
gaus_copy(g2, g2copy)
#set up extrinsincs cam dicts:
path1 = campose1
cameras_extrinsic_file1 = os.path.join(path1, "sparse/0", "images.bin")
cam_extrinsics1 = read_extrinsics_binary(cameras_extrinsic_file1)
path2 = campose2
cameras_extrinsic_file2 = os.path.join(path2, "sparse/0", "images.bin")
cam_extrinsics2 = read_extrinsics_binary(cameras_extrinsic_file2)
img1, img2, r1, t1, r2, t2, dep1, dep2, wei1, wei2 = campick(g1copy, g2copy, cam_extrinsics1, cam_extrinsics2, dataset, pipe)
npimg1 = tensor2np(img1)
npimg2 = tensor2np(img2)
cv2.imwrite('img1.png', npimg1)
cv2.imwrite('img2.png', npimg2)
cv2.imshow('Image', npimg1)
cv2.waitKey(0)
cv2.imshow('Image', npimg2)
cv2.waitKey(0)
cv2.destroyAllWindows()
return r1, t1, r2, t2