82 lines
2.6 KiB
Python
82 lines
2.6 KiB
Python
import os
|
|
import folder_paths
|
|
import numpy as np
|
|
import torch.nn.functional as F
|
|
import torch
|
|
from comfy.utils import ProgressBar
|
|
import cv2
|
|
from .utilities import Engine
|
|
from .engine_builder_node import DepthAnythingEngineBuilder
|
|
|
|
ENGINE_DIR = os.path.join(folder_paths.models_dir,"tensorrt", "depth-anything")
|
|
os.makedirs(ENGINE_DIR, exist_ok=True)
|
|
|
|
class DepthAnythingTensorrt:
|
|
def __init__(self):
|
|
self.engine = None
|
|
self.engine_label = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"engine": (os.listdir(ENGINE_DIR),),
|
|
}
|
|
}
|
|
RETURN_NAMES = ("IMAGE",)
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "main"
|
|
CATEGORY = "tensorrt"
|
|
|
|
def main(self, images, engine):
|
|
|
|
# setup tensorrt engine
|
|
if (not hasattr(self, 'engine') or self.engine_label != engine):
|
|
self.engine = Engine(os.path.join(ENGINE_DIR,engine))
|
|
self.engine.load()
|
|
self.engine.activate()
|
|
self.engine.allocate_buffers()
|
|
self.engine_label = engine
|
|
|
|
cudaStream = torch.cuda.current_stream().cuda_stream
|
|
pbar = ProgressBar(images.shape[0])
|
|
images = images.permute(0, 3, 1, 2)
|
|
images_resized = F.interpolate(images, size=(518,518), mode='bilinear', align_corners=False)
|
|
images_list = list(torch.split(images_resized, split_size_or_sections=1))
|
|
|
|
depth_frames = []
|
|
|
|
for img in images_list:
|
|
result = self.engine.infer({"input": img},cudaStream)
|
|
depth = result['output']
|
|
|
|
# Process the depth output
|
|
depth = np.reshape(depth.cpu().numpy(), (518,518))
|
|
depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
|
|
depth = depth.astype(np.uint8)
|
|
depth = cv2.resize(depth, (images.shape[3], images.shape[2]))
|
|
depth = cv2.cvtColor(depth, cv2.COLOR_RGB2BGR)
|
|
|
|
depth_frames.append(depth)
|
|
pbar.update(1)
|
|
|
|
depth_frames_np = np.array(depth_frames).astype(np.float32) / 255.0
|
|
if "metric" in engine:
|
|
result = 1 - torch.from_numpy(depth_frames_np)
|
|
else:
|
|
result = torch.from_numpy(depth_frames_np)
|
|
return (result,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DepthAnythingTensorrt" : DepthAnythingTensorrt,
|
|
"DepthAnythingEngineBuilder" : DepthAnythingEngineBuilder,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DepthAnythingTensorrt" : "Depth Anything Tensorrt",
|
|
"DepthAnythingEngineBuilder" : "Depth Anything Engine Builder",
|
|
}
|
|
|
|
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|