nan
b5a2f9690c
- Added depth.py for MoGe-2 depth sampling and balloon deformation functionality. - Introduced validate_balloon_asset.py to validate balloon shape modifiers and material properties in Blender. - Enhanced test_server.py with new unit tests for balloon depth analysis and job progress tracking. - Updated existing tests to improve code clarity and maintainability.
369 lines
10 KiB
Python
369 lines
10 KiB
Python
import functools
|
|
import platform
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
|
|
import bpy
|
|
|
|
from .runtime import runtime
|
|
from . import model_catalog as shared_model_catalog
|
|
|
|
|
|
BEN2_MODEL_KEY = "BEN2_BASE"
|
|
DEFAULT_DA3_MODEL_KEY = "BASE"
|
|
DEFAULT_MOGE2_MODEL_KEY = "MOGE2_VITS_NORMAL"
|
|
REQUIRED_MODEL_KEYS = (DEFAULT_MOGE2_MODEL_KEY, BEN2_MODEL_KEY)
|
|
|
|
|
|
DEVICE_LABELS = {
|
|
"cuda": ("CUDA", "NVIDIA CUDA"),
|
|
"directml": ("DirectML", "Windows GPU (DirectML)"),
|
|
"mps": ("MPS", "Apple Metal (MPS)"),
|
|
"cpu": ("CPU", "CPU"),
|
|
}
|
|
|
|
|
|
APPLE_SILICON_PATTERN = re.compile(r"\bApple M(\d+)\b", re.IGNORECASE)
|
|
|
|
|
|
def parse_apple_silicon_generation(brand):
|
|
match = APPLE_SILICON_PATTERN.search(str(brand))
|
|
return int(match.group(1)) if match else None
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def apple_silicon_generation():
|
|
if sys.platform != "darwin" or platform.machine().lower() not in {
|
|
"arm64",
|
|
"aarch64",
|
|
}:
|
|
return None
|
|
try:
|
|
result = subprocess.run(
|
|
("sysctl", "-n", "machdep.cpu.brand_string"),
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=2,
|
|
)
|
|
except (OSError, subprocess.SubprocessError):
|
|
return None
|
|
return parse_apple_silicon_generation(result.stdout)
|
|
|
|
|
|
def supports_mps(generation=None):
|
|
if generation is None:
|
|
generation = apple_silicon_generation()
|
|
return generation is not None and generation >= 4
|
|
|
|
|
|
def supported_devices():
|
|
devices = runtime.environment_manifest().get("devices")
|
|
if isinstance(devices, list) and devices:
|
|
names = [str(device).lower() for device in devices]
|
|
if not supports_mps():
|
|
names = [name for name in names if name != "mps"]
|
|
return names or ["cpu"]
|
|
if sys.platform == "win32":
|
|
return ["directml", "cpu"]
|
|
return ["mps", "cpu"] if supports_mps() else ["cpu"]
|
|
|
|
|
|
def device_enum_items(_owner, _context):
|
|
return [
|
|
(name.upper(), *DEVICE_LABELS[name], index)
|
|
for index, name in enumerate(supported_devices())
|
|
]
|
|
|
|
|
|
def models_status(refresh=False):
|
|
status = runtime.server_status().get("resources", {}).get("model_manager", {})
|
|
return {
|
|
"catalog": tuple(
|
|
shared_model_catalog.model_record(
|
|
model,
|
|
_model_files_ready(model),
|
|
)
|
|
for model in shared_model_catalog.DOWNLOADABLE_MODELS.values()
|
|
),
|
|
"loaded": status.get("loaded", {}),
|
|
}
|
|
|
|
|
|
def _model_files_ready(model):
|
|
directory = runtime.storage_root() / "models" / model.directory_name
|
|
return all(any(directory.glob(pattern)) for pattern in model.ready_patterns)
|
|
|
|
|
|
def model_ready(model_key):
|
|
return any(
|
|
model.get("key") == model_key and model.get("ready")
|
|
for model in models_status().get("catalog", ())
|
|
)
|
|
|
|
|
|
def required_models_ready():
|
|
return all(model_ready(model_key) for model_key in REQUIRED_MODEL_KEYS)
|
|
|
|
|
|
def model_catalog():
|
|
return models_status().get("catalog", ())
|
|
|
|
|
|
def model_record(model_key):
|
|
return next(
|
|
(model for model in model_catalog() if model.get("key") == model_key),
|
|
None,
|
|
)
|
|
|
|
|
|
def _model_items(models):
|
|
return [
|
|
(
|
|
model["key"],
|
|
model["label"] if model.get("ready") else f"{model['label']} (Missing)",
|
|
f"{model['description']}; license: {model['license']}",
|
|
index,
|
|
)
|
|
for index, model in enumerate(models)
|
|
]
|
|
|
|
|
|
def model_enum_items(_owner, _context):
|
|
return _model_items(
|
|
model
|
|
for model in model_catalog()
|
|
if model.get("family") == "da3" and model.get("enabled")
|
|
)
|
|
|
|
|
|
def moge2_model_enum_items(_owner, _context):
|
|
return _model_items(
|
|
model for model in model_catalog() if model.get("family") == "moge2"
|
|
)
|
|
|
|
|
|
def balloon_depth_model_enum_items(_owner, _context):
|
|
models = [
|
|
model
|
|
for model in model_catalog()
|
|
if model.get("family") == "moge2" and model.get("enabled")
|
|
]
|
|
models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
|
|
return _model_items(models)
|
|
|
|
def depth_plane_model_enum_items(_owner, _context):
|
|
models = [
|
|
model
|
|
for model in model_catalog()
|
|
if model.get("family") in {"moge2", "da3"} and model.get("enabled")
|
|
]
|
|
models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
|
|
items = _model_items(models)
|
|
items.append(
|
|
("ORIGINAL", "Original", "Use the original image as depth", len(items))
|
|
)
|
|
return items
|
|
|
|
|
|
def surface_model_enum_items(_owner, _context):
|
|
models = [
|
|
model
|
|
for model in model_catalog()
|
|
if model.get("family") in {"moge2", "da3"} and model.get("enabled")
|
|
]
|
|
models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
|
|
return _model_items(models)
|
|
|
|
|
|
BALLOON_MODE_ITEMS = (
|
|
(
|
|
"POISSON",
|
|
"Poisson",
|
|
"Use a global Poisson field for smooth inflation",
|
|
),
|
|
(
|
|
"TEDDY",
|
|
"Teddy / Chordal Axis",
|
|
"Use a Teddy-style chordal axis and local oval sections",
|
|
),
|
|
)
|
|
|
|
|
|
def balloon_mode_property():
|
|
return bpy.props.EnumProperty(
|
|
name="Mode",
|
|
description="Algorithm used to inflate the saved lasso region",
|
|
items=BALLOON_MODE_ITEMS,
|
|
default="POISSON",
|
|
)
|
|
|
|
|
|
def balloon_thickness_property():
|
|
return bpy.props.FloatProperty(
|
|
name="Thickness",
|
|
description="Scale the Balloon thickness while keeping the lasso region",
|
|
default=1.0,
|
|
min=0.05,
|
|
soft_max=5.0,
|
|
)
|
|
|
|
|
|
def balloon_resolution_property():
|
|
return bpy.props.FloatProperty(
|
|
name="Resolution (m)",
|
|
description=(
|
|
"Distance in meters between sampled mesh points; smaller values "
|
|
"create more geometry"
|
|
),
|
|
default=0.1,
|
|
min=0.005,
|
|
soft_min=0.02,
|
|
soft_max=0.5,
|
|
max=10.0,
|
|
)
|
|
|
|
|
|
def balloon_reuse_material_property():
|
|
return bpy.props.BoolProperty(
|
|
name="Reuse Material",
|
|
description="Reuse a compatible Balloon material made from the same image",
|
|
default=True,
|
|
)
|
|
|
|
|
|
def balloon_contour_analysis_property():
|
|
return bpy.props.BoolProperty(
|
|
name="Contour",
|
|
description=(
|
|
"Use BEN2 Alpha above 0.5 to analyze the lasso contour before "
|
|
"building the Balloon"
|
|
),
|
|
default=False,
|
|
)
|
|
|
|
|
|
def balloon_depth_analysis_property():
|
|
return bpy.props.BoolProperty(
|
|
name="Depth",
|
|
description=("Use MoGe-2 to add adjustable image depth to the Balloon mesh"),
|
|
default=False,
|
|
)
|
|
|
|
|
|
def balloon_depth_model_property():
|
|
return bpy.props.EnumProperty(
|
|
name="Depth Model",
|
|
description="MoGe-2 model used for Balloon depth",
|
|
items=balloon_depth_model_enum_items,
|
|
default=0,
|
|
)
|
|
|
|
def balloon_radial_zoom_property(scale):
|
|
return bpy.props.FloatVectorProperty(
|
|
size=2,
|
|
default=(scale, scale),
|
|
min=0.01,
|
|
options={"HIDDEN", "SKIP_SAVE"},
|
|
)
|
|
|
|
|
|
class SnapFormSettings(bpy.types.PropertyGroup):
|
|
balloon_mode: balloon_mode_property()
|
|
balloon_thickness: balloon_thickness_property()
|
|
balloon_resolution: balloon_resolution_property()
|
|
balloon_reuse_material: balloon_reuse_material_property()
|
|
balloon_contour_analysis: balloon_contour_analysis_property()
|
|
balloon_depth_analysis: balloon_depth_analysis_property()
|
|
balloon_depth_model: balloon_depth_model_property()
|
|
balloon_thickness_radial_zoom: balloon_radial_zoom_property(100.0)
|
|
balloon_resolution_radial_zoom: balloon_radial_zoom_property(1000.0)
|
|
input_path: bpy.props.StringProperty(
|
|
name="Image or Video",
|
|
description="A single image or video file used for reconstruction",
|
|
subtype="FILE_PATH",
|
|
)
|
|
background_input_path: bpy.props.StringProperty(
|
|
name="Image",
|
|
description="A single image used to test BEN2 background removal",
|
|
subtype="FILE_PATH",
|
|
)
|
|
video_frames: bpy.props.IntProperty(
|
|
name="Extract Frames",
|
|
description="Maximum number of consecutive frames extracted from the start of a video or frame sequence",
|
|
default=16,
|
|
min=1,
|
|
soft_max=64,
|
|
)
|
|
model: bpy.props.EnumProperty(
|
|
name="Model",
|
|
description="Depth Anything 3 model",
|
|
items=model_enum_items,
|
|
)
|
|
device: bpy.props.EnumProperty(
|
|
name="Debug Device",
|
|
description=(
|
|
"Inference device used by Debug tasks only; production tasks use "
|
|
"the Device setting in the Add-on Preferences"
|
|
),
|
|
items=device_enum_items,
|
|
default=0,
|
|
)
|
|
process_res: bpy.props.IntProperty(
|
|
name="Process Res",
|
|
description="DA3 base processing resolution",
|
|
default=504,
|
|
min=1,
|
|
soft_max=2048,
|
|
)
|
|
process_res_method: bpy.props.EnumProperty(
|
|
name="Process Res Method",
|
|
description="DA3 image resize method",
|
|
items=(
|
|
(
|
|
"upper_bound_resize",
|
|
"Upper Bound Resize",
|
|
"Make Process Res the longer side",
|
|
),
|
|
(
|
|
"lower_bound_resize",
|
|
"Lower Bound Resize",
|
|
"Make Process Res the shorter side",
|
|
),
|
|
),
|
|
default="upper_bound_resize",
|
|
)
|
|
moge_resolution_level: bpy.props.IntProperty(
|
|
name="Resolution Level",
|
|
description="MoGe-2 inference resolution level; higher is finer and slower",
|
|
default=5,
|
|
min=0,
|
|
max=9,
|
|
)
|
|
moge_model: bpy.props.EnumProperty(
|
|
name="Model",
|
|
description="MoGe-2 model used for depth estimation",
|
|
items=moge2_model_enum_items,
|
|
default=0,
|
|
)
|
|
ref_view_strategy: bpy.props.EnumProperty(
|
|
name="Ref View Strategy",
|
|
description="DA3 reference view selection strategy",
|
|
items=(
|
|
("first", "First", "Always use the first view"),
|
|
("middle", "Middle", "Use the middle view"),
|
|
(
|
|
"saddle_balanced",
|
|
"Saddle Balanced",
|
|
"Select a view with balanced feature metrics",
|
|
),
|
|
(
|
|
"saddle_sim_range",
|
|
"Saddle Sim Range",
|
|
"Select the view with the largest similarity range",
|
|
),
|
|
),
|
|
default="saddle_balanced",
|
|
)
|