202 lines
8.1 KiB
Python
202 lines
8.1 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import comfy.model_management as mm
|
|
|
|
from .utils import (
|
|
IMAGENET_MEAN, IMAGENET_STD, DEFAULT_PATCH_SIZE,
|
|
format_camera_params, process_tensor_to_image, process_tensor_to_mask,
|
|
resize_to_patch_multiple, check_model_capabilities,
|
|
imagenet_normalize, save_gaussians_to_ply,
|
|
)
|
|
from .normalization import (
|
|
apply_edge_antialiasing,
|
|
apply_standard_normalization,
|
|
apply_v2_style_normalization,
|
|
apply_raw_normalization,
|
|
)
|
|
|
|
|
|
class DepthAnything_V3():
|
|
@classmethod
|
|
def execute(cls, da3_model, images, normalization_mode="V2-Style", camera_params=None, resize_method="resize", invert_depth=False, keep_model_size=False):
|
|
device = mm.get_torch_device()
|
|
|
|
mm.load_models_gpu([da3_model])
|
|
model = da3_model.model
|
|
capabilities = da3_model.model_options.get("da3_capabilities", check_model_capabilities(model))
|
|
dtype = da3_model.model_options.get("da3_dtype", torch.float16)
|
|
B, H, W, C = images.shape
|
|
images_pt = images.permute(0, 3, 1, 2)
|
|
images_pt, orig_H, orig_W = resize_to_patch_multiple(images_pt, DEFAULT_PATCH_SIZE, resize_method)
|
|
normalized_images = imagenet_normalize(images_pt)
|
|
normalized_images = normalized_images.unsqueeze(1)
|
|
extrinsics_input = None
|
|
intrinsics_input = None
|
|
if camera_params is not None:
|
|
if capabilities["has_camera_conditioning"]:
|
|
extrinsics_input = camera_params["extrinsics"].to(device).to(dtype)
|
|
intrinsics_input = camera_params["intrinsics"].to(device).to(dtype)
|
|
if extrinsics_input.shape[0] == 1 and B > 1:
|
|
extrinsics_input = extrinsics_input.expand(B, -1, -1, -1)
|
|
intrinsics_input = intrinsics_input.expand(B, -1, -1, -1)
|
|
|
|
depth_out = []
|
|
conf_out = []
|
|
sky_out = []
|
|
ray_origin_out = []
|
|
ray_dir_out = []
|
|
extrinsics_list = []
|
|
intrinsics_list = []
|
|
gaussians_list = []
|
|
|
|
infer_gs = capabilities["has_3d_gaussians"]
|
|
for i in range(B):
|
|
img = normalized_images[i:i+1].to(device, dtype=dtype)
|
|
ext_i = extrinsics_input[i:i+1] if extrinsics_input is not None else None
|
|
int_i = intrinsics_input[i:i+1] if intrinsics_input is not None else None
|
|
output = model(img, extrinsics=ext_i, intrinsics=int_i, infer_gs=infer_gs)
|
|
|
|
depth = None
|
|
if hasattr(output, 'depth'):
|
|
depth = output.depth
|
|
elif isinstance(output, dict) and 'depth' in output:
|
|
depth = output['depth']
|
|
|
|
if depth is None or not torch.is_tensor(depth):
|
|
raise ValueError("Model output does not contain valid depth tensor")
|
|
|
|
conf = None
|
|
if hasattr(output, 'depth_conf'):
|
|
conf = output.depth_conf
|
|
elif isinstance(output, dict) and 'depth_conf' in output:
|
|
conf = output['depth_conf']
|
|
if conf is None or not torch.is_tensor(conf):
|
|
conf = torch.ones_like(depth)
|
|
sky = None
|
|
if hasattr(output, 'sky'):
|
|
sky = output.sky
|
|
elif isinstance(output, dict) and 'sky' in output:
|
|
sky = output['sky']
|
|
|
|
if sky is None or not torch.is_tensor(sky):
|
|
sky = torch.zeros_like(depth)
|
|
else:
|
|
sky_min, sky_max = sky.min(), sky.max()
|
|
if sky_max > sky_min:
|
|
sky = (sky - sky_min) / (sky_max - sky_min)
|
|
|
|
if normalization_mode == "Raw":
|
|
depth_processed = apply_raw_normalization(depth, invert_depth)
|
|
elif normalization_mode == "V2-Style":
|
|
depth_processed = apply_v2_style_normalization(depth, sky, device, invert_depth)
|
|
else: # "Standard"
|
|
depth_processed = apply_standard_normalization(depth, invert_depth)
|
|
|
|
conf_range = conf.max() - conf.min()
|
|
if conf_range > 1e-8:
|
|
conf = (conf - conf.min()) / conf_range
|
|
else:
|
|
conf = torch.ones_like(conf)
|
|
|
|
depth_out.append(depth_processed.cpu())
|
|
conf_out.append(conf.cpu())
|
|
sky_out.append(sky.cpu())
|
|
|
|
ray = None
|
|
if hasattr(output, 'ray'):
|
|
ray = output.ray
|
|
elif isinstance(output, dict) and 'ray' in output:
|
|
ray = output['ray']
|
|
|
|
if ray is not None and torch.is_tensor(ray):
|
|
ray = ray.squeeze(0).squeeze(0) # [6, H, W]
|
|
ray_origin = ray[:3]
|
|
ray_dir = ray[3:6]
|
|
ray_origin_out.append(ray_origin.cpu())
|
|
ray_dir_out.append(ray_dir.cpu())
|
|
else:
|
|
ray_origin_out.append(torch.zeros(3, depth.shape[-2], depth.shape[-1]))
|
|
ray_dir_out.append(torch.zeros(3, depth.shape[-2], depth.shape[-1]))
|
|
|
|
extr = None
|
|
if hasattr(output, 'extrinsics'):
|
|
extr = output.extrinsics
|
|
elif isinstance(output, dict) and 'extrinsics' in output:
|
|
extr = output['extrinsics']
|
|
|
|
if extr is not None and torch.is_tensor(extr):
|
|
extrinsics_list.append(extr.cpu())
|
|
else:
|
|
extrinsics_list.append(None)
|
|
|
|
intr = None
|
|
if hasattr(output, 'intrinsics'):
|
|
intr = output.intrinsics
|
|
elif isinstance(output, dict) and 'intrinsics' in output:
|
|
intr = output['intrinsics']
|
|
|
|
if intr is not None and torch.is_tensor(intr):
|
|
intrinsics_list.append(intr.cpu())
|
|
else:
|
|
intrinsics_list.append(None)
|
|
|
|
if infer_gs:
|
|
gs = None
|
|
if hasattr(output, 'gaussians'):
|
|
gs = output.gaussians
|
|
elif isinstance(output, dict) and 'gaussians' in output:
|
|
gs = output['gaussians']
|
|
|
|
if gs is not None and hasattr(gs, 'means') and torch.is_tensor(gs.means):
|
|
gaussians_list.append((gs, depth))
|
|
|
|
normalize_depth_output = (normalization_mode != "Raw")
|
|
depth_final = process_tensor_to_image(depth_out, orig_H, orig_W, normalize_output=normalize_depth_output, skip_resize=keep_model_size)
|
|
|
|
if not keep_model_size:
|
|
final_H = (orig_H // 2) * 2
|
|
final_W = (orig_W // 2) * 2
|
|
model_H, model_W = images_pt.shape[2], images_pt.shape[3]
|
|
|
|
if final_H != model_H or final_W != model_W:
|
|
scale_h = final_H / model_H
|
|
scale_w = final_W / model_W
|
|
|
|
for i, intr in enumerate(intrinsics_list):
|
|
if intr is not None and torch.is_tensor(intr):
|
|
intr_scaled = intr.squeeze().clone()
|
|
intr_scaled[0, 0] *= scale_w # fx
|
|
intr_scaled[1, 1] *= scale_h # fy
|
|
intr_scaled[0, 2] *= scale_w # cx
|
|
intr_scaled[1, 2] *= scale_h # cy
|
|
intrinsics_list[i] = intr_scaled
|
|
|
|
return depth_final
|
|
|
|
@staticmethod
|
|
def _process_ray_to_image(ray_list, orig_H, orig_W, normalize=True, skip_resize=False):
|
|
out = torch.cat([r.unsqueeze(0) for r in ray_list], dim=0)
|
|
|
|
if normalize:
|
|
for i in range(out.shape[0]):
|
|
ray_batch = out[i]
|
|
ray_min = ray_batch.min()
|
|
ray_max = ray_batch.max()
|
|
if ray_max > ray_min:
|
|
out[i] = (ray_batch - ray_min) / (ray_max - ray_min)
|
|
else:
|
|
out[i] = torch.zeros_like(ray_batch)
|
|
|
|
out = out.permute(0, 2, 3, 1).float()
|
|
|
|
if not skip_resize:
|
|
final_H = (orig_H // 2) * 2
|
|
final_W = (orig_W // 2) * 2
|
|
if out.shape[1] != final_H or out.shape[2] != final_W:
|
|
out = F.interpolate(out.permute(0, 3, 1, 2), size=(final_H, final_W), mode="bilinear").permute(0, 2, 3, 1)
|
|
|
|
if normalize:
|
|
return torch.clamp(out, 0, 1)
|
|
else:
|
|
return out
|