-
Notifications
You must be signed in to change notification settings - Fork 73
Expand file tree
/
Copy pathexample.py
More file actions
376 lines (302 loc) Β· 12.1 KB
/
example.py
File metadata and controls
376 lines (302 loc) Β· 12.1 KB
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
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
#!/usr/bin/env python3
"""
LingBot-Depth Example Script
Demonstrates depth refinement using the LingBot-Depth model.
Supports both local checkpoints and Hugging Face models.
Usage:
python example.py
python example.py --model robbyant/lingbot-depth-postrain-dc-vitl14
python example.py --example 0 --output custom_output
"""
import cv2
import torch
import numpy as np
import trimesh
import argparse
import time
from pathlib import Path
from mdm.model.v2 import MDMModel
def preprocess_input_image(image_path, device):
"""
Load and preprocess RGB image.
Args:
image_path (str): Path to RGB image
device (torch.device): Device to load tensor on
Returns:
tuple: (numpy_image, tensor_image)
- numpy_image: RGB numpy array (H, W, 3), uint8
- tensor_image: RGB tensor (1, 3, H, W), float32, [0,1]
"""
if not Path(image_path).exists():
raise FileNotFoundError(f"Image not found: {image_path}")
# Read image and convert BGR to RGB
image_np = cv2.imread(image_path)
if image_np is None:
raise ValueError(f"Failed to read image: {image_path}")
image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
# Convert to tensor and normalize to [0, 1]
image_tensor = torch.tensor(
image_np / 255.0,
dtype=torch.float32,
device=device
).permute(2, 0, 1).unsqueeze(0)
return image_np, image_tensor
def load_depth_map(depth_path, scale=1000.0):
"""
Load depth map from PNG file (16-bit) and convert to meters.
Args:
depth_path (str): Path to depth image
scale (float): Scale factor to convert to meters
- 1000.0 for millimeters
- 1.0 for meters
Returns:
np.ndarray: Depth map in meters (H, W), float32
"""
if not Path(depth_path).exists():
raise FileNotFoundError(f"Depth map not found: {depth_path}")
# Read depth map as 16-bit
depth_map = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)
if depth_map is None:
raise ValueError(f"Failed to read depth map: {depth_path}")
# Convert to meters
depth_map = depth_map.astype(np.float32) / scale
# Replace invalid values with 0
depth_map = np.nan_to_num(depth_map, nan=0.0, posinf=0.0, neginf=0.0)
return depth_map
def load_intrinsics(intrinsics_path, width, height):
"""
Load camera intrinsics and normalize by image dimensions.
Args:
intrinsics_path (str): Path to intrinsics file (.txt or .json)
width (int): Image width
height (int): Image height
Returns:
np.ndarray: Normalized intrinsics matrix (3, 3)
"""
if not Path(intrinsics_path).exists():
raise FileNotFoundError(f"Intrinsics not found: {intrinsics_path}")
# Load intrinsics
if intrinsics_path.endswith('.json'):
import json
with open(intrinsics_path, 'r') as f:
intrinsics = np.array(json.load(f), dtype=np.float32)
else:
intrinsics = np.loadtxt(intrinsics_path, dtype=np.float32)
# Normalize by image dimensions
intrinsics_normalized = intrinsics.copy()
intrinsics_normalized[0, 0] /= width # fx
intrinsics_normalized[0, 2] /= width # cx
intrinsics_normalized[1, 1] /= height # fy
intrinsics_normalized[1, 2] /= height # cy
return intrinsics_normalized
def depth_to_color_opencv(depth_map, vmin=None, vmax=None, colormap=cv2.COLORMAP_TURBO):
"""
Convert depth map to color visualization using OpenCV colormap.
Args:
depth_map (np.ndarray): Depth map (H, W)
vmin (float): Minimum depth for colormap (auto if None)
vmax (float): Maximum depth for colormap (auto if None)
colormap: OpenCV colormap (TURBO, JET, VIRIDIS, etc.)
Returns:
np.ndarray: Colored depth map (H, W, 3) in BGR format
"""
# Handle invalid values
valid_mask = np.isfinite(depth_map) & (depth_map > 0)
depth_clean = depth_map.copy()
depth_clean[~valid_mask] = 0
# Auto-range if not specified
if vmin is None:
vmin = depth_clean[valid_mask].min() if valid_mask.any() else 0
if vmax is None:
vmax = depth_clean[valid_mask].max() if valid_mask.any() else 1
# Normalize to [0, 255]
depth_normalized = np.clip(
(depth_clean - vmin) / (vmax - vmin + 1e-8) * 255,
0, 255
).astype(np.uint8)
# Apply colormap
depth_colored = cv2.applyColorMap(depth_normalized, colormap)
# Set invalid pixels to black
depth_colored[~valid_mask] = [0, 0, 0]
return depth_colored
def main():
"""Main function with argument parsing and execution."""
parser = argparse.ArgumentParser(
description='LingBot-Depth Example: Refine depth maps using the model',
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Run with default settings (example 0)
python example.py
# Use a different example directory
python example.py --example 1
# Use Hugging Face model
python example.py --model robbyant/lingbot-depth-postrain-dc-vitl14
# Use local checkpoint
python example.py --model ckpt/model.pt
# Custom output directory
python example.py --output my_results
"""
)
parser.add_argument(
'--example', type=str, default='0',
help='Example directory number (default: 0)'
)
parser.add_argument(
'--model', type=str,
default='robbyant/lingbot-depth-pretrain-vitl-14-v0.5',
help='Model path or Hugging Face ID (default: robbyant/lingbot-depth-pretrain-vitl-14)'
)
parser.add_argument(
'--output', type=str, default='result',
help='Output directory (default: result)'
)
parser.add_argument(
'--device', type=str, default='auto',
choices=['auto', 'cuda', 'cpu'],
help='Device to use (default: auto - uses CUDA if available)'
)
parser.add_argument(
'--no-mask', action='store_true',
help='Disable masking of invalid regions'
)
args = parser.parse_args()
# Print header
print("=" * 70)
print("LingBot-Depth Example: Depth Refinement".center(70))
print("=" * 70)
# Setup device
if args.device == 'auto':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
else:
device = torch.device(args.device)
print(f"\nπ± Device: {device}")
if device.type == 'cuda':
print(f" GPU: {torch.cuda.get_device_name(0)}")
# Setup paths
example_dir = Path('examples') / args.example
# Check if example directory exists first
if not example_dir.exists():
print(f"\nβ Error: Example directory not found: {example_dir}")
print(f"\n Available examples:")
examples_root = Path('examples')
if examples_root.exists():
for ex in sorted(examples_root.iterdir()):
if ex.is_dir():
# Check what files exist in this example
rgb_file = None
for ext in ['.png', '.jpg', '.jpeg']:
if (ex / f'rgb{ext}').exists():
rgb_file = f'rgb{ext}'
break
depth_ok = (ex / 'raw_depth.png').exists()
intrinsics_ok = (ex / 'intrinsics.txt').exists()
status = 'β' if all([rgb_file, depth_ok, intrinsics_ok]) else 'β '
print(f" {status} {ex.name}")
return 1
# Auto-detect RGB file format (.png or .jpg)
rgb_path = None
for ext in ['.png', '.jpg', '.jpeg']:
candidate = example_dir / f'rgb{ext}'
if candidate.exists():
rgb_path = candidate
break
depth_path = example_dir / 'raw_depth.png'
intrinsics_path = example_dir / 'intrinsics.txt'
# Check for required files
missing_files = []
if rgb_path is None:
missing_files.append('rgb.{png,jpg,jpeg}')
if not depth_path.exists():
missing_files.append('raw_depth.png')
if not intrinsics_path.exists():
missing_files.append('intrinsics.txt')
if missing_files:
print(f"\nβ Error: Missing required files in {example_dir}:")
for f in missing_files:
print(f" - {f}")
return 1
print(f"\nπ Input:")
print(f" RGB: {rgb_path}")
print(f" Depth: {depth_path}")
print(f" Intrinsics: {intrinsics_path}")
try:
# Load data
print(f"\nπ₯ Loading data...")
image_np, image_tensor = preprocess_input_image(str(rgb_path), device)
depth_np = load_depth_map(str(depth_path))
depth_tensor = torch.tensor(depth_np, dtype=torch.float32, device=device)
h, w = image_np.shape[:2]
print(f" Image size: {w}Γ{h}")
print(f" Input depth range: {depth_np[depth_np > 0].min():.2f} - {depth_np.max():.2f} meters")
# Load intrinsics
intrinsics = load_intrinsics(str(intrinsics_path), w, h)
intrinsics_tensor = torch.tensor(intrinsics, dtype=torch.float32, device=device).unsqueeze(0)
# Load model
print(f"\nπ€ Loading model: {args.model}")
start_time = time.time()
model = MDMModel.from_pretrained(args.model).to(device)
load_time = time.time() - start_time
print(f" β Model loaded in {load_time:.2f}s")
# Run inference
print(f"\nπ Running inference...")
start_time = time.time()
with torch.no_grad():
output = model.infer(
image_tensor,
depth_in=depth_tensor,
apply_mask=not args.no_mask,
intrinsics=intrinsics_tensor
)
inference_time = time.time() - start_time
depth_pred = output['depth'].squeeze().cpu().numpy()
points_pred = output['points'].squeeze().cpu().numpy()
print(f" β Inference completed in {inference_time:.3f}s")
print(f" Refined depth range: {depth_pred[depth_pred > 0].min():.2f} - {depth_pred.max():.2f} meters")
# Save results
output_dir = Path(args.output)
output_dir.mkdir(exist_ok=True, parents=True)
print(f"\nπΎ Saving results to: {output_dir}/")
# 1. Save depth maps as numpy arrays
np.save(output_dir / 'depth_input.npy', depth_np)
np.save(output_dir / 'depth_refined.npy', depth_pred)
print(f" β Depth arrays saved (.npy)")
# 2. Save depth visualizations
depth_raw_color = depth_to_color_opencv(depth_np)
depth_pred_color = depth_to_color_opencv(depth_pred)
depth_concat = np.concatenate([depth_raw_color, depth_pred_color], axis=1)
cv2.imwrite(str(output_dir / 'depth_input.png'), depth_raw_color)
cv2.imwrite(str(output_dir / 'depth_refined.png'), depth_pred_color)
cv2.imwrite(str(output_dir / 'depth_comparison.png'), depth_concat)
print(f" β Depth visualizations saved (.png)")
# 3. Save RGB image for reference
cv2.imwrite(str(output_dir / 'rgb.png'), cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR))
# 4. Save point cloud
valid_mask = np.isfinite(points_pred).all(axis=-1) & (points_pred[..., 2] > 0)
verts = points_pred[valid_mask]
verts_color = image_np[valid_mask]
# Downsample for reasonable file size
downsample = 2
verts = verts[::downsample]
verts_color = verts_color[::downsample]
point_cloud = trimesh.PointCloud(verts, verts_color)
point_cloud.export(output_dir / 'point_cloud.ply')
print(f" β Point cloud saved ({len(verts):,} points)")
# Print summary
print(f"\nπ Summary:")
print(f" Model load time: {load_time:.2f}s")
print(f" Inference time: {inference_time:.3f}s")
print(f" Valid points: {valid_mask.sum():,} / {valid_mask.size:,}")
print(f"\nβ
Done! Results saved to: {output_dir}/")
print("=" * 70)
return 0
except FileNotFoundError as e:
print(f"\nβ Error: {e}")
return 1
except Exception as e:
print(f"\nβ Error: {e}")
import traceback
traceback.print_exc()
return 1
if __name__ == '__main__':
exit(main())