nodes/killertofu86-ComfyUI-DukeStereo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

implemented polyline fill

Browse Source
This commit is contained in:
Lukas Schmidt
2026-02-09 21:18:41 +01:00
parent 69867c2056
commit a3629f2aa4
3 changed files with 205 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -0,0 +1 @@
__pycache__/nodes.cpython-311.pyc
__pycache__/nodes.cpython-311.pyc
BIN
View File
Binary file not shown.
nodes.py
+204 -51
View File
@@ -1,5 +1,6 @@
# my_node.py
import torch, numpy as np
# duke_stereo_nodes.py - DukeStereoSBS with Polylines Fill
import torch
import numpy as np
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForDepthEstimation
from time import time
@@ -7,18 +8,129 @@ from queue import Queue
from threading import Thread
import cv2
try:
from numba import njit, prange
except Exception as e:
print(f"WARNING! Numba failed to import! Stereo generation will be much slower! ({str(e)})")
from builtins import range as prange
def njit(parallel=False):
def Inner(func): return lambda *args, **kwargs: func(*args, **kwargs)
return Inner
def writer_thread(out_path, fps, width, height, queue):
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(out_path, fourcc, fps, (width*2, height))
out = cv2.VideoWriter(out_path, fourcc, fps, (width * 2, height))
while True:
item = queue.get()
if item is None: break
if item is None:
break
out.write(item)
out.release()
@njit(parallel=True)
def Das mit den Friends haben wir extra gemacht, dass er sie direkt wegschreibt, weil ich sonst mit meinem Ramm nicht hinkomme.apply_stereo_divergence_polylines(original_image, normalized_depth, divergence_px: float, separation_px: float, stereo_offset_exponent: float, fill_technique: str):
"""
Polylines-basierter Stereo-Warp mit intelligentem Lücken-Füllen.
Aus stereoimage_generation.py übernommen.
"""
EPSILON = 1e-7
PIXEL_HALF_WIDTH = 0.45 if fill_technique == 'polylines_sharp' else 0.0
h, w, c = original_image.shape
derived_image = np.zeros_like(original_image)
for row in prange(h):
pt = np.zeros((5 + 2 * w, 3), dtype=np.float32)
pt_end: int = 0
pt[pt_end] = [-1.0 * w, 0.0, 0.0]
pt_end += 1
for col in range(0, w):
coord_d = (normalized_depth[row][col] ** stereo_offset_exponent) * divergence_px
coord_x = col + 0.5 + coord_d + separation_px
if PIXEL_HALF_WIDTH < EPSILON:
pt[pt_end] = [coord_x, abs(coord_d), col]
pt_end += 1
else:
pt[pt_end] = [coord_x - PIXEL_HALF_WIDTH, abs(coord_d), col]
pt[pt_end + 1] = [coord_x + PIXEL_HALF_WIDTH, abs(coord_d), col]
pt_end += 2
pt[pt_end] = [2.0 * w, 0.0, w - 1]
pt_end += 1
sg_end: int = pt_end - 1
sg = np.zeros((sg_end, 6), dtype=np.float32)
for i in range(sg_end):
sg[i] += np.concatenate((pt[i], pt[i + 1]))
for i in range(1, sg_end):
u = i - 1
while pt[u][0] > pt[u + 1][0] and 0 <= u:
pt[u], pt[u + 1] = np.copy(pt[u + 1]), np.copy(pt[u])
sg[u], sg[u + 1] = np.copy(sg[u + 1]), np.copy(sg[u])
u -= 1
csg = np.zeros((5 * int(abs(divergence_px)) + 25, 6), dtype=np.float32)
csg_end: int = 0
sg_pointer: int = 0
pt_i: int = 0
for col in range(w):
color = np.full(c, 0.5, dtype=np.float32)
while pt[pt_i][0] < col:
pt_i += 1
pt_i -= 1
while pt[pt_i][0] < col + 1:
coord_from = max(col, pt[pt_i][0]) + EPSILON
coord_to = min(col + 1, pt[pt_i + 1][0]) - EPSILON
significance = coord_to - coord_from
coord_center = coord_from + 0.5 * significance
while sg_pointer < sg_end and sg[sg_pointer][0] < coord_center:
csg[csg_end] = sg[sg_pointer]
sg_pointer += 1
csg_end += 1
csg_i = 0
while csg_i < csg_end:
if csg[csg_i][3] < coord_center:
csg[csg_i] = csg[csg_end - 1]
csg_end -= 1
else:
csg_i += 1
best_csg_i: int = 0
if csg_end != 1:
best_csg_closeness: float = -EPSILON
for csg_i in range(csg_end):
ip_k = (coord_center - csg[csg_i][0]) / (csg[csg_i][3] - csg[csg_i][0])
closeness = (1.0 - ip_k) * csg[csg_i][1] + ip_k * csg[csg_i][4]
if best_csg_closeness < closeness and 0.0 < ip_k < 1.0:
best_csg_closeness = closeness
best_csg_i = csg_i
col_l: int = int(csg[best_csg_i][2] + EPSILON)
col_r: int = int(csg[best_csg_i][5] + EPSILON)
if col_l == col_r:
color += original_image[row][col_l] * significance
else:
ip_k = (coord_center - csg[best_csg_i][0]) / (csg[best_csg_i][3] - csg[best_csg_i][0])
color += (original_image[row][col_l] * (1.0 - ip_k) +
original_image[row][col_r] * ip_k
) * significance
pt_i += 1
derived_image[row][col] = np.asarray(color, dtype=np.uint8)
return derived_image
class DukeStereoSBS:
"""
Ein Node, der aus Bilder hauptsächlich auf der GPU SBS Stero bilder der selben auflösung (breite mal 2) macht.
Node für GPU-beschleunigte SBS Stereo-Bilder mit Polylines Fill-Technik.
"""
@classmethod
@@ -26,16 +138,19 @@ class DukeStereoSBS:
return {
"required": {
"output_path": ("STRING", {"default": "/ComfyUI/output/stereo_output.mp4"}),
"images": ("IMAGE",), # Bild-Input von anderem Node
"images": ("IMAGE",),
},
"optional": {
"model": (["Large", "Base"],), # COMBO/Dropdown
"model": (["Large", "Base"],),
"depth_size": ("INT", {"default": 518, "min": 128, "max": 1024}),
"depth_batch_size": ("INT", {"default": 64, "min": 4, "max": 256}),
"warp_batch_size": ("INT", {"default": 64, "min": 4, "max": 256}),
"depth_blur": ("INT", {"default": 6, "min": 0, "max": 10}),
"divergence": ("FLOAT", {"default": 2.5, "min": 0.0, "max": 10.0, "step": 0.1}),
"fps": ("INT", {"default": 24, "min": 1, "max": 120, "step": 1})
"fps": ("INT", {"default": 24, "min": 1, "max": 120, "step": 1}),
"fill_technique": (["polylines_soft", "polylines_sharp", "none"],),
"stereo_offset_exponent": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 4.0, "step": 0.1}),
"convergence_point": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05}),
}
}
@@ -45,75 +160,113 @@ class DukeStereoSBS:
CATEGORY = "Stereo"
@torch.no_grad()
def depth_batch(self,imgs, depth_batch_size,depth_size):
def depth_batch(self, imgs, depth_batch_size, depth_size):
depths = []
for i in range(0, len(imgs), depth_batch_size):
batch = imgs[i:i+depth_batch_size]
inputs = self.proc(images=batch, return_tensors="pt",size=depth_size).to("cuda")
batch = imgs[i:i + depth_batch_size]
inputs = self.proc(images=batch, return_tensors="pt", size=depth_size).to("cuda")
out = self.model(**inputs).predicted_depth
depths.extend([d.cpu() for d in out])
depths.extend([d.cpu().numpy() for d in out])
return depths
@torch.no_grad()
def stereo_warp_streaming(self,imgs, depths, div, warp_batch_size,depth_blur):
def normalize_depth(self, depth, h, w, convergence_point):
"""Depth auf Bildgröße interpolieren und normalisieren mit Convergence Point."""
depth_t = torch.from_numpy(depth).unsqueeze(0).unsqueeze(0).float()
depth_t = torch.nn.functional.interpolate(depth_t, size=(h, w), mode="bilinear", align_corners=True)
depth_np = depth_t.squeeze().numpy()
# Normalisieren auf 0-1
d_min, d_max = depth_np.min(), depth_np.max()
normalized = (depth_np - d_min) / (d_max - d_min + 1e-7)
# Convergence point anwenden (verschiebt wo Null-Parallax liegt)
normalized = normalized - convergence_point
return normalized
def stereo_warp_polylines(self, imgs, depths, divergence, depth_blur, warp_batch_size,
fill_technique, stereo_offset_exponent, convergence_point):
"""Stereo-Warp mit Polylines Fill-Technik."""
divergence_px = divergence * (self.w / 100)
for i in range(0, len(imgs), warp_batch_size):
batch_imgs = imgs[i:i+warp_batch_size]
batch_depths = depths[i:i+warp_batch_size]
img_t = torch.stack([torch.from_numpy(np.array(im)).permute(2,0,1) for im in batch_imgs]).float().cuda() / 255
depth_t = torch.stack([d.cuda() for d in batch_depths]).unsqueeze(1).float()
depth_t = torch.nn.functional.interpolate(depth_t, size=(self.h, self.w), mode="bilinear", align_corners=True)
if depth_blur > 0:
kernel = depth_blur if depth_blur % 2 == 1 else depth_blur + 1 # ungerade machen
pad = kernel // 2
depth_t = torch.nn.functional.avg_pool2d(
torch.nn.functional.pad(depth_t, (pad, pad, pad, pad), mode='replicate'),
kernel, stride=1, padding=0
)
depth_t = depth_t / depth_t.amax(dim=(2,3), keepdim=True)
disp = (depth_t - 0.5) * div * (self.w / 100)
grid = self.base_grid.expand(len(batch_imgs), -1, -1, -1).clone()
gl, gr = grid.clone(), grid.clone()
gl[..., 0] += disp.squeeze(1) * 2 / self.w
gr[..., 0] -= disp.squeeze(1) * 2 / self.w
left = torch.nn.functional.grid_sample(img_t, gl, padding_mode='border', align_corners=True)
right = torch.nn.functional.grid_sample(img_t, gr, padding_mode='border', align_corners=True)
# Optimiert: Alles auf GPU
sbs = torch.cat([left, right], dim=3) # SBS auf GPU
sbs = sbs.flip(dims=[1])
sbs = (sbs.permute(0,2,3,1) * 255).to(torch.uint8) # NCHW→NHWC, uint8
sbs_np = sbs.cpu().numpy() # Ein Transfer pro Batch
for frame in sbs_np:
self.write_queue.put(frame)
batch_imgs = imgs[i:i + warp_batch_size]
batch_depths = depths[i:i + warp_batch_size]
for img, depth in zip(batch_imgs, batch_depths):
img_np = np.array(img, dtype=np.uint8)
depth_norm = self.normalize_depth(depth, self.h, self.w, convergence_point)
# Optional: Depth Blur
if depth_blur > 0:
kernel = depth_blur if depth_blur % 2 == 1 else depth_blur + 1
depth_norm = cv2.blur(depth_norm, (kernel, kernel))
if fill_technique == "none":
# Fallback: einfacher grid_sample (alter Code)
left = self._simple_warp(img_np, depth_norm, divergence_px)
right = self._simple_warp(img_np, depth_norm, -divergence_px)
else:
# Polylines-basierter Warp
left = apply_stereo_divergence_polylines(
img_np, depth_norm, divergence_px, 0.0, stereo_offset_exponent, fill_technique
)
right = apply_stereo_divergence_polylines(
img_np, depth_norm, -divergence_px, 0.0, stereo_offset_exponent, fill_technique
)
# SBS zusammenfügen (BGR für OpenCV)
sbs = np.concatenate([left, right], axis=1)
sbs_bgr = cv2.cvtColor(sbs, cv2.COLOR_RGB2BGR)
self.write_queue.put(sbs_bgr)
def _simple_warp(self, img_np, depth_norm, divergence_px):
"""Einfacher GPU-Warp als Fallback (alte Methode)."""
img_t = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).float().cuda() / 255
depth_t = torch.from_numpy(depth_norm).unsqueeze(0).unsqueeze(0).float().cuda()
disp = depth_t * divergence_px * 2 / self.w
grid = self.base_grid.clone()
grid[..., 0] += disp.squeeze()
warped = torch.nn.functional.grid_sample(img_t, grid, padding_mode='border', align_corners=True)
warped_np = (warped.squeeze().permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
return warped_np
def load_model(self, model):
model_id = f"depth-anything/Depth-Anything-V2-{model}-hf"
self.proc = AutoImageProcessor.from_pretrained(model_id)
self.model = AutoModelForDepthEstimation.from_pretrained(model_id).cuda().half()
def execute(self, images, output_path, model="Large", fps=24, depth_size=518,
divergence=2.5, depth_blur=6, warp_batch_size=64, depth_batch_size=64,
fill_technique="polylines_soft", stereo_offset_exponent=1.0, convergence_point=0.5):
def execute(self, images,output_path, model="Large",fps=24, depth_size=518, divergence=2.5, depth_blur=6,warp_batch_size=64,depth_batch_size=64):
self.load_model(model)
# h, w aus erstem Bild
self.h, self.w = images.shape[1], images.shape[2]
# base_grid für Warping
gy, gx = torch.meshgrid(torch.linspace(-1,1,self.h), torch.linspace(-1,1,self.w), indexing='ij')
# base_grid für Fallback-Warp
gy, gx = torch.meshgrid(torch.linspace(-1, 1, self.h), torch.linspace(-1, 1, self.w), indexing='ij')
self.base_grid = torch.stack([gx, gy], dim=-1).unsqueeze(0).cuda()
self.write_queue = Queue(maxsize=16)
writer = Thread(target=writer_thread, args=(output_path, fps, self.w, self.h, self.write_queue))
writer.start()
imgs = [Image.fromarray((img.cpu().numpy() * 255).astype(np.uint8)) for img in images]
depths = self.depth_batch(imgs, depth_batch_size, depth_size)
self.stereo_warp_streaming(imgs, depths, divergence, warp_batch_size, depth_blur)
self.stereo_warp_polylines(imgs, depths, divergence, depth_blur, warp_batch_size,
fill_technique, stereo_offset_exponent, convergence_point)
self.write_queue.put(None)
writer.join()
return (output_path,)
# Mappings für die Registrierung der Node
NODE_CLASS_MAPPINGS = {
"DukeStereoSBS": DukeStereoSBS
}