Skip to content

Commit

Permalink
Blur with depth map
Browse files Browse the repository at this point in the history
  • Loading branch information
@chrisfreilich
chrisfreilich committed May 7, 2024
1 parent 1960650 commit daf6a2e
Show file tree
Hide file tree
Showing 2 changed files with 210 additions and 28 deletions.
8 changes: 7 additions & 1 deletion __init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from .colors import BlackAndWhite
from .colors import HueSatAdvanced, HueSat
from .colors import SolidColorRGB, SolidColorHSV, SolidColor
from .blur import MotionBlur, LensBlur, GaussianBlur
from .blur import MotionBlur, LensBlur, GaussianBlur, MotionBlurDepth, LensBlurDepth, GaussianBlurDepth

NODE_CLASS_MAPPINGS = {
"BlackAndWhite": BlackAndWhite,
Expand All @@ -21,6 +21,9 @@
"LensBlur": LensBlur,
"MotionBlur": MotionBlur,
"GaussianBlur": GaussianBlur,
"LensBlurDepth": LensBlurDepth,
"MotionBlurDepth": MotionBlurDepth,
"GaussianBlurDepth": GaussianBlurDepth,
"MergeRGB": MergeRGB,
"SplitRGB": SplitRGB,
"SelectiveColor": SelectiveColor,
Expand All @@ -41,6 +44,9 @@
"LensBlur": "Lens Blur",
"MotionBlur":"Motion Blur",
"GaussianBlur": "Gaussian Blur",
"LensBlurDepth": "Lens Blur with Depth Map",
"MotionBlurDepth":"Motion Blur with Depth Map",
"GaussianBlurDepth": "Gaussian Blur with Depth Map",
"MergeRGB": "Merge RGB",
"SplitRGB": "Split RGB",
"SelectiveColor": "Selective Color",
Expand Down
230 changes: 203 additions & 27 deletions blur.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,33 +118,209 @@ def INPUT_TYPES(s):

def do_blur(self, image, amount):
return blur(image, "gaussian", amount=amount)

def blur(image, type, **kwargs):
if image.shape[3] == 4:
original_alpha = image[:, :, :, 3]
image = image[:, :, :, :3]
else:
original_alpha = None
class MotionBlurDepth:

def __init__(self):
pass

@classmethod
def INPUT_TYPES(s):

blurred_images = []
for img in image:
img_cv2 = img.cpu().numpy()
img_cv2 = (img_cv2 * 255).astype(np.uint8)
img_cv2 = cv2.cvtColor(img_cv2, cv2.COLOR_RGB2BGR)
if type == "gaussian":
blurred_img_cv2 = gaussian_blur(img_cv2, kwargs["amount"])
elif type == "lens":
blurred_img_cv2 = lens_blur(img_cv2, kwargs["radius"], kwargs["components"], kwargs["exposure_gamma"])
elif type == "motion":
blurred_img_cv2 = motion_blur(img_cv2, kwargs["size"], kwargs["angle"])
blurred_img = cv2.cvtColor(blurred_img_cv2, cv2.COLOR_BGR2RGB)
blurred_img = torch.from_numpy(blurred_img.astype(np.float32) / 255.0).to(image.device)
blurred_images.append(blurred_img)

final_tensor = torch.stack(blurred_images)
return {
"required": {
"image": ("IMAGE",),
"depth_map": ("IMAGE",),
"angle": ("INT", {
"default": 0.0,
"min": 0.0,
"max": 360.0,
"step": 0.1,
"round": 0.01,
"display": "number"}),
"num_layers": ("INT", {
"default": 10,
"min": 1,
"max": 100,
"step": 1,
"round": 1,
"display": "number"}),
"min_blur": ("INT", {
"default": 1,
"min": 1,
"max": 100,
"step": 1,
"round": 1,
"display": "number"}),
"max_blur": ("INT", {
"default": 100,
"min": 1,
"max": 100,
"step": 1,
"round": 1,
"display": "number"})
}
}

RETURN_TYPES = ("IMAGE",)
FUNCTION = "do_blur"
CATEGORY = "Virtuoso/Blur"

def do_blur(self, image, depth_map, angle, num_layers, min_blur, max_blur):
return blur(image, "motion_depth", depth_map=depth_map, angle=angle, num_layers=num_layers, min_blur=min_blur, max_blur=max_blur)

class LensBlurDepth:

def __init__(self):
pass

@classmethod
def INPUT_TYPES(s):

if original_alpha is not None:
final_tensor = torch.cat((final_tensor, original_alpha.unsqueeze(3)), dim=3)
return {
"required": {
"image": ("IMAGE",),
"depth_map": ("IMAGE",),
"components": ("INT", {
"default": 4,
"min": 1,
"max": 6,
"step": 1,
"round": 1,
"display": "number"}),
"exposure_gamma": ("FLOAT", {
"default": 2,
"min": -100,
"max": 100,
"step": 0.01,
"round": 0.01,
"display": "number"}),
"num_layers": ("INT", {
"default": 10,
"min": 1,
"max": 100,
"step": 1,
"round": 1,
"display": "number"}),
"min_blur": ("INT", {
"default": 1,
"min": 1,
"max": 100,
"step": 1,
"round": 1,
"display": "number"}),
"max_blur": ("INT", {
"default": 100,
"min": 1,
"max": 100,
"step": 1,
"round": 1,
"display": "number"})
}
}

RETURN_TYPES = ("IMAGE",)
FUNCTION = "do_blur"
CATEGORY = "Virtuoso/Blur"

def do_blur(self, image, depth_map, components, exposure_gamma, num_layers, min_blur, max_blur):
return blur(image, "lens_depth", depth_map=depth_map, components=components, exposure_gamma=exposure_gamma, num_layers=num_layers, min_blur=min_blur, max_blur=max_blur)

class GaussianBlurDepth:

def __init__(self):
pass

@classmethod
def INPUT_TYPES(s):

return (final_tensor, )
return {
"required": {
"image": ("IMAGE",),
"depth_map": ("IMAGE",),
"sigma": ("INT", {
"default": 5,
"min": 1,
"max": 100,
"step": 1,
"round": 1,
"display": "number"}),
"num_layers": ("INT", {
"default": 10,
"min": 1,
"max": 100,
"step": 1,
"round": 1,
"display": "number"}),
"min_blur": ("INT", {
"default": 1,
"min": 1,
"max": 100,
"step": 1,
"round": 1,
"display": "number"}),
"max_blur": ("INT", {
"default": 100,
"min": 1,
"max": 100,
"step": 1,
"round": 1,
"display": "number"})
}
}

RETURN_TYPES = ("IMAGE",)
FUNCTION = "do_blur"
CATEGORY = "Virtuoso/Blur"

def do_blur(self, image, depth_map, sigma, num_layers, min_blur, max_blur):
return blur(image, "gaussian_depth", depth_map=depth_map, sigma=sigma, num_layers=num_layers, min_blur=min_blur, max_blur=max_blur)


def blur(image, type, depth_map=None, **kwargs):

if type in ["lens_depth", "gaussian_depth", "motion_depth"]:
if depth_map.shape[3] == 4:
depth_map = depth_map[:, :, :, :3] # Remove alpha channel if it exists

# If batch size of depth_map does not match with image, use the first depth_map for all images
if depth_map.shape[0] != image.shape[0]:
depth_map = depth_map[0].unsqueeze(0).repeat(image.shape[0], 1, 1, 1)

if image.shape[3] == 4:
original_alpha = image[:, :, :, 3]
image = image[:, :, :, :3]
else:
original_alpha = None

blurred_images = []
for i, img in enumerate(image):
img_cv2 = img.cpu().numpy()
img_cv2 = (img_cv2 * 255).astype(np.uint8)
img_cv2 = cv2.cvtColor(img_cv2, cv2.COLOR_RGB2BGR)

if type == "gaussian":
blurred_img_cv2 = gaussian_blur(img_cv2, kwargs["amount"])
elif type == "lens":
blurred_img_cv2 = lens_blur(img_cv2, kwargs["radius"], kwargs["components"], kwargs["exposure_gamma"])
elif type == "motion":
blurred_img_cv2 = motion_blur(img_cv2, kwargs["size"], kwargs["angle"])
else: #depth map blur
depth_map_img = depth_map[i].cpu().numpy()
depth_map_img = (depth_map_img * 255).astype(np.uint8)
if type == "motion_depth":
blurred_img_cv2 = motion_blur_with_depth_map(img_cv2, depth_map=depth_map_img, angle=kwargs["angle"], num_layers=kwargs["num_layers"], min_blur=kwargs["min_blur"], max_blur=kwargs["max_blur"])
elif type == "lens_depth":
blurred_img_cv2 = lens_blur_with_depth_map(img_cv2, depth_map=depth_map_img, components=kwargs["components"], exposure_gamma=kwargs["exposure_gamma"], num_layers=kwargs["num_layers"], min_blur=kwargs["min_blur"], max_blur=kwargs["max_blur"])
elif type == "gaussian_depth":
blurred_img_cv2 = gaussian_blur_with_depth_map(img_cv2, depth_map=depth_map_img, sigma=kwargs["sigma"], num_layers=kwargs["num_layers"], min_blur=kwargs["min_blur"], max_blur=kwargs["max_blur"])
blurred_img = cv2.cvtColor(blurred_img_cv2, cv2.COLOR_BGR2RGB)
blurred_img = torch.from_numpy(blurred_img.astype(np.float32) / 255.0).to(image.device)
blurred_images.append(blurred_img)

final_tensor = torch.stack(blurred_images)

if original_alpha is not None:
final_tensor = torch.cat((final_tensor, original_alpha.unsqueeze(3)), dim=3)

return (final_tensor, )

0 comments on commit daf6a2e

Please sign in to comment.