diff --git a/README.md b/README.md index 70c421d..dd516fb 100644 --- a/README.md +++ b/README.md @@ -19,8 +19,8 @@ - 使用一个导出类型下拉菜单和一个 `Export` 按钮 - 只显示当前格式的专属参数；GLB 与 PLY 设置不会同时出现 - Debug 的模型选项按格式过滤；PLY 只提供两个带 Gaussian head 的 Giant 模型 -- Debug 在模型选项同一行用勾选或警告图标表示下载状态，不显示状态文字、License 和 - 已完成任务的文件计数 +- Debug 与 Convert to Depth Plane 的模型选项直接显示下载状态；未下载模型带 + `(Missing)`，已下载模型保持原名 - 分别调用 DA3 的 `glb`、`gs_ply`、`depth_vis` 和 `mini_npz` exporter - Image Empty 右键菜单可直接运行 DA3，根据 `mini_npz` 深度生成面片 - 使用浮点深度贴图、Simple Subdivision、Displace Modifier 和原图 Bump @@ -29,6 +29,7 @@ - Debug 导出完成后只打开对应 job 文件夹，不向当前 Blender 文件添加任何数据 - 完整保留每次任务生成的所有官方文件 - 安装、模型下载和推理均在后台子进程运行，状态栏实时显示阶段、百分比进度条与取消按钮 +- 后台任务运行时，Install、Download 和 Export 按钮保持原位置并变灰，任务完成后恢复 ## 使用 @@ -44,21 +45,27 @@ ## Image Empty 转深度面片 1. 在 3D View 选中 `Object Type = Empty`、`Display As = Image` 的对象。 -2. 右键选择 `Convert to Depth Plane`，在对话框中选择 Model、Use Ray Pose、Process Resolution、Subdivision， - 以及是否删除原 Image Empty。默认分辨率为 `1024`、Subdivision 为 `5`，并默认删除原对象。 - Model 可选择 `DA3 Small`、`DA3 Base`、`DA3 Large 1.1`、`DA3 Giant 1.1` 或 - `DA3 Nested Giant Large 1.1`，默认 `DA3 Base`；Use Ray Pose - 默认开启。转换导出 `mini_npz`，使用前需要在 `Setup` 下载所选模型，并固定使用 - `process_res_method="lower_bound_resize"`；设备和参考视图策略继续使用场景设置。 -3. 推理完成后会自动创建面片。面片几何先由 `mini_npz` 的相机内参与深度确定，再通过对象 - Transform 匹配原 Image Empty 的可见尺寸、位置、旋转和 Offset。XY 拟合得到的统一倍率 - 同时应用于 Z 轴，使面片厚度与宽高按同一比例变化。开启 `Delete Original` 时，成功生成后 +2. 右键选择 `Convert to Depth Plane`，通过 `Depth from` 选择深度来源： + - `AI Generated`：显示 Model、Use Ray Pose 和 Resolution Scale。实际 Process Res 为图片长边像素数 + 乘以 Resolution Scale；范围为 `0.1–1.0`，默认 `0.5`。计算后的 Process Res 会直接显示在 + `Resolution Scale (… px)` 标签中。 + 转换导出 `mini_npz`，使用前需要在 + `Setup` 下载所选模型，并固定使用 `process_res_method="upper_bound_resize"`；设备和参考视图策略 + 继续使用场景设置。 + - `Original`：直接把原图作为深度贴图，不显示估算参数，也不启动 Runtime 或要求源图片存在于磁盘。 + `Original` 是默认深度来源。选择 `AI Generated` 时，其 Model、Use Ray Pose 和 Resolution Scale + 参数会集中显示在独立边框中。两种模式始终显示 Displace via、Subdivision 与 Delete Original。 + Displace via 可选 `Modifier`（默认）或 `Material`； + Subdivision 默认为 `5`，并默认删除原对象。 +3. AI Generated 完成后会自动创建面片。面片几何先由 `mini_npz` 的相机内参与深度确定，再通过对象 + Transform 匹配原 Image Empty 的可见尺寸、位置、旋转和 Offset。Original 模式直接使用源 Empty + 的可见尺寸、Offset 与 Transform。开启 `Delete Original` 时，成功生成后 删除源 Empty。生成面片 Object 使用源 Empty 的原名；关闭 `Delete Original` 时不处理 同名冲突，由 Blender 自动为生成面片追加编号。 4. 短边分段数为 `2^Subdivision`，长边分段数根据图像宽高比取最接近的整数。例如高:宽为 `2:3`、Subdivision 为 `1` 时，直接生成 `2×3` 个面；Subdivision 为 `3` 时生成 `8×12` 个面。随后的 Simple Subdivision 修改器使用原生名称，Viewport Levels 固定为 - `1`，Render Levels 固定为 `2`，不再随 Process Resolution 或 Subdivision 动态变化。 + `1`，Render Levels 固定为 `2`，不再随 Resolution Scale 或 Subdivision 动态变化。 所有任务使用统一的状态栏进度显示。变化中的阶段说明显示在右侧的紧凑进度条内部，取消按钮 位于进度条右侧。Image Depth 的阶段包括 Runtime 启动、输入准备、模型加载、深度估算、缓存 @@ -86,12 +93,18 @@ Height = sqrt(Area / Aspect) `snapform_camera_plane_height`，便于检查实际采用的 参考深度和尺寸。 -Displace Modifier 的深度贴图使用 `Extension = Extend`。材质不再使用生成的深度图执行 -Displacement，其 `Strength` 固定为 `-1`、`Midlevel` 固定为 `1`；原始彩色贴图同时连接 +选择 `Modifier` 时，AI Generated 的 Displace Modifier 深度贴图使用 `Extension = Extend`， +`Strength` 固定为 `-1`、`Midlevel` 固定为 `1`。Original 把原图指定给 Displace Texture，并设置 +`Texture Coordinates = UV`、`Midlevel = 0`；其余修改器设置保留 Blender 默认值。 +选择 `Material` 时，不创建 Displace Modifier；深度 Image Texture 的 Color 连接到 Displacement +节点的 Height，再连接 Material Output 的 Displacement，并把材质 Surface Displacement 设置为 +`Displacement Only`。AI Generated 使用 `Midlevel = 1`、`Scale = -1`；Original 使用 +`Midlevel = 0`、`Scale = 1`，与对应修改器值一致。无论选择哪种呈现方式，原始彩色贴图都会连接 Principled BSDF 的 `Base Color`，并通过 `Bump.Height → Bump.Normal → Principled.Normal` 提供表面细节。Bump 的基础 Distance 为 `0.001`，会乘以面片 Transform 的统一缩放倍率。原图 Color 还会经过 `Float Curve` 连接到 -Principled BSDF 的 `Roughness`，曲线两端输出分别为 `0.5` 和 `0.1`；材质 IOR 默认为 `1.2`。 +Principled BSDF 的 `Roughness`，曲线两端输出分别为 `0.5` 和 `0.1`；材质 IOR 保持 +Principled BSDF 默认值 `1.5`。 这里的“原始距离”是 DA3 模型输出的尺度。普通 Small、Base、Large 和 Giant 模型的深度不应 被解释为现实世界的米；需要现实尺度时应选择带 metric depth 的 @@ -99,15 +112,18 @@ Principled BSDF 的 `Roughness`，曲线两端输出分别为 `0.5` 和 `0.1`； 相机内参逐像素反投影得到的完整 camera-space 网格；中位深度截面为现有 Displace Modifier 工作流提供正确的整体高宽尺度。 -生成完成后会直接创建面片，不提供缓存重建入口或单独的导入参数。 -重新选择 `Convert to Depth Plane` 会再次运行 DA3。 +生成完成后会直接创建面片，不提供缓存重建入口或单独的导入参数。重新选择 +`Convert to Depth Plane` 会按当前 Depth from 再次转换。 面片、深度 Image、Texture 与 Material 会在同一个 Blender Undo 步骤中创建；转换完成后执行 `Ctrl+Z` 会完整撤销本次转换，不会越过转换去回退更早的 Image Empty 操作。 -当前转换入口要求 Image Empty 引用仍然存在的本地图像文件。DA3 的原始 +Image Empty 使用 Packed Image 时，SnapForm 会通过 Blender 写出临时 PNG 交给 DA3， +不会 Unpack 或修改源 Image；任务结束后自动删除临时文件。普通图片仍直接读取源文件。DA3 的原始 `results.npz` 保留在对应 job 的 `exports/mini_npz/` 中，生成的浮点深度 Image 会 Pack 到 -当前 `.blend`。源 Empty 名为 `img` 时，深度 Image 与 Displace Texture 都命名为 -`img.depth`。Pack 使用 OpenEXR，因而不会把大于 `1.0` 的原始深度截断。 +当前 `.blend`。Depth Image 与 Displace Texture 根据源 Image datablock 命名。例如 Empty Object +名为 `obj_img`、引用的 Image 名为 `img.png` 时，Depth Image 与 Texture 都命名为 `img.depth`； +生成的 Object、Mesh 与 Material 仍命名为 `obj_img`。Pack 使用 OpenEXR，因而不会把大于 `1.0` +的原始深度截断。已有同名 datablock 时由 Blender 自动追加 `.001` 等编号。 选择 PLY 后，模型菜单只会提供 `DA3 Giant 1.1` 和 `DA3 Nested Giant Large 1.1`。Small、Base 和 Large 没有 Gaussian head，无法由 @@ -145,7 +161,7 @@ CC BY-NC 4.0，仅限非商业用途。 `gs_views_interval` 仅在选择 `PLY (3DGS)` 时显示： - `process_res=1024` -- `process_res_method="lower_bound_resize"` +- `process_res_method="upper_bound_resize"` - `use_ray_pose=False` - `ref_view_strategy="saddle_balanced"` - `conf_thresh_percentile=40.0` diff --git a/pack.py b/pack.py index 4725cd6..f60ad12 100644 --- a/pack.py +++ b/pack.py @@ -12,10 +12,8 @@ def build_extension(): manifest = tomllib.loads( (SOURCE_DIR / "blender_manifest.toml").read_text(encoding="utf-8") ) - blender_version = manifest["blender_version_min"].split(".") - blender_tag = f"b{blender_version[0]}{blender_version[1]}" output_path = DIST_DIR / ( - f"{manifest['name']}.v{manifest['version']}.{blender_tag}.zip" + f"{manifest['name']}.v{manifest['version']}.zip" ) DIST_DIR.mkdir(exist_ok=True) diff --git a/src/snapform/image_depth.py b/src/snapform/image_depth.py index 290d9fb..586645a 100644 --- a/src/snapform/image_depth.py +++ b/src/snapform/image_depth.py @@ -5,36 +5,29 @@ import bpy DEPTH_KEY_CANDIDATES = ("depth", "depths") INTRINSICS_KEY_CANDIDATES = ("intrinsics", "ixts") -IMAGE_DEPTH_MODEL_ITEMS = ( - ( - "SMALL", - "DA3 Small", - "Fastest model with Apache 2.0 weights", - ), - ( - "BASE", - "DA3 Base", - "Balanced model with Apache 2.0 weights", - ), - ( - "LARGE_1_1", - "DA3 Large 1.1", - "Higher quality model with CC BY-NC 4.0 weights", - ), - ( - "GIANT_1_1", - "DA3 Giant 1.1", - "Giant model with Gaussian head and CC BY-NC 4.0 weights", - ), - ( - "NESTED_GIANT_LARGE_1_1", - "DA3 Nested Giant Large 1.1", - "Nested Giant model with metric depth and CC BY-NC 4.0 weights", - ), -) DISPLACEMENT_STRENGTH = -1.0 DISPLACEMENT_MIDLEVEL = 1.0 +ORIGINAL_DISPLACEMENT_STRENGTH = 1.0 +ORIGINAL_DISPLACEMENT_MIDLEVEL = 0.0 BUMP_BASE_DISTANCE = 0.001 +IMAGE_NAME_SUFFIXES = ( + ".jpeg", + ".jpg", + ".png", + ".webp", + ".tiff", + ".tif", + ".exr", + ".hdr", + ".bmp", + ".tga", +) + + +def image_depth_model_items(owner, context): + from .models import model_enum_items + + return model_enum_items(owner, context) def is_image_empty(candidate): @@ -58,6 +51,60 @@ def image_path(image): return Path(bpy.path.abspath(filepath)).expanduser() +def image_is_packed(image): + if getattr(image, "packed_file", None) is not None: + return True + return bool(getattr(image, "packed_files", ())) + + +def image_base_name(image): + name = image.name + lower_name = name.lower() + for suffix in IMAGE_NAME_SUFFIXES: + suffix_index = lower_name.rfind(suffix) + if suffix_index <= 0: + continue + trailing = name[suffix_index + len(suffix) :] + if trailing and not ( + trailing.startswith(".") and trailing[1:].isdigit() + ): + continue + return f"{name[:suffix_index]}{trailing}" + return name + + +def depth_data_name(source_object): + return f"{image_base_name(source_object.data)}.depth" + + +def scaled_process_resolution(image, scale): + width, height = image.size + if width <= 0 or height <= 0: + raise ValueError("The source image has no readable dimensions") + if scale < 0.1 or scale > 1.0: + raise ValueError("Resolution Scale must be between 0.1 and 1.0") + return int(max(width, height) * scale + 0.5) + + +def save_packed_image(image, scene): + from .runtime import temporary_input_path + + output_path = temporary_input_path() + image_settings = scene.render.image_settings + original_file_format = image_settings.file_format + try: + image_settings.file_format = "PNG" + image.save_render(str(output_path), scene=scene) + except (OSError, RuntimeError) as error: + output_path.unlink(missing_ok=True) + raise RuntimeError("Unable to export the packed image") from error + finally: + image_settings.file_format = original_file_format + if not output_path.is_file(): + raise RuntimeError("Blender did not export the packed image") + return output_path + + def plane_bounds(source_object): image_width, image_height = source_object.data.size if image_width <= 0 or image_height <= 0: @@ -352,7 +399,7 @@ def create_plane_mesh(source_object, bounds=None, subdivision_level=5): ) ) - mesh = bpy.data.meshes.new(f"{source_object.name}.DA3Depth") + mesh = bpy.data.meshes.new(source_object.name) mesh.from_pydata(vertices, (), faces) mesh.update() for polygon in mesh.polygons: @@ -369,6 +416,9 @@ def create_material( name, color_image, depth_scale, + displacement_image=None, + displacement_midlevel=0.0, + displacement_scale=1.0, ): material = bpy.data.materials.new(name) material.use_nodes = True @@ -387,8 +437,6 @@ def create_material( color_texture.interpolation = "Linear" color_texture.extension = "EXTEND" bump.inputs["Distance"].default_value = BUMP_BASE_DISTANCE * depth_scale - shader.inputs["IOR"].default_value = 1.2 - roughness_curve.mapping.initialize() curve_points = roughness_curve.mapping.curves[0].points curve_points[0].location = (0.0, 0.5) @@ -415,15 +463,118 @@ def create_material( shader.inputs["Roughness"], ) links.new(shader.outputs["BSDF"], output.inputs["Surface"]) + + if displacement_image is not None: + depth_texture = nodes.new("ShaderNodeTexImage") + displacement = nodes.new("ShaderNodeDisplacement") + depth_texture.image = displacement_image + depth_texture.interpolation = "Linear" + depth_texture.extension = "EXTEND" + displacement.inputs["Midlevel"].default_value = ( + displacement_midlevel + ) + displacement.inputs["Scale"].default_value = displacement_scale + depth_texture.location = (-300.0, -360.0) + displacement.location = (80.0, -260.0) + links.new( + depth_texture.outputs["Color"], + displacement.inputs["Height"], + ) + links.new( + displacement.outputs["Displacement"], + output.inputs["Displacement"], + ) + set_displacement_only(material) return material +def set_displacement_only(material): + for owner, property_name in ( + (material, "surface_displacement_method"), + (material, "displacement_method"), + (getattr(material, "cycles", None), "displacement_method"), + ): + if owner is not None and hasattr(owner, property_name): + setattr(owner, property_name, "DISPLACEMENT") + return + + +def finish_depth_plane( + context, + source_object, + plane_object, + delete_source, +): + source_name = source_object.name + for selected in context.selected_objects: + selected.select_set(False) + plane_object.select_set(True) + context.view_layer.objects.active = plane_object + if delete_source: + bpy.data.objects.remove(source_object, do_unlink=True) + plane_object.name = source_name + return plane_object + + +def build_image_depth_plane( + context, + source_object, + subdivision_level, + delete_source, + displacement_via="MODIFIER", +): + source_name = source_object.name + mesh = create_plane_mesh( + source_object, + subdivision_level=subdivision_level, + ) + plane_object = bpy.data.objects.new(source_name, mesh) + context.collection.objects.link(plane_object) + plane_object.matrix_world = source_object.matrix_world.copy() + + subdivision = plane_object.modifiers.new("Subdivision", "SUBSURF") + subdivision.subdivision_type = "SIMPLE" + subdivision.levels = 1 + subdivision.render_levels = 2 + + if displacement_via == "MODIFIER": + texture = bpy.data.textures.new( + depth_data_name(source_object), + type="IMAGE", + ) + texture.image = source_object.data + displace = plane_object.modifiers.new("Displace", "DISPLACE") + displace.texture = texture + displace.texture_coords = "UV" + displace.mid_level = ORIGINAL_DISPLACEMENT_MIDLEVEL + + depth_scale = abs(float(plane_object.scale.z)) + material = create_material( + source_name, + source_object.data, + depth_scale, + displacement_image=( + source_object.data if displacement_via == "MATERIAL" else None + ), + displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL, + displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH, + ) + mesh.materials.append(material) + return finish_depth_plane( + context, + source_object, + plane_object, + delete_source, + ) + + def build_depth_plane( context, source_object, npz_path, subdivision_level, delete_source, + displacement_via="MODIFIER", ): source_name = source_object.name depth = extract_depth_array(npz_path) @@ -435,7 +586,7 @@ def build_depth_plane( intrinsics, DISPLACEMENT_STRENGTH, ) - depth_name = f"{source_name}.depth" + depth_name = depth_data_name(source_object) depth_image = create_depth_image( depth_name, depth_values, @@ -457,23 +608,29 @@ def build_depth_plane( subdivision.levels = 1 subdivision.render_levels = 2 - texture = bpy.data.textures.new( - depth_name, - type="IMAGE", - ) - texture.image = depth_image - texture.extension = "EXTEND" - displace = plane_object.modifiers.new("Displace", "DISPLACE") - displace.texture = texture - displace.texture_coords = "UV" - displace.direction = "NORMAL" - displace.strength = DISPLACEMENT_STRENGTH - displace.mid_level = DISPLACEMENT_MIDLEVEL + if displacement_via == "MODIFIER": + texture = bpy.data.textures.new( + depth_name, + type="IMAGE", + ) + texture.image = depth_image + texture.extension = "EXTEND" + displace = plane_object.modifiers.new("Displace", "DISPLACE") + displace.texture = texture + displace.texture_coords = "UV" + displace.direction = "NORMAL" + displace.strength = DISPLACEMENT_STRENGTH + displace.mid_level = DISPLACEMENT_MIDLEVEL material = create_material( - f"{source_name}.DA3DepthMaterial", + source_name, source_object.data, depth_scale, + displacement_image=( + depth_image if displacement_via == "MATERIAL" else None + ), + displacement_midlevel=DISPLACEMENT_MIDLEVEL, + displacement_scale=DISPLACEMENT_STRENGTH, ) mesh.materials.append(material) plane_object["snapform_reference_depth"] = reference_depth @@ -481,14 +638,12 @@ def build_depth_plane( plane_object["snapform_camera_plane_height"] = bounds[3] - bounds[2] plane_object["snapform_depth_scale"] = depth_scale - for selected in context.selected_objects: - selected.select_set(False) - plane_object.select_set(True) - context.view_layer.objects.active = plane_object - if delete_source: - bpy.data.objects.remove(source_object, do_unlink=True) - plane_object.name = source_name - return plane_object + return finish_depth_plane( + context, + source_object, + plane_object, + delete_source, + ) class BuildDepthPlane(bpy.types.Operator): @@ -512,6 +667,14 @@ class BuildDepthPlane(bpy.types.Operator): default=True, options={"HIDDEN", "SKIP_SAVE"}, ) + snapform_displacement_via: bpy.props.EnumProperty( + items=( + ("MODIFIER", "Modifier", "Use a Displace modifier"), + ("MATERIAL", "Material", "Use material displacement"), + ), + default="MODIFIER", + options={"HIDDEN", "SKIP_SAVE"}, + ) def execute(self, context): source_object = bpy.data.objects.get(self.snapform_source_object_name) @@ -526,7 +689,7 @@ class BuildDepthPlane(bpy.types.Operator): if not npz_path.is_file(): self.report( {"ERROR"}, - f"DA3 Mini NPZ does not exist: {npz_path}", + f"Depth result does not exist: {npz_path}", ) return {"CANCELLED"} @@ -537,6 +700,7 @@ class BuildDepthPlane(bpy.types.Operator): npz_path, self.snapform_subdivision_level, self.snapform_delete_source, + self.snapform_displacement_via, ) except (OSError, RuntimeError, ValueError) as error: self.report({"ERROR"}, str(error)) @@ -547,27 +711,49 @@ class BuildDepthPlane(bpy.types.Operator): class ConvertImageToDepthPlane(bpy.types.Operator): bl_idname = "snapform.convert_image_to_depth_plane" bl_label = "Convert to Depth Plane" - bl_description = "Run DA3 and convert the image to a displaced depth plane" + bl_description = "Convert an image to a displaced depth plane" bl_options = {"REGISTER"} + snapform_depth_source: bpy.props.EnumProperty( + name="Depth from", + description="Source used to displace the plane", + items=( + ( + "ESTIMATE", + "AI Generated", + "Generate depth with an AI model", + ), + ("ORIGINAL", "Original", "Use the original image as depth"), + ), + default="ORIGINAL", + ) + snapform_displacement_via: bpy.props.EnumProperty( + name="Displace via", + description="How the depth displacement is presented", + items=( + ("MODIFIER", "Modifier", "Use a Displace modifier"), + ("MATERIAL", "Material", "Use material displacement"), + ), + default="MODIFIER", + ) snapform_model: bpy.props.EnumProperty( name="Model", - description="DA3 model used to estimate image depth", - items=IMAGE_DEPTH_MODEL_ITEMS, - default="BASE", + description="Model used to estimate image depth", + items=image_depth_model_items, ) snapform_use_ray_pose: bpy.props.BoolProperty( name="Use Ray Pose", description="Use ray-based camera pose estimation", default=True, ) - snapform_process_res: bpy.props.IntProperty( - name="Process Resolution", - description="DA3 processing resolution", - default=1024, - min=14, - max=2016, - step=14, + snapform_process_res_factor: bpy.props.FloatProperty( + name="Resolution Scale", + description="Scale applied to the source image long-edge resolution", + default=0.5, + min=0.1, + max=1.0, + step=10, + precision=2, ) snapform_delete_original: bpy.props.BoolProperty( name="Delete Original", @@ -589,13 +775,27 @@ class ConvertImageToDepthPlane(bpy.types.Operator): return is_image_empty(context.object) and active_job() is None def invoke(self, context, _event): + if not self.properties.is_property_set("snapform_model"): + self.snapform_model = "BASE" return context.window_manager.invoke_props_dialog(self, width=340) - def draw(self, _context): + def draw(self, context): layout = self.layout - layout.prop(self, "snapform_model") - layout.prop(self, "snapform_use_ray_pose") - layout.prop(self, "snapform_process_res") + layout.prop(self, "snapform_depth_source") + if self.snapform_depth_source == "ESTIMATE": + estimate_box = layout.box() + estimate_box.prop(self, "snapform_model") + estimate_box.prop(self, "snapform_use_ray_pose") + process_res = scaled_process_resolution( + context.object.data, + self.snapform_process_res_factor, + ) + estimate_box.prop( + self, + "snapform_process_res_factor", + text=f"Resolution Scale ({process_res} px)", + ) + layout.prop(self, "snapform_displacement_via") layout.prop(self, "snapform_subdivision_level") layout.prop(self, "snapform_delete_original") @@ -603,7 +803,47 @@ class ConvertImageToDepthPlane(bpy.types.Operator): from .operators import IMAGE_DEPTH_JOB source_object = context.object - source_path = image_path(source_object.data) + if self.snapform_depth_source == "ORIGINAL": + try: + build_image_depth_plane( + context, + source_object, + self.snapform_subdivision_level, + self.snapform_delete_original, + self.snapform_displacement_via, + ) + except (OSError, RuntimeError, ValueError) as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + undo_result = bpy.ops.ed.undo_push( + message="Convert to Depth Plane", + ) + if "FINISHED" not in undo_result: + self.report( + {"ERROR"}, + "Unable to create the depth plane Undo step", + ) + return {"CANCELLED"} + return {"FINISHED"} + + image = source_object.data + try: + process_res = scaled_process_resolution( + image, + self.snapform_process_res_factor, + ) + except ValueError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + source_path = image_path(image) + delete_input = False + if image_is_packed(image): + try: + source_path = save_packed_image(image, context.scene) + except RuntimeError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + delete_input = True if not source_path.is_file(): self.report( {"ERROR"}, @@ -617,11 +857,15 @@ class ConvertImageToDepthPlane(bpy.types.Operator): snapform_source_object_name=source_object.name, snapform_model_key=self.snapform_model, snapform_use_ray_pose=self.snapform_use_ray_pose, - snapform_process_res=self.snapform_process_res, + snapform_process_res=process_res, snapform_subdivision_level=self.snapform_subdivision_level, snapform_delete_source=self.snapform_delete_original, + snapform_delete_input=delete_input, + snapform_displacement_via=self.snapform_displacement_via, ) if "CANCELLED" in result: + if delete_input: + source_path.unlink(missing_ok=True) return {"CANCELLED"} return {"FINISHED"} diff --git a/src/snapform/models.py b/src/snapform/models.py index 8e86095..a4d6d72 100644 --- a/src/snapform/models.py +++ b/src/snapform/models.py @@ -53,18 +53,62 @@ MODELS = { supports_gs=True, ), } +MODEL_NUMBERS = { + model_key: number for number, model_key in enumerate(MODELS) +} -def enum_items(supports_gs=None): - return [ - ( - model.key, - model.label, - f"{model.description}; license: {model.license_name}", - ) +_ENUM_ITEMS_CACHE = {} + + +def enum_items(supports_gs=None, show_download_status=False): + model_ready = None + if show_download_status: + from .runtime import model_ready + + models = [ + model for model in MODELS.values() if supports_gs is None or model.supports_gs == supports_gs ] + download_status = ( + tuple(model_ready(model.key) for model in models) + if model_ready is not None + else () + ) + cache_key = ( + supports_gs, + show_download_status, + download_status, + ) + cached_items = _ENUM_ITEMS_CACHE.get(cache_key) + if cached_items is not None: + return cached_items + + items = [] + for index, model in enumerate(models): + label = model.label + if download_status and not download_status[index]: + label = f"{label} (Missing)" + items.append( + ( + model.key, + label, + f"{model.description}; license: {model.license_name}", + MODEL_NUMBERS[model.key], + ) + ) + + _ENUM_ITEMS_CACHE[cache_key] = items + return items + + +def model_enum_items(_owner, _context): + return enum_items(show_download_status=True) + + +def gs_model_enum_items(_owner, _context): + return enum_items(supports_gs=True, show_download_status=True) def selected_model_key(settings, export_format=None): diff --git a/src/snapform/operators.py b/src/snapform/operators.py index 052d709..1d2e870 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -11,7 +11,7 @@ from .models import get_model, selected_model_key RECONSTRUCTION_JOB = "RECONSTRUCT" IMAGE_DEPTH_JOB = "IMAGE_DEPTH" IMAGE_DEPTH_DEFAULT_PROCESS_RES = 1024 -IMAGE_DEPTH_PROCESS_RES_METHOD = "lower_bound_resize" +IMAGE_DEPTH_PROCESS_RES_METHOD = "upper_bound_resize" def open_job_directory(result_directory): @@ -36,6 +36,8 @@ class ActiveJob: process_res: int = IMAGE_DEPTH_DEFAULT_PROCESS_RES subdivision_level: int = 5 delete_source: bool = True + displacement_via: str = "MODIFIER" + temporary_input_path: Path | None = None cancelled: bool = False @@ -62,6 +64,11 @@ def close_active_job(): if _active_job is None: return _active_job.log_file.close() + if _active_job.temporary_input_path is not None: + try: + _active_job.temporary_input_path.unlink(missing_ok=True) + except OSError: + pass _active_job = None @@ -98,20 +105,32 @@ class RunJob(bpy.types.Operator): ) snapform_process_res: bpy.props.IntProperty( default=IMAGE_DEPTH_DEFAULT_PROCESS_RES, - min=14, - max=2016, + min=1, + max=65_536, options={"HIDDEN", "SKIP_SAVE"}, ) snapform_delete_source: bpy.props.BoolProperty( default=True, options={"HIDDEN", "SKIP_SAVE"}, ) + snapform_delete_input: bpy.props.BoolProperty( + default=False, + options={"HIDDEN", "SKIP_SAVE"}, + ) snapform_subdivision_level: bpy.props.IntProperty( default=5, min=0, max=10, options={"HIDDEN", "SKIP_SAVE"}, ) + snapform_displacement_via: bpy.props.EnumProperty( + items=( + ("MODIFIER", "Modifier", ""), + ("MATERIAL", "Material", ""), + ), + default="MODIFIER", + options={"HIDDEN", "SKIP_SAVE"}, + ) def execute(self, context): global _active_job @@ -138,7 +157,7 @@ class RunJob(bpy.types.Operator): ) except OSError as error: log_file.close() - self.report({"ERROR"}, f"Unable to start DA3 task: {error}") + self.report({"ERROR"}, f"Unable to start task: {error}") return {"CANCELLED"} _active_job = ActiveJob( @@ -157,6 +176,13 @@ class RunJob(bpy.types.Operator): process_res=self.snapform_process_res, subdivision_level=self.snapform_subdivision_level, delete_source=self.snapform_delete_source, + displacement_via=self.snapform_displacement_via, + temporary_input_path=( + Path(self.snapform_input_path) + if self.job_type == IMAGE_DEPTH_JOB + and self.snapform_delete_input + else None + ), ) settings.status = self._starting_message(settings) settings.progress = 0.0 @@ -245,9 +271,9 @@ class RunJob(bpy.types.Operator): "DOWNLOAD_MODEL": "Starting model download", } if self.job_type == RECONSTRUCTION_JOB: - return f"Starting DA3 {settings.export_format} export" + return f"Starting {settings.export_format} export" if self.job_type == IMAGE_DEPTH_JOB: - return "Starting DA3 image depth conversion" + return "Starting depth estimation" if self.job_type == "DOWNLOAD_MODEL": return f"Starting {get_model(self.model_key).label} download" return labels[self.job_type] @@ -298,7 +324,7 @@ class RunJob(bpy.types.Operator): settings.progress = 0.0 self.report( {"ERROR"}, - f"{error or 'DA3 subprocess failed'}; log: {job.log_path}", + f"{error or 'Task subprocess failed'}; log: {job.log_path}", ) close_active_job() self._redraw_ui(context, force=True) @@ -326,7 +352,7 @@ class RunJob(bpy.types.Operator): close_active_job() self._redraw_ui(context, force=True) return {"CANCELLED"} - status["message"] = "DA3 depth plane is ready" + status["message"] = "Depth plane is ready" settings.status = status.get("message", "Task complete") settings.progress = 1.0 @@ -346,6 +372,7 @@ class RunJob(bpy.types.Operator): snapform_npz_path=str(npz_path), snapform_subdivision_level=job.subdivision_level, snapform_delete_source=job.delete_source, + snapform_displacement_via=job.displacement_via, ) if "FINISHED" not in result: raise RuntimeError("Unable to build the depth plane") diff --git a/src/snapform/panel.py b/src/snapform/panel.py index 24b8d17..445018e 100644 --- a/src/snapform/panel.py +++ b/src/snapform/panel.py @@ -1,6 +1,6 @@ import bpy -from .models import MODELS, selected_model_key +from .models import MODELS from .operators import active_job from .runtime import model_ready, runtime_ready @@ -18,13 +18,14 @@ class RuntimeModelsPanel(bpy.types.Panel): if not runtime_ready(): layout.label(text="Runtime is not installed", icon="ERROR") - if not running: - operator = layout.operator( - "snapform.run_job", - text="Install Runtime", - icon="IMPORT", - ) - operator.job_type = "INSTALL_RUNTIME" + install_row = layout.row() + install_row.enabled = not running + operator = install_row.operator( + "snapform.run_job", + text="Install Runtime", + icon="IMPORT", + ) + operator.job_type = "INSTALL_RUNTIME" for model in MODELS.values(): self._draw_model(layout, model, running) @@ -34,16 +35,16 @@ class RuntimeModelsPanel(bpy.types.Panel): row.label(text=model.label) if model_ready(model.key): row.label(text="Ready", icon="CHECKMARK") - elif runtime_ready() and not running: - operator = row.operator( + else: + download_row = row.row(align=True) + download_row.enabled = runtime_ready() and not running + operator = download_row.operator( "snapform.run_job", text="Download", icon="IMPORT", ) operator.job_type = "DOWNLOAD_MODEL" operator.model_key = model.key - else: - row.label(text="Not downloaded", icon="ERROR") class DebugPanel(bpy.types.Panel): @@ -59,16 +60,10 @@ class DebugPanel(bpy.types.Panel): settings = context.scene.snapform_settings layout.prop(settings, "export_format") - model_key = selected_model_key(settings) model_property = ( "gs_model" if settings.export_format == "gs_ply" else "model" ) - model_row = layout.row(align=True) - model_row.prop(settings, model_property) - if model_ready(model_key): - model_row.label(text="", icon="CHECKMARK") - else: - model_row.label(text="", icon="ERROR") + layout.prop(settings, model_property) layout.prop(settings, "input_path") input_path = bpy.path.abspath(settings.input_path) @@ -103,10 +98,11 @@ class DebugPanel(bpy.types.Panel): ply_box.label(text="PLY (3DGS) Settings") ply_box.prop(settings, "gs_views_interval") - if active_job() is None: - operator = layout.operator( - "snapform.run_job", - text="Export", - icon="EXPORT", - ) - operator.job_type = "RECONSTRUCT" + export_row = layout.row() + export_row.enabled = active_job() is None + operator = export_row.operator( + "snapform.run_job", + text="Export", + icon="EXPORT", + ) + operator.job_type = "RECONSTRUCT" diff --git a/src/snapform/preferences.py b/src/snapform/preferences.py index 3c2fa23..e3e573d 100644 --- a/src/snapform/preferences.py +++ b/src/snapform/preferences.py @@ -33,6 +33,7 @@ class SnapFormPreferences(bpy.types.AddonPreferences): ) def draw(self, _context): + from .operators import active_job from .runtime import runtime_directory, runtime_ready layout = self.layout @@ -44,7 +45,9 @@ class SnapFormPreferences(bpy.types.AddonPreferences): if runtime_ready(): layout.label(text="Runtime is ready", icon="CHECKMARK") else: - operator = layout.operator( + install_row = layout.row() + install_row.enabled = active_job() is None + operator = install_row.operator( "snapform.run_job", text="Install Runtime", icon="IMPORT", diff --git a/src/snapform/properties.py b/src/snapform/properties.py index a3ec254..bd60a12 100644 --- a/src/snapform/properties.py +++ b/src/snapform/properties.py @@ -1,6 +1,6 @@ import bpy -from .models import enum_items +from .models import gs_model_enum_items, model_enum_items class SnapFormSettings(bpy.types.PropertyGroup): @@ -19,14 +19,12 @@ class SnapFormSettings(bpy.types.PropertyGroup): model: bpy.props.EnumProperty( name="Model", description="Depth Anything 3 model", - items=enum_items(), - default="BASE", + items=model_enum_items, ) gs_model: bpy.props.EnumProperty( name="Model", description="Depth Anything 3 model with a Gaussian head", - items=enum_items(supports_gs=True), - default="GIANT_1_1", + items=gs_model_enum_items, ) device: bpy.props.EnumProperty( name="Device", @@ -60,7 +58,7 @@ class SnapFormSettings(bpy.types.PropertyGroup): "Make Process Res the shorter side", ), ), - default="lower_bound_resize", + default="upper_bound_resize", ) conf_thresh_percentile: bpy.props.FloatProperty( name="Conf Thresh Percentile", diff --git a/src/snapform/runtime.py b/src/snapform/runtime.py index 58130ef..da3219c 100644 --- a/src/snapform/runtime.py +++ b/src/snapform/runtime.py @@ -78,6 +78,12 @@ def jobs_directory(): return path +def temporary_input_path(): + path = jobs_directory() / "inputs" + path.mkdir(parents=True, exist_ok=True) + return path / f"{uuid.uuid4().hex}.png" + + def new_job_paths(job_type): timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") token = uuid.uuid4().hex[:8] diff --git a/src/snapform/worker.py b/src/snapform/worker.py index 5398ab7..9a41b39 100644 --- a/src/snapform/worker.py +++ b/src/snapform/worker.py @@ -229,7 +229,7 @@ def wait_for_exported_paths(output_directory, export_format, timeout=30.0): def reconstruct(args): - write_status(args.status, 0.02, "Starting DA3 runtime") + write_status(args.status, 0.02, "Starting AI runtime") import torch from depth_anything_3.api import DepthAnything3 @@ -291,7 +291,7 @@ def reconstruct(args): ) if not result_paths: raise RuntimeError( - f"DA3 did not export any {args.export_format} files" + f"No {args.export_format} files were exported" ) write_status( args.status, @@ -373,7 +373,7 @@ def main(): else: reconstruct(args) except Exception as error: - write_status(args.status, 1.0, "DA3 task failed", str(error)) + write_status(args.status, 1.0, "Task failed", str(error)) raise diff --git a/tests/test_addon.py b/tests/test_addon.py index 4417877..5cedce3 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -146,14 +146,14 @@ class AddonTest(unittest.TestCase): def test_conversion_dialog_exposes_requested_settings(self): operator = self.addon.ConvertImageToDepthPlane - self.assertEqual( - operator.__annotations__["snapform_model"]["default"], - "BASE", + self.assertNotIn( + "default", + operator.__annotations__["snapform_model"], ) self.assertTrue(operator.__annotations__["snapform_use_ray_pose"]["default"]) self.assertEqual( - operator.__annotations__["snapform_process_res"]["default"], - 1024, + operator.__annotations__["snapform_process_res_factor"]["default"], + 0.5, ) self.assertTrue( operator.__annotations__["snapform_delete_original"]["default"] @@ -165,14 +165,62 @@ class AddonTest(unittest.TestCase): self.assertEqual( set(operator.__annotations__), { + "snapform_depth_source", + "snapform_displacement_via", "snapform_model", "snapform_use_ray_pose", - "snapform_process_res", + "snapform_process_res_factor", "snapform_delete_original", "snapform_subdivision_level", }, ) + depth_sources = operator.__annotations__["snapform_depth_source"] + self.assertEqual(depth_sources["default"], "ORIGINAL") + self.assertEqual( + [item[0] for item in depth_sources["items"]], + ["ESTIMATE", "ORIGINAL"], + ) + self.assertEqual(depth_sources["items"][0][1], "AI Generated") + self.assertEqual(depth_sources["name"], "Depth from") + displacement_via = operator.__annotations__[ + "snapform_displacement_via" + ] + self.assertEqual(displacement_via["default"], "MODIFIER") + self.assertEqual(displacement_via["name"], "Displace via") + self.assertEqual( + [item[0] for item in displacement_via["items"]], + ["MODIFIER", "MATERIAL"], + ) + draw_source = inspect.getsource(operator.draw) + self.assertIn( + 'if self.snapform_depth_source == "ESTIMATE":', + draw_source, + ) + self.assertLess( + draw_source.index( + '"snapform_process_res_factor",' + ), + draw_source.index( + 'layout.prop(self, "snapform_subdivision_level")' + ), + ) + self.assertIn( + 'text=f"Resolution Scale ({process_res} px)"', + draw_source, + ) + self.assertIn("estimate_box = layout.box()", draw_source) + self.assertIn( + 'estimate_box.prop(self, "snapform_model")', + draw_source, + ) + self.assertIn( + 'estimate_box.prop(self, "snapform_use_ray_pose")', + draw_source, + ) + self.assertNotIn("layout.separator", draw_source) + self.assertNotIn("layout.label", draw_source) model_items = operator.__annotations__["snapform_model"]["items"] + model_items = model_items(None, None) self.assertEqual( [item[0] for item in model_items], [ @@ -184,6 +232,127 @@ class AddonTest(unittest.TestCase): ], ) self.assertIn("invoke_props_dialog", inspect.getsource(operator.invoke)) + self.assertNotIn( + '"DA3"', + inspect.getsource(operator), + ) + + def test_depth_conversion_progress_is_provider_neutral(self): + starting_source = inspect.getsource( + self.addon.RunJob._starting_message + ) + finish_source = inspect.getsource(self.addon.RunJob._finish) + + self.assertIn('return "Starting depth estimation"', starting_source) + self.assertIn('"Depth plane is ready"', finish_source) + self.assertNotIn("DA3 image depth", starting_source) + self.assertNotIn("DA3 depth plane", finish_source) + + def test_model_options_mark_only_models_that_are_not_downloaded(self): + items_callback = self.addon.SnapFormSettings.__annotations__["model"][ + "items" + ] + with patch.object( + self.addon.runtime, + "model_ready", + side_effect=lambda key: key == "BASE", + ): + items = items_callback(None, None) + + labels = {item[0]: item[1] for item in items} + self.assertEqual(labels["BASE"], "DA3 Base") + self.assertEqual(labels["SMALL"], "DA3 Small (Missing)") + + def test_packed_image_is_exported_without_unpacking(self): + operator = self.addon.ConvertImageToDepthPlane() + operator.snapform_depth_source = "ESTIMATE" + operator.snapform_model = "BASE" + operator.snapform_use_ray_pose = True + operator.snapform_process_res_factor = 0.5 + operator.snapform_displacement_via = "MATERIAL" + operator.snapform_subdivision_level = 5 + operator.snapform_delete_original = True + image = SimpleNamespace( + filepath="", + packed_file=object(), + size=(1920, 1080), + ) + context = SimpleNamespace( + object=SimpleNamespace(name="Reference", data=image), + scene=object(), + ) + calls = [] + self.fake_bpy.ops = SimpleNamespace( + snapform=SimpleNamespace( + run_job=lambda **options: calls.append(options) + or {"RUNNING_MODAL"}, + ), + ) + + with TemporaryDirectory() as directory: + packed_path = Path(directory) / "packed.png" + packed_path.touch() + with patch.object( + self.addon.image_depth, + "save_packed_image", + return_value=packed_path, + ) as save_packed_image: + result = operator.execute(context) + + self.assertEqual(result, {"FINISHED"}) + save_packed_image.assert_called_once_with(image, context.scene) + self.assertEqual(calls[0]["snapform_input_path"], str(packed_path)) + self.assertEqual(calls[0]["snapform_process_res"], 960) + self.assertEqual( + calls[0]["snapform_displacement_via"], + "MATERIAL", + ) + self.assertTrue(calls[0]["snapform_delete_input"]) + + def test_packed_image_export_writes_a_png_copy(self): + class PackedImage: + def save_render(self, filepath, scene): + self.filepath = filepath + self.scene = scene + self.file_format = scene.render.image_settings.file_format + Path(filepath).write_bytes(b"packed image") + + image = PackedImage() + image_settings = SimpleNamespace(file_format="OPEN_EXR") + scene = SimpleNamespace( + render=SimpleNamespace(image_settings=image_settings), + ) + with TemporaryDirectory() as directory: + output_path = Path(directory) / "input.png" + with patch.object( + self.addon.runtime, + "temporary_input_path", + return_value=output_path, + ): + result = self.addon.image_depth.save_packed_image(image, scene) + + self.assertEqual(result, output_path) + self.assertEqual(output_path.read_bytes(), b"packed image") + self.assertEqual(image.filepath, str(output_path)) + self.assertIs(image.scene, scene) + self.assertEqual(image.file_format, "PNG") + self.assertEqual(image_settings.file_format, "OPEN_EXR") + + def test_temporary_packed_input_is_removed_when_job_closes(self): + closed = [] + with TemporaryDirectory() as directory: + input_path = Path(directory) / "packed.png" + input_path.touch() + self.addon.operators._active_job = SimpleNamespace( + log_file=SimpleNamespace(close=lambda: closed.append(True)), + temporary_input_path=input_path, + ) + + self.addon.operators.close_active_job() + + self.assertFalse(input_path.exists()) + self.assertEqual(closed, [True]) + self.assertIsNone(self.addon.operators.active_job()) def test_image_depth_model_and_ray_pose_are_independent(self): settings = SimpleNamespace( @@ -255,7 +424,7 @@ class AddonTest(unittest.TestCase): ) self.assertEqual( reconstruct.call_args.kwargs["process_res_method"], - "lower_bound_resize", + "upper_bound_resize", ) self.assertEqual( reconstruct.call_args.kwargs["use_ray_pose"], @@ -279,7 +448,7 @@ class AddonTest(unittest.TestCase): str(runtime.DEFAULT_STORAGE_ROOT), ) - def test_process_res_method_defaults_to_lower_bound_resize(self): + def test_process_res_method_defaults_to_upper_bound_resize(self): process_res = self.addon.SnapFormSettings.__annotations__[ "process_res" ] @@ -290,7 +459,7 @@ class AddonTest(unittest.TestCase): ) self.assertEqual(process_res["default"], 1024) - self.assertEqual(process_res_method["default"], "lower_bound_resize") + self.assertEqual(process_res_method["default"], "upper_bound_resize") def test_reconstruct_command_preserves_da3_parameter_names(self): runtime = importlib.import_module("snapform.runtime") @@ -404,7 +573,7 @@ class AddonTest(unittest.TestCase): gs_model = self.addon.SnapFormSettings.__annotations__[ "gs_model" ] - identifiers = [item[0] for item in gs_model["items"]] + identifiers = [item[0] for item in gs_model["items"](None, None)] self.assertEqual( identifiers, @@ -422,13 +591,13 @@ class AddonTest(unittest.TestCase): self.addon.DebugPanel.draw ) - self.assertEqual(draw_source.count("layout.operator("), 1) + self.assertEqual(draw_source.count("export_row.operator("), 1) self.assertIn('settings.export_format == "glb"', draw_source) self.assertIn('settings.export_format == "gs_ply"', draw_source) self.assertIn('"gs_model" if settings.export_format', draw_source) - self.assertIn("model_row.prop(settings, model_property)", draw_source) - self.assertIn('model_row.label(text="", icon="CHECKMARK")', draw_source) - self.assertIn('model_row.label(text="", icon="ERROR")', draw_source) + self.assertIn("layout.prop(settings, model_property)", draw_source) + self.assertNotIn("model_row", draw_source) + self.assertIn("export_row.enabled = active_job() is None", draw_source) self.assertNotIn('text="Ready"', draw_source) self.assertNotIn('text="Missing"', draw_source) self.assertNotIn("License:", draw_source) @@ -444,6 +613,11 @@ class AddonTest(unittest.TestCase): self.assertIn("for model in MODELS.values()", draw_source) self.assertIn("operator.model_key = model.key", model_source) + self.assertIn( + "download_row.enabled = runtime_ready() and not running", + model_source, + ) + self.assertNotIn('label(text="Not downloaded"', model_source) self.assertNotIn("Isolated Runtime ready", draw_source) self.assertNotIn("License:", model_source) self.assertEqual( @@ -451,6 +625,14 @@ class AddonTest(unittest.TestCase): "Setup", ) + def test_preferences_install_button_is_disabled_while_busy(self): + draw_source = inspect.getsource(self.addon.SnapFormPreferences.draw) + + self.assertIn( + "install_row.enabled = active_job() is None", + draw_source, + ) + def test_completed_export_opens_job_directory(self): calls = [] self.fake_bpy.ops = SimpleNamespace( diff --git a/tests/test_image_depth.py b/tests/test_image_depth.py index 06ce5e9..0ee0b32 100644 --- a/tests/test_image_depth.py +++ b/tests/test_image_depth.py @@ -234,7 +234,7 @@ class ImageDepthTest(unittest.TestCase): self.assertAlmostEqual(value, expected) def test_source_empty_is_only_removed_when_requested(self): - source = inspect.getsource(self.image_depth.build_depth_plane) + source = inspect.getsource(self.image_depth.finish_depth_plane) self.assertIn("source_name = source_object.name", source) self.assertIn("if delete_source:", source) @@ -243,27 +243,88 @@ class ImageDepthTest(unittest.TestCase): source, ) self.assertNotIn("source_object.name =", source) - self.assertIn( - "bpy.data.objects.new(\n source_name,", - source, - ) self.assertIn( "bpy.data.objects.remove(source_object, do_unlink=True)\n" " plane_object.name = source_name", source, ) + def test_original_depth_mode_uses_uv_and_zero_midlevel(self): + source = inspect.getsource(self.image_depth.build_image_depth_plane) + + self.assertIn("texture.image = source_object.data", source) + self.assertIn('modifiers.new("Displace", "DISPLACE")', source) + self.assertIn("displace.texture = texture", source) + self.assertIn('displace.texture_coords = "UV"', source) + self.assertNotIn("displace.direction", source) + self.assertNotIn("displace.strength", source) + self.assertIn( + "displace.mid_level = ORIGINAL_DISPLACEMENT_MIDLEVEL", + source, + ) + self.assertIn("plane_object.matrix_world =", source) + + def test_process_resolution_scales_the_image_long_edge(self): + image = SimpleNamespace(size=(1080, 1920)) + + self.assertEqual( + self.image_depth.scaled_process_resolution(image, 0.5), + 960, + ) + + def test_mesh_uses_source_object_name(self): + source = inspect.getsource(self.image_depth.create_plane_mesh) + + self.assertIn("bpy.data.meshes.new(source_object.name)", source) + + def test_depth_data_uses_source_image_name(self): + source_object = SimpleNamespace( + name="obj_img", + data=SimpleNamespace(name="img.png"), + ) + + self.assertEqual( + self.image_depth.depth_data_name(source_object), + "img.depth", + ) + source_object.data.name = "img.png.001" + self.assertEqual( + self.image_depth.depth_data_name(source_object), + "img.001.depth", + ) + estimate_source = inspect.getsource( + self.image_depth.build_depth_plane + ) + original_source = inspect.getsource( + self.image_depth.build_image_depth_plane + ) + self.assertIn("depth_name = depth_data_name(source_object)", estimate_source) + self.assertIn("depth_data_name(source_object)", original_source) + + def test_material_uses_source_object_name(self): + da3_source = inspect.getsource(self.image_depth.build_depth_plane) + image_source = inspect.getsource( + self.image_depth.build_image_depth_plane + ) + + self.assertIn( + "create_material(\n source_name,", + da3_source, + ) + self.assertIn("create_material(\n source_name,", image_source) + self.assertNotIn("DA3DepthMaterial", da3_source) + def test_displace_texture_uses_extend(self): source = inspect.getsource( self.image_depth.build_depth_plane ) - self.assertIn('depth_name = f"{source_name}.depth"', source) + self.assertIn("depth_name = depth_data_name(source_object)", source) self.assertIn("create_depth_image(\n depth_name,", source) - self.assertIn("bpy.data.textures.new(\n depth_name,", source) + self.assertIn("bpy.data.textures.new(\n depth_name,", source) self.assertIn('texture.extension = "EXTEND"', source) - def test_original_image_drives_bump_instead_of_displacement(self): + def test_material_supports_bump_and_optional_displacement(self): source = inspect.getsource(self.image_depth.create_material) self.assertIn('nodes.new("ShaderNodeBump")', source) @@ -297,16 +358,94 @@ class ImageDepthTest(unittest.TestCase): 'curve_points[-1].location = (1.0, 0.1)', source, ) - self.assertIn( - 'shader.inputs["IOR"].default_value = 1.2', - source, - ) + self.assertNotIn('shader.inputs["IOR"]', source) self.assertIn( 'roughness_curve.location = (-300.0, 320.0)', source, ) self.assertIn('bump.location = (-300.0, -80.0)', source) - self.assertNotIn("ShaderNodeDisplacement", source) + self.assertIn('nodes.new("ShaderNodeDisplacement")', source) + self.assertIn( + 'depth_texture.outputs["Color"],\n' + ' displacement.inputs["Height"]', + source, + ) + self.assertIn( + 'displacement.outputs["Displacement"],\n' + ' output.inputs["Displacement"]', + source, + ) + self.assertIn( + 'displacement.inputs["Midlevel"].default_value', + source, + ) + self.assertIn( + 'displacement.inputs["Scale"].default_value', + source, + ) + self.assertIn("set_displacement_only(material)", source) + + def test_displacement_values_follow_depth_source(self): + estimate_source = inspect.getsource( + self.image_depth.build_depth_plane + ) + original_source = inspect.getsource( + self.image_depth.build_image_depth_plane + ) + + self.assertIn( + "displacement_midlevel=DISPLACEMENT_MIDLEVEL", + estimate_source, + ) + self.assertIn( + "displacement_scale=DISPLACEMENT_STRENGTH", + estimate_source, + ) + self.assertIn( + "displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL", + original_source, + ) + self.assertIn( + "displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH", + original_source, + ) + self.assertEqual(self.image_depth.DISPLACEMENT_MIDLEVEL, 1.0) + self.assertEqual(self.image_depth.DISPLACEMENT_STRENGTH, -1.0) + self.assertEqual( + self.image_depth.ORIGINAL_DISPLACEMENT_MIDLEVEL, + 0.0, + ) + self.assertEqual( + self.image_depth.ORIGINAL_DISPLACEMENT_STRENGTH, + 1.0, + ) + + def test_material_displacement_mode_skips_displace_modifier(self): + estimate_source = inspect.getsource( + self.image_depth.build_depth_plane + ) + original_source = inspect.getsource( + self.image_depth.build_image_depth_plane + ) + + for source in (estimate_source, original_source): + self.assertIn('if displacement_via == "MODIFIER":', source) + self.assertIn( + 'if displacement_via == "MATERIAL" else None', + source, + ) + + def test_surface_displacement_is_set_to_displacement_only(self): + class Material: + surface_displacement_method = "BUMP" + + material = Material() + self.image_depth.set_displacement_only(material) + + self.assertEqual( + material.surface_displacement_method, + "DISPLACEMENT", + ) def test_plane_transform_scales_depth_axis(self): source = inspect.getsource(self.image_depth.apply_plane_transform) diff --git a/tests/test_pack.py b/tests/test_pack.py index b84f113..bac218c 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -29,7 +29,7 @@ class PackTest(unittest.TestCase): finally: pack.DIST_DIR = original_dist - self.assertEqual(output_path.name, "SnapForm.v0.9.27.b45.zip") + self.assertEqual(output_path.name, "SnapForm.v0.9.27.zip") with ZipFile(output_path) as archive: names = set(archive.namelist()) diff --git a/tests/test_worker.py b/tests/test_worker.py index 601a7f7..17a1bab 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -43,7 +43,7 @@ class WorkerTest(unittest.TestCase): source = inspect.getsource(self.worker.reconstruct) self.assertLess( - source.index('"Starting DA3 runtime"'), + source.index('"Starting AI runtime"'), source.index("import torch"), )