Zoe Depth Anything 预处理器类

报告文档：Zoe Depth Anything 检测器和预处理器类

导入语句

from ..utils import common_annotator_call, create_node_input_types
import comfy.model_management as model_management

这部分代码导入了通用的辅助函数以及模型管理模块，用于处理图像注释和模型的设备管理。

类 ZoeDepthAnythingDetector 的定义

构造函数

def __init__(self, model):
    self.model = model
    self.device = "cpu"

构造函数接收一个模型对象，并默认将设备设置为 “cpu”。

从预训练模型加载

@classmethod
def from_pretrained(cls, pretrained_model_or_path=DEPTH_ANYTHING_MODEL_NAME,
                    filename="depth_anything_metric_depth_indoor.pt"):
    model_path = custom_hf_download(pretrained_model_or_path, filename, subfolder="checkpoints_metric_depth",
                                    repo_type="space")
    conf = get_config("zoedepth", "infer")
    model = ZoeDepthAnything.build_from_config(conf)
    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'))['model'])
    model.eval()
    return cls(model)

这个类方法用于从预训练模型加载 ZoeDepthAnythingDetector。它下载模型文件，构建模型配置，并加载模型状态。

更改设备

def to(self, device):
    self.model.to(device)
    self.device = device
    return self

此方法用于将模型移动到指定的设备（例如 GPU 或 CPU）。

模型调用

def __call__(self, input_image, detect_resolution=512, output_type=None, upscale_method="INTER_CUBIC", **kwargs):
    input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
    input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
    image_depth = input_image
    with torch.no_grad():
        image_depth = torch.from_numpy(image_depth).float().to(self.device)
        image_depth = image_depth / 255.0
        image_depth = rearrange(image_depth, 'h w c -> 1 c h w')
        depth = self.model.infer(image_depth)
        depth = depth[0, 0].cpu().numpy()
        vmin = np.percentile(depth, 2)
        vmax = np.percentile(depth, 85)
        depth -= vmin
        depth /= vmax - vmin
        depth = 1.0 - depth
        depth_image = (depth * 255.0).clip(0, 255).astype(np.uint8)
    detected_map = remove_pad(HWC3(depth_image))
    if output_type == "pil":
        detected_map = Image.fromarray(detected_map)
    return detected_map

这个方法是类的主功能实现，用于处理输入图像，进行深度估计，并输出处理后的图像。它处理输入图像，调整尺寸，计算深度图，并根据指定输出格式返回结果。

类 Zoe_Depth_Anything_Preprocessor 的定义

输入类型定义

@classmethod
def INPUT_TYPES(s):
    return create_node_input_types(
        environment=(["indoor", "outdoor"], {"default": "indoor"})
    )

定义了预处理器需要的输入类型，包括环境参数。

执行函数

def execute(self, image, environment, resolution=512, **kwargs):
    ckpt_name = "depth_anything_metric_depth_indoor.pt" if environment == "indoor" else "depth_anything_metric_depth_outdoor.pt"
    model = ZoeDepthAnythingDetector.from_pretrained(filename=ckpt_name).to(model_management.get_torch_device())
    out = common_annotator_call(model, image, resolution=resolution)
    del model
    return (out, )

此函数加载适当的深度检测模型，并使用该模型处理图像，生成深度估计图。根据环境选择适当的预训练模型，并将模型移至适当的设备进行处理。

节点类和显示名称的注册

NODE_CLASS_MAPPINGS = {
    "Zoe_DepthAnythingPreprocessor": Zoe_Depth_Anything_Preprocessor
}
NODE_DISPLAY_NAME

_MAPPINGS = {
    "Zoe_DepthAnythingPreprocessor": "Zoe Depth Anything"
}

这些代码行将上述定义的类注册为可用的节点类型，并设置它们的显示名称。这使得这些功能可以在更大的系统或框架中通过指定名称被调用。