-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdisprefine.py
executable file
·102 lines (82 loc) · 4.06 KB
/
disprefine.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
import numpy as np
import h5py
import cv2
from DL import inference_single, inference_single_v2, InferDepthPro, InferDAM
def image_resize(img, sz):
return cv2.resize(img, sz, interpolation=cv2.INTER_LANCZOS4)
def load_disp_from_mat(path, dataset_name="data"):
"""load left disparity from .mat file."""
with h5py.File(path, "r") as file:
print("Keys: %s" % file.keys())
dataset = file[dataset_name]
disp = dataset[:]
return disp
def save_numpy_array_to_matlab(arr, save_path):
with h5py.File(save_path, 'w') as file:
file.create_dataset('data', data=arr)
print("===> done !")
def save_numpy_array(arr, save_path):
np.save(save_path, arr)
print("===> done !")
def make_colormap(m:np.ndarray) -> np.ndarray:
m_norm = cv2.normalize(m, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
m_cm = cv2.applyColorMap(m_norm, cv2.COLORMAP_VIRIDIS)
return m_cm
if __name__ == "__main__":
import os
m_types = ['vits', 'vitb', 'vitl']
# model_path="D:/william/checkpoints/depth-anything/v2/depth_anything_v2_vitb.pth"
# subfolders = [11, 12, 19]
# top_dir = "/home/william/extdisk/data/test_mono/raw/"
# predicted_save_dir = f"{top_dir}/output"
# os.makedirs(predicted_save_dir, exist_ok=True)
# subfolders = [f for f in os.listdir(top_dir) if os.path.isdir(os.path.join(top_dir, f))]
# subfiles = [f for f in os.listdir(top_dir) if os.path.isfile(os.path.join(top_dir, f)) and f.endswith(".png")]
# for type in m_types:
# print(f"processing {type} right now ...")
# model_path = f"/home/william/extdisk/checkpoints/depth-anything/depth_anything_v2_{type}.pth"
# for sub in subfiles:
# print(f"processing {sub} ...")
# # path = f"{top_dir}/{sub}/output_0222_agg_mask.mat"
# # disp = load_disp_from_mat(path, "out/s1out/cleanDispL")
# # print(disp.T.shape)
# # disp_save_path = f'{top_dir}/{sub}/output_0222_agg.npy'
# # save_numpy_array(disp.T, disp_save_path)
# # image_path = f"{top_dir}/{sub}/left.bmp"
# image_path = f"{top_dir}/{sub}"
# image = cv2.imread(image_path)
# # predicted = inference_single(image, model_path)
# predicted = inference_single_v2(image, model_path, encoder=f'{type}')
# # print(predicted)
# # print(predicted.shape)
# # save_path = "data/11/output_0222_DL.mat"
# # save_numpy_array_to_matlab(predicted, save_path)
# predicted_save_path = f"{predicted_save_dir}/dam_v2_{type}_out_{sub}.npy"
# save_numpy_array(predicted, predicted_save_path)
dp_model_path="/home/william/extdisk/checkpoints/depth-pro/depth_pro.pt"
dam_model_path="/home/william/extdisk/checkpoints/depth-anything/depth_anything_v2_vitl.pth"
dp_save_path = "/home/william/extdisk/data/mono/dp_out"
dam_save_path = "/home/william/extdisk/data/mono/dam_out"
os.makedirs(dp_save_path, exist_ok=True)
os.makedirs(dam_save_path, exist_ok=True)
dp_infer = InferDepthPro()
dam_infer = InferDAM()
dp_infer.initialize(model_path=dp_model_path, is_half=False)
dam_infer.initialize(model_path=dam_model_path)
dp_infer.get_device_info()
dam_infer.get_params_count()
image_path = "/home/william/extdisk/data/mono/left.png"
depth, est_focal = dp_infer.infer(image_path)
depth = depth.detach().cpu().numpy()
print(f"out depth shape is {depth.shape}")
predicted_save_path = f"{dp_save_path}/dp_12_f_{est_focal}.npy"
save_numpy_array(depth, predicted_save_path)
print(f"estimated focal length is: {est_focal}")
depth_dam = dam_infer.infer(image_path)
predicted_save_path = f"{dam_save_path}/dam_12_vitl.npy"
save_numpy_array(depth_dam, predicted_save_path)
dp_cm = make_colormap(depth)
dam_cm = make_colormap(depth_dam)
cv2.imwrite(f"{dp_save_path}/dp_12_f_{est_focal}.jpg", dp_cm)
cv2.imwrite(f"{dam_save_path}/dp_12_f_{est_focal}.jpg", dam_cm)
print("=====================> done! <=====================")