feat: Implement background depth distance feature for depth processing

- Added `background_depth_distance` property to `ActiveJob` and `RunJob` classes to manage background depth settings.
- Updated job execution to handle background depth distance, allowing for fixed negative displacement in depth calculations.
- Introduced new functions in the worker module to apply background depth distance to depth frames and manage original background depth.
- Enhanced tests to cover new functionality, ensuring correct behavior of background depth processing and reference depth extraction.
- Refactored job submission to run in a separate thread, improving responsiveness during execution.

src/snapform/worker.py

@@ -29,6 +29,7 @@ HF_ENDPOINT = "https://huggingface.co"
 MODELSCOPE_ENDPOINT = "https://modelscope.cn"
 MODELSCOPE_REVISION = "master"
 MODELSCOPE_TIMEOUT = 60.0
+REFERENCE_DEPTH_KEY = "snapform_reference_depth"
 
 
 class TaskCancelled(RuntimeError):
@@ -40,6 +41,144 @@ def check_cancelled(cancel_check):
         raise TaskCancelled("Task cancelled")
 
 
+def apply_foreground_depth_distance(
+    prediction,
+    foreground_paths,
+    background_depth_distance,
+):
+    """Set fully transparent pixels to a fixed negative displacement."""
+    import numpy as np
+    from PIL import Image
+
+    if len(foreground_paths) != len(prediction.depth):
+        raise RuntimeError("Foreground mask count does not match depth frames")
+    masks = []
+    height, width = prediction.depth.shape[-2:]
+    for foreground_path in foreground_paths:
+        with Image.open(foreground_path) as source:
+            alpha = source.convert("RGBA").getchannel("A")
+        alpha = alpha.resize((width, height), resample=Image.Resampling.NEAREST)
+        masks.append(np.asarray(alpha) > 0)
+    apply_depth_background_array(
+        prediction,
+        np.stack(masks),
+        background_depth_distance,
+    )
+
+
+def apply_depth_background_array(
+    prediction,
+    foreground_mask,
+    background_depth_distance,
+):
+    import numpy as np
+
+    depth = np.asarray(prediction.depth)
+    masked_depth = depth.copy()
+    background_depth = -background_depth_distance
+    for frame_index in range(len(masked_depth)):
+        np.maximum(
+            masked_depth[frame_index],
+            background_depth,
+            out=masked_depth[frame_index],
+        )
+        masked_depth[frame_index][~foreground_mask[frame_index]] = (
+            background_depth
+        )
+        set_depth_border(masked_depth[frame_index], background_depth)
+    prediction.depth = masked_depth
+
+
+def set_depth_border(depth, background_depth):
+    """Set the outermost pixel ring to the fixed background depth."""
+    if depth.ndim != 2 or not depth.size:
+        raise ValueError("Depth frame must be a non-empty 2D array")
+    depth[0, :] = background_depth
+    depth[-1, :] = background_depth
+    depth[:, 0] = background_depth
+    depth[:, -1] = background_depth
+
+
+def convert_ai_depth_for_displacement(prediction):
+    """Convert DA3 distance into Blender's positive displacement direction."""
+    import numpy as np
+
+    prediction.depth = 1.0 - np.asarray(prediction.depth)
+
+
+def depth_reference_value(depth):
+    """Return the unmasked depth used to size the imported camera plane."""
+    import numpy as np
+
+    values = np.asarray(depth)
+    valid = np.isfinite(values) & (values > 0.0)
+    if not valid.any():
+        raise ValueError("Depth array contains no positive finite values")
+    return float(np.median(values[valid]))
+
+
+def write_npz_reference_depth(npz_path, reference_depth):
+    """Add SnapForm metadata without changing exported prediction arrays."""
+    import numpy as np
+
+    npz_path = Path(npz_path)
+    temporary_path = npz_path.with_name(f"{npz_path.stem}.tmp.npz")
+    with np.load(npz_path, allow_pickle=False) as archive:
+        values = {key: archive[key] for key in archive.files}
+    values[REFERENCE_DEPTH_KEY] = np.asarray(reference_depth, dtype=np.float32)
+    np.savez_compressed(temporary_path, **values)
+    temporary_path.replace(npz_path)
+
+
+def write_original_background_depth(
+    image_paths,
+    foreground_paths,
+    background_depth_distance,
+    output_directory,
+):
+    """Write original-image luminance with a fixed negative background depth."""
+    import numpy as np
+    from PIL import Image
+
+    if len(image_paths) != len(foreground_paths):
+        raise RuntimeError("Foreground mask count does not match source frames")
+    depth_frames = []
+    for image_path, foreground_path in zip(image_paths, foreground_paths):
+        with Image.open(image_path) as source:
+            luminance = np.asarray(source.convert("L"), dtype=np.float32) / 255.0
+        with Image.open(foreground_path) as foreground:
+            alpha = np.asarray(
+                foreground.convert("RGBA").getchannel("A"),
+                dtype=np.uint8,
+            )
+        if luminance.shape != alpha.shape:
+            raise RuntimeError("Foreground mask size does not match source frame")
+        depth_frames.append(
+            compose_original_background_depth(
+                luminance,
+                alpha,
+                background_depth_distance,
+            )
+        )
+    np.savez_compressed(
+        Path(output_directory) / "original-depth.npz",
+        depth=np.stack(depth_frames),
+    )
+
+
+def compose_original_background_depth(
+    luminance,
+    alpha,
+    background_depth_distance,
+):
+    import numpy as np
+
+    background_depth = -background_depth_distance
+    depth = np.where(alpha > 0, luminance, background_depth)
+    set_depth_border(depth, background_depth)
+    return depth
+
+
 def write_status(
     path,
     progress,
@@ -757,6 +896,21 @@ def remove_background(args, model_cache=None, cancel_check=None):
                 cancel_check=cancel_check,
             )
             copy_frames(images, output_directory)
