| import os | |
| from accelerate import Accelerator | |
| from accelerate.logging import get_logger | |
| from accelerate.utils import ( | |
| InitProcessGroupKwargs, | |
| ProjectConfiguration, | |
| set_seed | |
| ) | |
| import torch | |
| from contextlib import nullcontext | |
| import trimesh | |
| import gradio as gr | |
| from gradio_imageslider import ImageSlider | |
| from da2.utils.base import load_config | |
| from da2.utils.model import load_model | |
| from da2.utils.io import ( | |
| read_cv2_image, | |
| torch_transform, | |
| tensorize | |
| ) | |
| from da2.utils.vis import colorize_distance | |
| from da2.utils.d2pc import distance2pointcloud | |
| from datetime import ( | |
| timedelta, | |
| datetime | |
| ) | |
| import cv2 | |
| import numpy as np | |
| last_glb_path = None | |
| def prepare_to_run_demo(): | |
| config = load_config('configs/infer.json') | |
| kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=config['accelerator']['timeout'])) | |
| output_dir = f'output/infer' | |
| if not os.path.exists(output_dir): os.makedirs(output_dir, exist_ok=True) | |
| accu_steps = config['accelerator']['accumulation_nsteps'] | |
| accelerator = Accelerator( | |
| gradient_accumulation_steps=accu_steps, | |
| mixed_precision=config['accelerator']['mixed_precision'], | |
| log_with=config['accelerator']['report_to'], | |
| project_config=ProjectConfiguration(project_dir=output_dir), | |
| kwargs_handlers=[kwargs] | |
| ) | |
| logger = get_logger(__name__, log_level='INFO') | |
| config['env']['logger'] = logger | |
| set_seed(config['env']['seed']) | |
| return config, accelerator | |
| def read_mask_demo(mask_path, shape): | |
| if mask_path is None: | |
| return np.ones(shape[1:]) > 0 | |
| mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) | |
| mask = mask > 0 | |
| return mask | |
| def load_infer_data_demo(image, mask, model_dtype, device): | |
| cv2_image = read_cv2_image(image) | |
| image = torch_transform(cv2_image) | |
| mask = read_mask_demo(mask, image.shape) | |
| image = tensorize(image, model_dtype, device) | |
| return image, cv2_image, mask | |
| def ply2glb(ply_path, glb_path): | |
| pcd = trimesh.load(ply_path) | |
| points = np.asarray(pcd.vertices) | |
| colors = np.asarray(pcd.visual.vertex_colors) | |
| cloud = trimesh.points.PointCloud(vertices=points, colors=colors) | |
| cloud.export(glb_path) | |
| os.remove(ply_path) | |
| def fn(image_path, mask_path): | |
| global last_glb_path | |
| config, accelerator = prepare_to_run_demo() | |
| model = load_model(config, accelerator) | |
| image, cv2_image, mask = load_infer_data_demo(image_path, mask_path, | |
| model_dtype=config['spherevit']['dtype'], device=accelerator.device) | |
| if torch.backends.mps.is_available(): | |
| autocast_ctx = nullcontext() | |
| else: | |
| autocast_ctx = torch.autocast(accelerator.device.type) | |
| with autocast_ctx, torch.no_grad(): | |
| distance = model(image).cpu().numpy()[0] | |
| if last_glb_path is not None: | |
| os.remove(last_glb_path) | |
| distance_vis = colorize_distance(distance, mask) | |
| save_path = f'cache/tmp_{datetime.now().strftime("%Y%m%d_%H%M%S")}.glb' | |
| last_glb_path = save_path | |
| normal_image = distance2pointcloud(distance, cv2_image, mask, save_path=save_path.replace('.glb', '.ply'), return_normal=True, save_distance=False) | |
| ply2glb(save_path.replace('.glb', '.ply'), save_path) | |
| return save_path, [distance_vis, normal_image] | |
| inputs = [ | |
| gr.Image(label="Input Image", type="filepath"), | |
| gr.Image(label="Input Mask", type="filepath"), | |
| ] | |
| outputs = [ | |
| gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="3D Point Cloud"), | |
| gr.ImageSlider( | |
| label="Output Depth / Normal (transformed from the depth)", | |
| type="pil", | |
| slider_position=75, | |
| ) | |
| ] | |
| demo = gr.Interface( | |
| fn=fn, | |
| title="DA<sup>2</sup>: <u>D</u>epth <u>A</u>nything in <u>A</u>ny <u>D</u>irection", | |
| description=""" | |
| <p align="center"> | |
| <a title="Project Page" href="https://depth-any-in-any-dir.github.io/" target="_blank" rel="noopener noreferrer" style="display: inline-block;"> | |
| <img src="https://img.shields.io/badge/Project-Website-pink?logo=googlechrome&logoColor=white"> | |
| </a> | |
| <a title="arXiv" href="http://arxiv.org/abs/2509.26618" target="_blank" rel="noopener noreferrer" style="display: inline-block;"> | |
| <img src="https://img.shields.io/badge/arXiv-Paper-b31b1b?logo=arxiv&logoColor=white"> | |
| </a> | |
| <a title="Github" href="https://github.com/EnVision-Research/DA-2" target="_blank" rel="noopener noreferrer" style="display: inline-block;"> | |
| <img src="https://img.shields.io/github/stars/EnVision-Research/DA-2?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars"> | |
| </a> | |
| <a title="Social" href="https://x.com/_akhaliq/status/1973283687652606411" target="_blank" rel="noopener noreferrer" style="display: inline-block;"> | |
| <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social"> | |
| </a> | |
| <a title="Social" href="https://x.com/haodongli00/status/1973287870317338747" target="_blank" rel="noopener noreferrer" style="display: inline-block;"> | |
| <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social"> | |
| </a> | |
| <br> | |
| <strong>Please consider starring <span style="color: orange">★</span> our <a href="https://github.com/EnVision-Research/DA-2" target="_blank" rel="noopener noreferrer">GitHub Repo</a> if you find this demo useful!</strong> | |
| </p> | |
| <p><strong>Note: the "Input Mask" is optional, all pixels are assumed to be valid if mask is None.</strong></p> | |
| """, | |
| inputs=inputs, | |
| outputs=outputs, | |
| examples=[ | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/a1.png"), None], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/a2.png"), None], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/a3.png"), None], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/a4.png"), None], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/b0.png"), | |
| os.path.join(os.path.dirname(__file__), "assets/masks/b0.png")], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/b1.png"), | |
| os.path.join(os.path.dirname(__file__), "assets/masks/b1.png")], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/a5.png"), None], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/a6.png"), None], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/a7.png"), None], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/a8.png"), None], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/b2.png"), | |
| os.path.join(os.path.dirname(__file__), "assets/masks/b2.png")], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/b3.png"), | |
| os.path.join(os.path.dirname(__file__), "assets/masks/b3.png")], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/a9.png"), None], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/a10.png"), None], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/a11.png"), None], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/a0.png"), None], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/b4.png"), | |
| os.path.join(os.path.dirname(__file__), "assets/masks/b4.png")], | |
| [os.path.join(os.path.dirname(__file__), "assets/demos/b5.png"), | |
| os.path.join(os.path.dirname(__file__), "assets/masks/b5.png")], | |
| ], | |
| examples_per_page=20 | |
| ) | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=6381, | |
| ) | |