95 lines
2.7 KiB
Python
95 lines
2.7 KiB
Python
from ..geometry import unproject_camera_depth
|
|
from ..geometry import GeometryFrame
|
|
|
|
|
|
def merge_frame_predictions(predictions):
|
|
import numpy as np
|
|
|
|
if not predictions:
|
|
raise ValueError("DA3 returned no frame predictions")
|
|
merged = predictions[0]
|
|
for name in (
|
|
"depth",
|
|
"conf",
|
|
"extrinsics",
|
|
"intrinsics",
|
|
"processed_images",
|
|
):
|
|
setattr(
|
|
merged,
|
|
name,
|
|
np.concatenate(
|
|
[getattr(prediction, name) for prediction in predictions],
|
|
axis=0,
|
|
),
|
|
)
|
|
return merged
|
|
|
|
|
|
def infer(
|
|
model_manager,
|
|
parameters,
|
|
images,
|
|
*,
|
|
progress_callback=None,
|
|
cancel_check=None,
|
|
):
|
|
from .onnx_da3 import infer as infer_onnx
|
|
from .onnx_runtime import provider_label
|
|
|
|
model_directory = model_manager.directory(parameters["model"])
|
|
device = parameters["device"]
|
|
model = model_manager.get_da3(model_directory, device)
|
|
options = {
|
|
"process_res": parameters["process_res"],
|
|
"process_res_method": parameters["process_res_method"],
|
|
"ref_view_strategy": parameters["ref_view_strategy"],
|
|
}
|
|
|
|
def run(paths):
|
|
try:
|
|
return infer_onnx(model, paths, **options)
|
|
except UnicodeDecodeError as error:
|
|
if provider_label(model) == "DirectML":
|
|
raise RuntimeError(
|
|
"DA3 inference failed on DirectML. Switch the device to "
|
|
"CPU in AnyImage Preferences and retry. "
|
|
f"Original error: {error}"
|
|
) from error
|
|
raise
|
|
|
|
if parameters.get("inference_mode", "joint") != "per_frame":
|
|
return run([str(path) for path in images])
|
|
|
|
predictions = []
|
|
for index, image in enumerate(images):
|
|
if cancel_check is not None:
|
|
cancel_check()
|
|
if progress_callback is not None:
|
|
progress_callback(index, len(images))
|
|
predictions.append(run([str(image)]))
|
|
return merge_frame_predictions(predictions)
|
|
|
|
|
|
def surface_frame(prediction):
|
|
import numpy as np
|
|
|
|
depth = np.asarray(prediction.depth, dtype=np.float32)
|
|
confidence = np.asarray(prediction.conf, dtype=np.float32)
|
|
intrinsics = np.asarray(prediction.intrinsics, dtype=np.float32)
|
|
if depth.shape[0] != 1:
|
|
raise ValueError("Convert to Surface currently supports one image")
|
|
points = unproject_camera_depth(depth[0], intrinsics[0])
|
|
valid = (
|
|
np.isfinite(depth[0])
|
|
& (depth[0] > 0.0)
|
|
& np.isfinite(points).all(axis=-1)
|
|
)
|
|
return GeometryFrame(
|
|
depth=depth[0],
|
|
points=points,
|
|
intrinsics=intrinsics[0],
|
|
valid_mask=valid,
|
|
confidence=confidence[0],
|
|
)
|