import shutil from pathlib import Path import bpy from .job_runtime import JobOperatorBase, runtime from .properties import ( BEN2_MODEL_KEY, DEFAULT_MOGE2_MODEL_KEY, REQUIRED_MODEL_KEYS, model_ready, model_record, ) IMAGE_DEPTH_DEFAULT_PROCESS_RES = 1024 IMAGE_DEPTH_PROCESS_RES_METHOD = "upper_bound_resize" BEN2_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"} BEN2_VIDEO_SUFFIXES = { ".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v", ".mts", ".mpg", ".mpeg", } HIDDEN = {"HIDDEN", "SKIP_SAVE"} def input_path(value): path = Path(bpy.path.abspath(value)) if path.is_file() or path.is_dir(): return path raise ValueError("Select a valid image, video, or frame directory") def background_input_path(value): path = input_path(value) if path.is_dir() or path.suffix.lower() in ( BEN2_IMAGE_SUFFIXES | BEN2_VIDEO_SUFFIXES ): return path raise ValueError( "BEN2 requires an image, video, or directory of image frames" ) def source_start_frame(source_object_name): from .blender.depth_plane import extract_frame_range source_object = bpy.data.objects.get(source_object_name) return extract_frame_range(source_object)[0] if source_object else 1 def production_device(): addon = bpy.context.preferences.addons.get(__package__) return getattr(getattr(addon, "preferences", None), "device", "AUTO") def require_environment(): if not runtime.environment_ready(): raise RuntimeError("Install the SnapForm Environment first") def require_model(model_key): model = model_record(model_key) if model is None: raise RuntimeError(f"Unknown model: {model_key}") if not model_ready(model_key): raise RuntimeError(f"Download {model['label']} first") return model def cleanup_temporary_input(path, delete_input): if not delete_input or not path: return temporary_path = Path(path) try: if temporary_path.is_dir(): shutil.rmtree(temporary_path, ignore_errors=True) return temporary_path.unlink(missing_ok=True) temporary_path.parent.rmdir() except OSError: pass def open_job_directory(directory): directory = Path(directory) if not directory.is_dir(): raise RuntimeError(f"Job directory does not exist: {directory}") result = bpy.ops.wm.path_open(filepath=str(directory)) if "FINISHED" not in result: raise RuntimeError(f"Unable to open job directory: {directory}") def push_undo(message): result = bpy.ops.ed.undo_push(message=message) if "FINISHED" not in result: raise RuntimeError(f"Unable to create the {message} Undo step") def build_da3_depth_response( context, result, *, source_object_name, depth_data_name, remove_background, delete_source, unshaded, ): from .blender.depth_plane import ( build_depth_plane, create_removed_background_image, create_sequence_image, image_base_name, is_image_empty, ) directory = result.directory source_object = bpy.data.objects.get(source_object_name) if not is_image_empty(source_object): raise RuntimeError("The source image empty is no longer available") foreground_paths = sorted(directory.glob("foreground.*.png")) foreground_image = None if remove_background: foreground_image = create_removed_background_image( source_object, directory if foreground_paths else directory / "foreground.png", assign_to_source=False, ) color_image = None if list(directory.glob("frame.*")) and foreground_image is None: color_image = create_sequence_image( f"{image_base_name(source_object.data)}.frame", directory, ) build_depth_plane( context, source_object, result.file("depth"), delete_source, unshaded=unshaded, depth_name=depth_data_name, color_image=color_image, foreground_image=foreground_image, depth_frames_directory=( directory if list(directory.glob("frame.*")) else None ), ) push_undo("Convert to Depth Plane") def build_moge2_depth_response( context, result, *, source_object_name, remove_background, delete_source, unshaded, ): from .blender.depth_plane import ( build_moge2_depth_plane, create_removed_background_image, create_sequence_image, image_base_name, is_image_empty, ) directory = result.directory source_object = bpy.data.objects.get(source_object_name) if not is_image_empty(source_object): raise RuntimeError("The source image empty is no longer available") foreground_paths = sorted(directory.glob("foreground.*.png")) foreground_image = None if remove_background: foreground_image = create_removed_background_image( source_object, directory if foreground_paths else directory / "foreground.png", assign_to_source=False, ) color_image = None if list(directory.glob("frame.*")) and foreground_image is None: color_image = create_sequence_image( f"{image_base_name(source_object.data)}.frame", directory, ) build_moge2_depth_plane( context, source_object, result.file("depth"), delete_source, unshaded=unshaded, color_image=color_image, foreground_image=foreground_image, data_frames_directory=directory, ) push_undo("Convert to Depth Plane") def build_original_depth_response( context, result, *, source_object_name, background, delete_source, unshaded, ): from .blender.depth_plane import ( build_image_depth_plane, create_removed_background_image, create_sequence_image, image_base_name, is_image_empty, ) directory = result.directory source_object = bpy.data.objects.get(source_object_name) if not is_image_empty(source_object): raise RuntimeError("The source image empty is no longer available") foreground_image = None if background: foreground_paths = sorted(directory.glob("foreground.*.png")) foreground_image = create_removed_background_image( source_object, directory if foreground_paths else directory / "foreground.png", assign_to_source=False, ) color_image = None if list(directory.glob("frame.*")) and foreground_image is None: color_image = create_sequence_image( f"{image_base_name(source_object.data)}.frame", directory, ) build_image_depth_plane( context, source_object, delete_source, unshaded=unshaded, color_image=color_image, foreground_image=foreground_image, depth_frames_directory=( directory if list(directory.glob("frame.*")) else None ), ) push_undo("Convert to Depth Plane") def build_surface_response( context, result, *, source_object_name, delete_source, create_camera, unshaded, moge2, ): from .blender.depth_plane import is_image_empty from .blender.surface import build_image_mesh source_object = bpy.data.objects.get(source_object_name) if not is_image_empty(source_object): raise RuntimeError("The source Image Empty no longer exists") build_image_mesh( context, source_object, result.file("surface"), delete_source, create_camera, unshaded, moge2, ) push_undo("Convert to Surface") def apply_background_response(result, source_object_name): from .blender.depth_plane import ( create_removed_background_image, is_image_empty, ) source_object = bpy.data.objects.get(source_object_name) if not is_image_empty(source_object): raise RuntimeError("The source image empty is no longer available") foreground_paths = sorted(result.directory.glob("foreground.*.png")) create_removed_background_image( source_object, result.directory if foreground_paths else result.file("foreground"), ) push_undo("Remove Background") def da3_parameters(settings, path, model, process_res, start_frame): return { "input": str(path), "model": model, "process_res": process_res, "process_res_method": "upper_bound_resize", "conf_thresh_percentile": settings.conf_thresh_percentile, "num_max_points": settings.num_max_points, "show_cameras": 1 if settings.show_cameras else 0, "gs_views_interval": settings.gs_views_interval, "video_frames": settings.video_frames, "start_frame": start_frame, "device": production_device().lower(), "ref_view_strategy": settings.ref_view_strategy, } def moge2_parameters(settings, path, model, resolution_level, start_frame): return { "input": str(path), "model": model, "resolution_level": int(resolution_level), "video_frames": settings.video_frames, "start_frame": start_frame, "device": production_device().lower(), } class DownloadModel(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.download_model" bl_label = "Download Model" job_type = "download-model" model: bpy.props.StringProperty(options=HIDDEN) def request(self, _context): require_environment() if not bpy.app.online_access: raise RuntimeError("Online Access is disabled in Blender Preferences") model = model_record(self.model) label = model["label"] if model else self.model self.starting_message = f"Starting {label} download" return {"model": self.model} class DownloadRequiredModels(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.download_required_models" bl_label = "Download Required Models" job_type = "download-required-models" starting_message = "Starting required model downloads" def request(self, _context): require_environment() if not bpy.app.online_access: raise RuntimeError("Online Access is disabled in Blender Preferences") return {"models": list(REQUIRED_MODEL_KEYS)} class DebugDa3(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.debug_da3" bl_label = "Debug DA3" job_type = "debug-da3" def request(self, context): require_environment() settings = context.scene.snapform_settings model = settings.gs_model if settings.export_format == "gs_ply" else settings.model spec = require_model(model) if settings.export_format == "gs_ply" and not spec.get("supports_gs"): raise RuntimeError("PLY export requires a DA3 model with a Gaussian head") self.export_format = settings.export_format self.starting_message = f"Starting {self.export_format} export" return { **da3_parameters( settings, input_path(settings.input_path), model, settings.process_res, 1, ), "device": settings.device.lower(), "export_format": settings.export_format, "inference_mode": "joint", "remove_background_model": "", "mesh_filter_edges": settings.mesh_filter_edges, "mesh_max_edge_length": settings.mesh_max_edge_length, } def response(self, _context, result): open_job_directory(result.directory) return f"Exported {self.export_format} to {result.directory}" class DebugMoge2(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.debug_moge2" bl_label = "Debug MoGe-2" job_type = "debug-moge2" starting_message = "Starting MoGe-2 inference" def request(self, context): require_environment() settings = context.scene.snapform_settings require_model(settings.moge_model) return { **moge2_parameters( settings, input_path(settings.input_path), settings.moge_model, settings.moge_resolution_level, 1, ), "device": settings.device.lower(), "export_format": settings.moge_export_format, "include_points": True, } def response(self, _context, result): open_job_directory(result.directory) return f"Exported MoGe-2 data to {result.directory}" class DebugBen2(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.debug_ben2" bl_label = "Debug BEN2" job_type = "debug-ben2" starting_message = "Starting BEN2 background removal" def request(self, context): require_environment() settings = context.scene.snapform_settings require_model(BEN2_MODEL_KEY) return { "input": str(background_input_path(settings.background_input_path)), "model": BEN2_MODEL_KEY, "device": settings.device.lower(), "video_frames": settings.video_frames, "start_frame": 1, } def response(self, _context, result): open_job_directory(result.directory) return f"Removed background to {result.directory}" class GenerateDa3Depth(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_da3_depth" bl_label = "Generate DA3 Depth" job_type = "generate-da3-depth" starting_message = "Starting depth estimation" input_path: bpy.props.StringProperty(options=HIDDEN) source_object_name: bpy.props.StringProperty(options=HIDDEN) depth_data_name: bpy.props.StringProperty(options=HIDDEN) model: bpy.props.StringProperty(default="BASE", options=HIDDEN) process_res: bpy.props.IntProperty(default=1024, min=1, options=HIDDEN) delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN) remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) def request(self, context): require_environment() require_model(self.model) if self.remove_background: require_model(BEN2_MODEL_KEY) settings = context.scene.snapform_settings return { **da3_parameters( settings, input_path(self.input_path), self.model, self.process_res, source_start_frame(self.source_object_name), ), "export_format": "mini_npz", "inference_mode": "per_frame", "remove_background_model": ( BEN2_MODEL_KEY if self.remove_background else "" ), } def response(self, context, result): build_da3_depth_response( context, result, source_object_name=self.source_object_name, depth_data_name=self.depth_data_name, remove_background=self.remove_background, delete_source=self.delete_source, unshaded=self.unshaded, ) return "Depth plane is ready" def cleanup(self): cleanup_temporary_input(self.input_path, self.delete_input) class GenerateMoge2Depth(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_moge2_depth" bl_label = "Generate MoGe-2 Depth" job_type = "generate-moge2-depth" starting_message = "Starting MoGe-2 inference" input_path: bpy.props.StringProperty(options=HIDDEN) source_object_name: bpy.props.StringProperty(options=HIDDEN) model: bpy.props.StringProperty(default=DEFAULT_MOGE2_MODEL_KEY, options=HIDDEN) resolution_level: bpy.props.IntProperty(default=5, min=0, max=9, options=HIDDEN) delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN) remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) def request(self, context): require_environment() require_model(self.model) if self.remove_background: require_model(BEN2_MODEL_KEY) settings = context.scene.snapform_settings return { **moge2_parameters( settings, input_path(self.input_path), self.model, self.resolution_level, source_start_frame(self.source_object_name), ), "remove_background_model": ( BEN2_MODEL_KEY if self.remove_background else "" ), } def response(self, context, result): build_moge2_depth_response( context, result, source_object_name=self.source_object_name, remove_background=self.remove_background, delete_source=self.delete_source, unshaded=self.unshaded, ) return "MoGe-2 depth plane is ready" def cleanup(self): cleanup_temporary_input(self.input_path, self.delete_input) class GenerateDa3Surface(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_da3_surface" bl_label = "Generate DA3 Surface" job_type = "generate-da3-surface" starting_message = "Starting perspective surface generation" input_path: bpy.props.StringProperty(options=HIDDEN) source_object_name: bpy.props.StringProperty(options=HIDDEN) model: bpy.props.StringProperty(default="BASE", options=HIDDEN) process_res: bpy.props.IntProperty(default=1024, min=1, options=HIDDEN) delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) create_camera: bpy.props.BoolProperty(default=False, options=HIDDEN) unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN) remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) def request(self, context): require_environment() require_model(self.model) if self.remove_background: require_model(BEN2_MODEL_KEY) settings = context.scene.snapform_settings return { **da3_parameters( settings, input_path(self.input_path), self.model, self.process_res, source_start_frame(self.source_object_name), ), "export_format": "mesh_data", "inference_mode": "per_frame", "remove_background_model": ( BEN2_MODEL_KEY if self.remove_background else "" ), "mesh_filter_edges": settings.mesh_filter_edges, "mesh_max_edge_length": settings.mesh_max_edge_length, } def response(self, context, result): build_surface_response( context, result, source_object_name=self.source_object_name, delete_source=self.delete_source, create_camera=self.create_camera, unshaded=self.unshaded, moge2=False, ) return "Perspective surface is ready" def cleanup(self): cleanup_temporary_input(self.input_path, self.delete_input) class GenerateMoge2Surface(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_moge2_surface" bl_label = "Generate MoGe-2 Surface" job_type = "generate-moge2-surface" starting_message = "Starting MoGe-2 inference" input_path: bpy.props.StringProperty(options=HIDDEN) source_object_name: bpy.props.StringProperty(options=HIDDEN) model: bpy.props.StringProperty(default=DEFAULT_MOGE2_MODEL_KEY, options=HIDDEN) resolution_level: bpy.props.IntProperty(default=5, min=0, max=9, options=HIDDEN) delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) create_camera: bpy.props.BoolProperty(default=False, options=HIDDEN) unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN) remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) def request(self, context): require_environment() require_model(self.model) if self.remove_background: require_model(BEN2_MODEL_KEY) settings = context.scene.snapform_settings return { **moge2_parameters( settings, input_path(self.input_path), self.model, self.resolution_level, source_start_frame(self.source_object_name), ), "remove_background_model": ( BEN2_MODEL_KEY if self.remove_background else "" ), "surface_mesh": True, } def response(self, context, result): build_surface_response( context, result, source_object_name=self.source_object_name, delete_source=self.delete_source, create_camera=self.create_camera, unshaded=self.unshaded, moge2=True, ) return "MoGe-2 surface is ready" def cleanup(self): cleanup_temporary_input(self.input_path, self.delete_input) class ProcessImageBackground(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.process_image_background" bl_label = "Process Image Background" job_type = "remove-background" starting_message = "Starting Image Empty background removal" input_path: bpy.props.StringProperty(options=HIDDEN) source_object_name: bpy.props.StringProperty(options=HIDDEN) delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) def request(self, context): require_environment() require_model(BEN2_MODEL_KEY) settings = context.scene.snapform_settings return { "input": str(background_input_path(self.input_path)), "model": BEN2_MODEL_KEY, "device": production_device().lower(), "video_frames": settings.video_frames, "start_frame": source_start_frame(self.source_object_name), } def response(self, _context, result): apply_background_response(result, self.source_object_name) return "Image background is removed" def cleanup(self): cleanup_temporary_input(self.input_path, self.delete_input) class GenerateOriginalBackground(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_original_background" bl_label = "Generate Original Background" job_type = "remove-background" starting_message = "Starting background removal" input_path: bpy.props.StringProperty(options=HIDDEN) source_object_name: bpy.props.StringProperty(options=HIDDEN) delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN) def request(self, context): require_environment() require_model(BEN2_MODEL_KEY) settings = context.scene.snapform_settings return { "input": str(background_input_path(self.input_path)), "model": BEN2_MODEL_KEY, "device": production_device().lower(), "video_frames": settings.video_frames, "start_frame": source_start_frame(self.source_object_name), } def response(self, context, result): build_original_depth_response( context, result, source_object_name=self.source_object_name, background=True, delete_source=self.delete_source, unshaded=self.unshaded, ) return "Depth plane is ready" def cleanup(self): cleanup_temporary_input(self.input_path, self.delete_input) class ExtractOriginalFrames(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.extract_original_frames" bl_label = "Extract Original Frames" job_type = "extract-frames" starting_message = "Extracting source frames" input_path: bpy.props.StringProperty(options=HIDDEN) source_object_name: bpy.props.StringProperty(options=HIDDEN) delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN) def request(self, context): require_environment() settings = context.scene.snapform_settings return { "input": str(input_path(self.input_path)), "video_frames": settings.video_frames, "start_frame": source_start_frame(self.source_object_name), } def response(self, context, result): build_original_depth_response( context, result, source_object_name=self.source_object_name, background=False, delete_source=self.delete_source, unshaded=self.unshaded, ) return "Depth plane is ready" def cleanup(self): cleanup_temporary_input(self.input_path, self.delete_input) class ClearModels(bpy.types.Operator): bl_idname = "snapform.clear_models" bl_label = "Clear Models" def execute(self, _context): try: runtime.clear_resource("model_manager") except RuntimeError as error: self.report({"ERROR"}, str(error)) return {"CANCELLED"} return {"FINISHED"} JOB_OPERATOR_CLASSES = ( DownloadModel, DownloadRequiredModels, DebugDa3, DebugMoge2, DebugBen2, GenerateDa3Depth, GenerateMoge2Depth, GenerateDa3Surface, GenerateMoge2Surface, ProcessImageBackground, GenerateOriginalBackground, ExtractOriginalFrames, ClearModels, )