349 lines
15 KiB
Python
349 lines
15 KiB
Python
import torch
|
|
import numpy as np
|
|
import os
|
|
import sys
|
|
from typing import Union, List
|
|
from PIL import Image
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from typing import Union, List
|
|
from PIL import Image, ImageFilter, ImageEnhance, ImageOps, ImageDraw, ImageChops, ImageFont
|
|
|
|
file_dir = os.path.dirname(__file__)
|
|
sys.path.append(file_dir)
|
|
|
|
import stereoimage_generation as sig
|
|
|
|
from comfy.utils import ProgressBar
|
|
|
|
def tensor2np(tensor: torch.Tensor) -> np.ndarray:
|
|
if tensor.dim() == 4: # Batch of images
|
|
tensor = tensor[0] # Assuming we take the first image in the batch
|
|
np_array = tensor.cpu().numpy()
|
|
np_array = np.clip(255.0 * np_array, 0, 255).astype(np.uint8)
|
|
if np_array.shape[0] == 3: # Convert from (3, H, W) to (H, W, 3)
|
|
np_array = np_array.transpose(1, 2, 0)
|
|
return np_array
|
|
|
|
def np2tensor(img_np: Union[np.ndarray, List[np.ndarray]]) -> torch.Tensor:
|
|
if isinstance(img_np, list):
|
|
return torch.cat([np2tensor(img) for img in img_np], dim=0)
|
|
return torch.from_numpy(img_np.astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
class StereoImageNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"depth_map": ("IMAGE",),
|
|
"modes": (["left-right", "right-left", "top-bottom", "bottom-top", "red-cyan-anaglyph"],),
|
|
"fill_technique": ([
|
|
'No fill', 'No fill - Reverse projection', 'Imperfect fill - Hybrid Edge', 'Fill - Naive',
|
|
'Fill - Naive interpolating', 'Fill - Polylines Soft', 'Fill - Polylines Sharp',
|
|
'Fill - Post-fill', 'Fill - Reverse projection with Post-fill', 'Fill - Hybrid Edge with fill'
|
|
], {"default": "Fill - Polylines Soft"}),
|
|
#"return_basic_mask": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
"divergence": ("FLOAT", {"default": 3.5, "min": 0.05, "max": 15, "step": 0.01}),
|
|
"separation": ("FLOAT", {"default": 0, "min": -5, "max": 5, "step": 0.01}),
|
|
"stereo_balance": ("FLOAT", {"default": 0, "min": -0.95, "max": 0.95, "step": 0.05}),
|
|
"stereo_offset_exponent": ("FLOAT", {"default": 2, "min": 1, "max": 2, "step": 1}),
|
|
"depth_blur_sigma": ("FLOAT", {"default": 0, "min": 0, "max": 10, "step": 0.1}),
|
|
"depth_blur_edge_threshold": ("FLOAT", {"default": 40, "min": 0.1, "max": 100, "step": 0.1})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "MASK")
|
|
RETURN_NAMES = ("stereoscope", "modified_depthmap_left", "modified_depthmap_right", "no_fill_imperfect_mask")
|
|
FUNCTION = "generate"
|
|
|
|
def generate(self, image, depth_map, divergence, separation, modes,
|
|
stereo_balance, stereo_offset_exponent, fill_technique, depth_blur_sigma, depth_blur_edge_threshold):# return_basic_mask):
|
|
|
|
fill_technique_mapping = {
|
|
'No fill': 'none',
|
|
'No fill - Reverse projection': 'inverse',
|
|
'Imperfect fill - Hybrid Edge': 'hybrid_edge',
|
|
'Fill - Naive': 'naive',
|
|
'Fill - Naive interpolating': 'naive_interpolating',
|
|
'Fill - Polylines Soft': 'polylines_soft',
|
|
'Fill - Polylines Sharp': 'polylines_sharp',
|
|
'Fill - Post-fill': 'none_post',
|
|
'Fill - Reverse projection with Post-fill': 'inverse_post',
|
|
'Fill - Hybrid Edge with fill': 'hybrid_edge_plus'
|
|
}
|
|
|
|
fill_technique = fill_technique_mapping.get(fill_technique, 'none')
|
|
|
|
results_final = []
|
|
modified_depthmap_final_left = []
|
|
modified_depthmap_final_right = []
|
|
mask_final = []
|
|
total_steps = len(image)
|
|
pbar = ProgressBar(total_steps)
|
|
|
|
for i in range(len(image)):
|
|
img = tensor2np(image[i:i+1])
|
|
dm = tensor2np(depth_map[i:i+1])
|
|
|
|
if len(dm.shape) == 3 and dm.shape[2] == 3:
|
|
dm = np.dot(dm[..., :3], [0.2989, 0.5870, 0.1140]).astype(np.uint8)
|
|
|
|
if img.shape[:2] != dm.shape:
|
|
dm = np.array(Image.fromarray(dm).resize((img.shape[1], img.shape[0])))
|
|
|
|
output = sig.create_stereoimages(img, dm, divergence, separation,
|
|
[modes], stereo_balance, stereo_offset_exponent,
|
|
fill_technique, depth_blur_sigma, depth_blur_edge_threshold)
|
|
|
|
if len(output) == 3:
|
|
results, left_modified_depthmap, right_modified_depthmap = output
|
|
else:
|
|
results, modified_depthmap = output
|
|
left_modified_depthmap = modified_depthmap
|
|
right_modified_depthmap = modified_depthmap
|
|
|
|
results_final.append(convertResult(results))
|
|
modified_depthmap_final_left.append(convertResult(left_modified_depthmap))
|
|
modified_depthmap_final_right.append(convertResult(right_modified_depthmap))
|
|
|
|
mask = self.generate_mask(results)
|
|
mask_final.append(mask)
|
|
|
|
pbar.update(1)
|
|
|
|
return (torch.cat(results_final), torch.cat(modified_depthmap_final_left), torch.cat(modified_depthmap_final_right), torch.cat(mask_final))
|
|
|
|
def generate_mask(self, stereoscope_img):
|
|
np_img = np.array(stereoscope_img)
|
|
mask = (np_img.sum(axis=-1) == 0).astype(np.uint8) * 255 # Black areas become white in mask
|
|
return np2tensor(mask)
|
|
|
|
|
|
|
|
def convertResult(results):
|
|
# Save the result for debugging
|
|
if isinstance(results, list):
|
|
for idx, result in enumerate(results):
|
|
result = np.array(result)
|
|
#Image.fromarray(result).save(f"debug_result_{idx}.png")
|
|
else:
|
|
results = np.array(results)
|
|
#Image.fromarray(results).save("debug_result.png")
|
|
|
|
# Convert the result to a tensor
|
|
if isinstance(results, list):
|
|
results = np.array(results[0])
|
|
else:
|
|
results = np.array(results)
|
|
|
|
# Ensure the results are in the correct shape (H, W, C)
|
|
if len(results.shape) == 2: # Convert grayscale to RGB
|
|
results = np.stack([results]*3, axis=-1)
|
|
|
|
results_tensor = np2tensor(results)
|
|
|
|
return results_tensor
|
|
|
|
|
|
def tensor2cv2(image: torch.Tensor) -> List[tuple[np.ndarray, tuple]]:
|
|
original_shape = image.shape
|
|
images = []
|
|
|
|
if image.dim() == 4: # Batch of images
|
|
for i in range(image.shape[0]):
|
|
img = image[i]
|
|
if img.shape[0] == 3: # If in [C, H, W] format
|
|
img = img.permute(1, 2, 0)
|
|
npimage = (img * 255).byte().cpu().numpy()
|
|
if npimage.shape[-1] == 3:
|
|
images.append((cv2.cvtColor(npimage, cv2.COLOR_RGB2BGR), original_shape))
|
|
elif npimage.shape[-1] == 1:
|
|
images.append((npimage.squeeze(-1), original_shape)) # For grayscale images
|
|
else:
|
|
images.append((npimage, original_shape)) # Return as-is for other channel configurations
|
|
else: # Single image
|
|
if image.shape[0] == 3: # If in [C, H, W] format
|
|
image = image.permute(1, 2, 0)
|
|
npimage = (image * 255).byte().cpu().numpy()
|
|
if npimage.shape[-1] == 3:
|
|
images.append((cv2.cvtColor(npimage, cv2.COLOR_RGB2BGR), original_shape))
|
|
elif npimage.shape[-1] == 1:
|
|
images.append((npimage.squeeze(-1), original_shape)) # For grayscale images
|
|
else:
|
|
images.append((npimage, original_shape)) # Return as-is for other channel configurations
|
|
|
|
return images
|
|
|
|
def cv22tensor(cv2_imgs: List[tuple[np.ndarray, tuple]]) -> torch.Tensor:
|
|
tensors = []
|
|
for cv2_img, original_shape in cv2_imgs:
|
|
if cv2_img.ndim == 2: # grayscale
|
|
tensor = torch.from_numpy(cv2_img).float() / 255.0
|
|
tensor = tensor.unsqueeze(-1) # Add channel dimension
|
|
else: # color
|
|
tensor = torch.from_numpy(cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB)).float() / 255.0
|
|
|
|
if len(original_shape) == 4:
|
|
tensor = tensor.permute(2, 0, 1).unsqueeze(0) # [H, W, C] -> [1, C, H, W]
|
|
elif len(original_shape) == 3 and original_shape[0] == 3:
|
|
tensor = tensor.permute(2, 0, 1) # [H, W, C] -> [C, H, W]
|
|
|
|
tensors.append(tensor)
|
|
|
|
return torch.stack(tensors)
|
|
|
|
def merge_channels(red_channel, green_channel, blue_channel):
|
|
if red_channel.dim() == 4: # Batch of images
|
|
images = []
|
|
for i in range(red_channel.shape[0]):
|
|
image = mix_rgb_channels(tensor2pil(red_channel[i]).convert('L'), tensor2pil(
|
|
green_channel[i]).convert('L'), tensor2pil(blue_channel[i]).convert('L'))
|
|
images.append(pil2tensor(image).squeeze(0)) # Remove batch dimension added by pil2tensor
|
|
return (torch.stack(images),)
|
|
else: # Single image
|
|
image = mix_rgb_channels(tensor2pil(red_channel).convert('L'), tensor2pil(
|
|
green_channel).convert('L'), tensor2pil(blue_channel).convert('L'))
|
|
return (pil2tensor(image),)
|
|
|
|
def mix_rgb_channels(red, green, blue):
|
|
# Create an empty image with the same size as the channels
|
|
width, height = red.size
|
|
merged_img = Image.new('RGB', (width, height))
|
|
|
|
# Merge the channels into the new image
|
|
merged_img = Image.merge('RGB', (red, green, blue))
|
|
|
|
return merged_img
|
|
|
|
# PIL to Tensor
|
|
def pil2tensor(image):
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
# Tensor to PIL
|
|
def tensor2pil(image):
|
|
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
|
|
# class LazyStereo:
|
|
# @classmethod
|
|
# def INPUT_TYPES(s):
|
|
# return {"required": {
|
|
# "image": ("IMAGE",),
|
|
# "depth_map": ("IMAGE",),
|
|
# "shift_amount": ("INT", {
|
|
# "default": 10,
|
|
# "min": 5, #Minimum value
|
|
# "max": 200, #Maximum value
|
|
# "step": 1, #Slider's step
|
|
# "display": "number" # Cosmetic only: display as "number" or "slider"
|
|
# }),
|
|
# "mode": (["Cross-eyed", "Parallel"],),
|
|
# },
|
|
# }
|
|
|
|
# RETURN_TYPES = ("IMAGE",)
|
|
# FUNCTION = "generate_cross_eyed_image"
|
|
|
|
# def generate_cross_eyed_image(self, image, depth_map, shift_amount=10, mode="Cross-eyed"):
|
|
# images = tensor2cv2(image)
|
|
# depth_maps = tensor2cv2(depth_map)
|
|
|
|
# results = []
|
|
# total_steps = len(image) # Total number of images to process
|
|
# pbar = ProgressBar(total_steps)
|
|
|
|
# for (img, original_img_shape), (depth_map, original_depth_shape) in zip(images, depth_maps):
|
|
# if len(depth_map.shape) == 3 and depth_map.shape[2] > 1:
|
|
# depth_map = cv2.cvtColor(depth_map, cv2.COLOR_BGR2GRAY)
|
|
# elif len(depth_map.shape) == 2:
|
|
# # It's already grayscale, no need to convert
|
|
# pass
|
|
# else:
|
|
# raise ValueError(f"Unexpected depth map shape: {depth_map.shape}")
|
|
|
|
# if img is None or depth_map is None:
|
|
# raise ValueError("Image or depth map not found.")
|
|
|
|
# height, width = img.shape[:2]
|
|
|
|
# # Calculate gradients to determine high-contrast regions
|
|
# grad_x = cv2.Sobel(depth_map, cv2.CV_64F, 1, 0, ksize=3)
|
|
# grad_y = cv2.Sobel(depth_map, cv2.CV_64F, 0, 1, ksize=3)
|
|
# gradient_magnitude = np.sqrt(grad_x**2 + grad_y**2)
|
|
# max_gradient = np.max(gradient_magnitude)
|
|
# gradient_magnitude = gradient_magnitude / max_gradient # Normalize to [0, 1]
|
|
|
|
# # Initialize the left and right images and masks for inpainting
|
|
# left_image = np.zeros_like(img)
|
|
# right_image = np.zeros_like(img)
|
|
# left_mask = np.zeros((height, width), dtype=np.uint8)
|
|
# right_mask = np.zeros((height, width), dtype=np.uint8)
|
|
|
|
# # Generate the left and right images by shifting pixels according to the depth map
|
|
# for y in range(height):
|
|
# for x in range(width):
|
|
# depth = depth_map[y, x]
|
|
# local_gradient = gradient_magnitude[y, x]
|
|
# adaptive_shift = shift_amount * (1 - local_gradient) # Reduce shift in high-contrast areas
|
|
# shift = int((depth / 255.0) * adaptive_shift)
|
|
# if x - shift >= 0:
|
|
# left_image[y, x] = img[y, x - shift]
|
|
# else:
|
|
# left_image[y, x] = img[y, 0]
|
|
# if x + shift < width:
|
|
# right_image[y, x] = img[y, x + shift]
|
|
# else:
|
|
# right_image[y, x] = img[y, width - 1]
|
|
|
|
# # Update masks for inpainting
|
|
# if x - shift < 0 or x + shift >= width:
|
|
# left_mask[y, x] = 255
|
|
# right_mask[y, x] = 255
|
|
|
|
# # Apply edge detection to create a more precise inpainting mask
|
|
# edges = cv2.Canny(depth_map, 100, 200)
|
|
# left_mask = cv2.bitwise_or(left_mask, edges)
|
|
# right_mask = cv2.bitwise_or(right_mask, edges)
|
|
|
|
# # Inpaint the left and right images to fill in gaps
|
|
# left_image = cv2.inpaint(left_image, left_mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
|
|
# right_image = cv2.inpaint(right_image, right_mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
|
|
|
|
# if mode == "Cross-eyed":
|
|
# # Combine the left and right images side by side
|
|
# cross_eyed_image = np.hstack((right_image, left_image))
|
|
# else:
|
|
# cross_eyed_image = np.hstack((left_image, right_image))
|
|
|
|
# # Separate the channels
|
|
# red_channel = cross_eyed_image[:, :, 0]
|
|
# green_channel = cross_eyed_image[:, :, 1]
|
|
# blue_channel = cross_eyed_image[:, :, 2]
|
|
|
|
# # Convert channels back to tensors
|
|
# red_tensor = cv22tensor([(red_channel, (1, 1, height, width))])
|
|
# green_tensor = cv22tensor([(green_channel, (1, 1, height, width))])
|
|
# blue_tensor = cv22tensor([(blue_channel, (1, 1, height, width))])
|
|
|
|
# merged_image = merge_channels(blue_tensor, green_tensor, red_tensor)
|
|
|
|
# results.append(merged_image[0])
|
|
# pbar.update(1)
|
|
|
|
# # Concatenate the results into a single tensor
|
|
# return (torch.cat(results),)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
# "LazyStereo": LazyStereo,
|
|
"StereoImageNode": StereoImageNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
# "LazyStereo": "LazyStereo",
|
|
"StereoImageNode": "Stereo Image Node"
|
|
}
|