+        background_depth_distance = float(
+            getattr(args, "background_depth_distance", 0.0)
+        )
+        if background_depth_distance > 0.0:
+            foreground_paths = (
+                [output_directory / "foreground.png"]
+                if len(images) == 1
+                else sorted(output_directory.glob("foreground.*.png"))
+            )
+            write_original_background_depth(
+                images,
+                foreground_paths,
+                background_depth_distance,
+                output_directory,
+            )
     finally:
         if temporary_directory is not None:
             shutil.rmtree(temporary_directory, ignore_errors=True)
@@ -878,6 +1032,10 @@ def reconstruct(args, model_cache=None, cancel_check=None):
             "remove_background_model_dir",
             "",
         )
+        background_depth_distance = float(
+            getattr(args, "background_depth_distance", 0.0)
+        )
+        foreground_images = []
         if remove_background_model_dir:
             write_status(args.status, 0.12, "Removing image background")
             if len(images) == 1:
@@ -897,6 +1055,7 @@ def reconstruct(args, model_cache=None, cancel_check=None):
                 )
                 if not (output_directory / "foreground.png").is_file():
                     raise RuntimeError("BEN2 did not write a foreground image")
+                foreground_images = [output_directory / "foreground.png"]
             else:
                 load_ms, inference_ms = infer_foreground_frames(
                     images,
@@ -910,6 +1069,10 @@ def reconstruct(args, model_cache=None, cancel_check=None):
                     model_cache=model_cache,
                     cancel_check=cancel_check,
                 )
+                foreground_images = sorted(
+                    output_directory.glob("foreground.*.png")
+                )
+        if remove_background_model_dir:
             print(
                 f"BEN2 inference {inference_ms:.0f} ms; "
                 f"model load {load_ms:.0f} ms"
@@ -950,6 +1113,16 @@ def reconstruct(args, model_cache=None, cancel_check=None):
                 ref_view_strategy=args.ref_view_strategy,
             )
         check_cancelled(cancel_check)
+        reference_depth = None
+        if args.export_format == "mini_npz":
+            reference_depth = depth_reference_value(prediction.depth[0])
+            convert_ai_depth_for_displacement(prediction)
+        if background_depth_distance > 0.0:
+            apply_foreground_depth_distance(
+                prediction,
+                foreground_images,
+                background_depth_distance,
+            )
         write_status(args.status, 0.88, f"Exporting {args.export_format}")
         export(
             prediction,
@@ -981,6 +1154,8 @@ def reconstruct(args, model_cache=None, cancel_check=None):
                         directory.rmdir()
                     except OSError:
                         pass
+            if reference_depth is not None:
+                write_npz_reference_depth(result_path, reference_depth)
         check_cancelled(cancel_check)
         copy_frames(images, output_directory)
         write_status(