-
Notifications
You must be signed in to change notification settings - Fork 52
/
Copy patheval_baseline.py
165 lines (138 loc) · 8.33 KB
/
eval_baseline.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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
import os
import sys
from pathlib import Path
if (_package_root := str(Path(__file__).absolute().parents[2])) not in sys.path:
sys.path.insert(0, _package_root)
import json
from typing import *
import importlib
import importlib.util
import click
@click.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True}, help='Evaluation script.')
@click.option('--baseline', 'baseline_code_path', type=click.Path(), required=True, help='Path to the baseline model python code.')
@click.option('--config', 'config_path', type=click.Path(), default='configs/eval/all_benchmarks.json', help='Path to the evaluation configurations. '
'Defaults to "configs/eval/all_benchmarks.json".')
@click.option('--output', '-o', 'output_path', type=click.Path(), required=True, help='Path to the output json file.')
@click.option('--oracle', 'oracle_mode', is_flag=True, help='Use oracle mode for evaluation, i.e., use the GT intrinsics input.')
@click.option('--dump_pred', is_flag=True, help='Dump predition results.')
@click.option('--dump_gt', is_flag=True, help='Dump ground truth.')
@click.pass_context
def main(ctx: click.Context, baseline_code_path: str, config_path: str, oracle_mode: bool, output_path: Union[str, Path], dump_pred: bool, dump_gt: bool):
# Lazy import
import cv2
import numpy as np
from tqdm import tqdm
import torch
import torch.nn.functional as F
import utils3d
from moge.test.baseline import MGEBaselineInterface
from moge.test.dataloader import EvalDataLoaderPipeline
from moge.test.metrics import compute_metrics
from moge.utils.geometry_torch import intrinsics_to_fov
from moge.utils.vis import colorize_depth, colorize_normal
from moge.utils.tools import key_average, flatten_nested_dict, timeit, import_file_as_module
# Load the baseline model
module = import_file_as_module(baseline_code_path, Path(baseline_code_path).stem)
baseline_cls: Type[MGEBaselineInterface] = getattr(module, 'Baseline')
baseline : MGEBaselineInterface = baseline_cls.load.main(ctx.args, standalone_mode=False)
# Load the evaluation configurations
with open(config_path, 'r') as f:
config = json.load(f)
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
all_metrics = {}
# Iterate over the dataset
for benchmark_name, benchmark_config in tqdm(list(config.items()), desc='Benchmarks'):
filenames, metrics_list = [], []
with (
EvalDataLoaderPipeline(**benchmark_config) as eval_data_pipe,
tqdm(total=len(eval_data_pipe), desc=benchmark_name, leave=False) as pbar
):
# Iterate over the samples in the dataset
for i in range(len(eval_data_pipe)):
sample = eval_data_pipe.get()
sample = {k: v.to(baseline.device) if isinstance(v, torch.Tensor) else v for k, v in sample.items()}
image = sample['image']
gt_intrinsics = sample['intrinsics']
# Inference
torch.cuda.synchronize()
with torch.inference_mode(), timeit('_inference_timer', verbose=False) as timer:
if oracle_mode:
pred = baseline.infer_for_evaluation(image, gt_intrinsics)
else:
pred = baseline.infer_for_evaluation(image)
torch.cuda.synchronize()
# Compute metrics
metrics, misc = compute_metrics(pred, sample, vis=dump_pred or dump_gt)
metrics['inference_time'] = timer.time
metrics_list.append(metrics)
# Dump results
dump_path = Path(output_path.replace(".json", f"_dump"), f'{benchmark_name}', sample['filename'].replace('.zip', ''))
if dump_pred:
dump_path.joinpath('pred').mkdir(parents=True, exist_ok=True)
cv2.imwrite(str(dump_path / 'pred' / 'image.jpg'), cv2.cvtColor((image.cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR))
with Path(dump_path, 'pred', 'metrics.json').open('w') as f:
json.dump(metrics, f, indent=4)
if 'pred_points' in misc:
points = misc['pred_points'].cpu().numpy()
cv2.imwrite(str(dump_path / 'pred' / 'points.exr'), cv2.cvtColor(points.astype(np.float32), cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT])
if 'pred_depth' in misc:
depth = misc['pred_depth'].cpu().numpy()
if 'mask' in pred:
mask = pred['mask'].cpu().numpy()
depth = np.where(mask, depth, np.inf)
cv2.imwrite(str(dump_path / 'pred' / 'depth.png'), cv2.cvtColor(colorize_depth(depth), cv2.COLOR_RGB2BGR))
if 'mask' in pred:
mask = pred['mask'].cpu().numpy()
cv2.imwrite(str(dump_path / 'pred' / 'mask.png'), (mask * 255).astype(np.uint8))
if 'normal' in pred:
normal = pred['normal'].cpu().numpy()
cv2.imwrite(str(dump_path / 'pred' / 'normal.png'), cv2.cvtColor(colorize_normal(normal), cv2.COLOR_RGB2BGR))
if 'intrinsics' in pred:
intrinsics = pred['intrinsics']
fov_x, fov_y = intrinsics_to_fov(intrinsics)
with open(dump_path / 'pred' / 'fov.json', 'w') as f:
json.dump({
'fov_x': np.rad2deg(fov_x.item()),
'fov_y': np.rad2deg(fov_y.item()),
'intrinsics': intrinsics.cpu().numpy().tolist(),
}, f)
if dump_gt:
dump_path.joinpath('gt').mkdir(parents=True, exist_ok=True)
cv2.imwrite(str(dump_path / 'gt' / 'image.jpg'), cv2.cvtColor((image.cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR))
if 'points' in sample:
points = sample['points']
cv2.imwrite(str(dump_path / 'gt' / 'points.exr'), cv2.cvtColor(points.cpu().numpy().astype(np.float32), cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT])
if 'depth' in sample:
depth = sample['depth']
mask = sample['depth_mask']
cv2.imwrite(str(dump_path / 'gt' / 'depth.png'), cv2.cvtColor(colorize_depth(depth.cpu().numpy(), mask=mask.cpu().numpy()), cv2.COLOR_RGB2BGR))
if 'normal' in sample:
normal = sample['normal']
cv2.imwrite(str(dump_path / 'gt' / 'normal.png'), cv2.cvtColor(colorize_normal(normal.cpu().numpy()), cv2.COLOR_RGB2BGR))
if 'depth_mask' in sample:
mask = sample['depth_mask']
cv2.imwrite(str(dump_path / 'gt' /'mask.png'), (mask.cpu().numpy() * 255).astype(np.uint8))
if 'intrinsics' in sample:
intrinsics = sample['intrinsics']
fov_x, fov_y = intrinsics_to_fov(intrinsics)
with open(dump_path / 'gt' / 'info.json', 'w') as f:
json.dump({
'fov_x': np.rad2deg(fov_x.item()),
'fov_y': np.rad2deg(fov_y.item()),
'intrinsics': intrinsics.cpu().numpy().tolist(),
}, f)
# Save intermediate results
if i % 100 == 0 or i == len(eval_data_pipe) - 1:
Path(output_path).write_text(
json.dumps({
**all_metrics,
benchmark_name: key_average(metrics_list)
}, indent=4)
)
pbar.update(1)
all_metrics[benchmark_name] = key_average(metrics_list)
# Save final results
all_metrics['mean'] = key_average(list(all_metrics.values()))
Path(output_path).write_text(json.dumps(all_metrics, indent=4))
if __name__ == '__main__':
main()