From df897fab294e0bcfda5a39724bfbd9fd9dd8f4e5 Mon Sep 17 00:00:00 2001 From: BennyKok Date: Thu, 5 Oct 2023 18:32:22 +0800 Subject: [PATCH] fix vector and add ImageAlphaMaskMerge --- blender/b_vector3d.py | 6 ++-- blender/b_vector4d.py | 6 ++-- blender/image_alpha_mask_merge.py | 31 +++++++++++++++++ blender/load_image_with_alpha.py | 55 +++++++++++++++++++++++++++++++ blender/mesh_utils.py | 5 ++- blender/save_image_exclude.py | 2 +- 6 files changed, 97 insertions(+), 8 deletions(-) create mode 100644 blender/image_alpha_mask_merge.py create mode 100644 blender/load_image_with_alpha.py diff --git a/blender/b_vector3d.py b/blender/b_vector3d.py index 3be4b83..a0327b1 100644 --- a/blender/b_vector3d.py +++ b/blender/b_vector3d.py @@ -5,17 +5,17 @@ class VECTOR3D: "required": { "x": ("FLOAT", { "default": 0, - "step": 0.1, + "step": 0.01, "display": "number" }), "y": ("FLOAT", { "default": 0, - "step": 0.1, + "step": 0.01, "display": "number" }), "z": ("FLOAT", { "default": 0, - "step": 0.1, + "step": 0.01, "display": "number" }), }, diff --git a/blender/b_vector4d.py b/blender/b_vector4d.py index 46ec2c4..c00f0c6 100644 --- a/blender/b_vector4d.py +++ b/blender/b_vector4d.py @@ -5,17 +5,17 @@ class VECTOR4D: "required": { "x": ("FLOAT", { "default": 0, - "step": 0.1, + "step": 0.01, "display": "number" }), "y": ("FLOAT", { "default": 0, - "step": 0.1, + "step": 0.01, "display": "number" }), "z": ("FLOAT", { "default": 0, - "step": 0.1, + "step": 0.01, "display": "number" }), "u": ("FLOAT", { diff --git a/blender/image_alpha_mask_merge.py b/blender/image_alpha_mask_merge.py new file mode 100644 index 0000000..d1f57ec --- /dev/null +++ b/blender/image_alpha_mask_merge.py @@ -0,0 +1,31 @@ +import folder_paths +import os +from PIL import Image, ImageOps +from PIL.PngImagePlugin import PngInfo +import numpy as np +import torch +import json + +class ImageAlphaMaskMerge: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"image": ("IMAGE",) , + "mask": ("MASK",) }, + } + + CATEGORY = "image" + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "load_image" + def load_image(self, image, mask): + image = torch.cat((image, 1 - mask.unsqueeze(0).unsqueeze(3)), dim=3) + return (image, ) + +NODE_CLASS_MAPPINGS = { + "ImageAlphaMaskMerge": ImageAlphaMaskMerge, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageAlphaMaskMerge": "ImageAlphaMaskMerge" +} \ No newline at end of file diff --git a/blender/load_image_with_alpha.py b/blender/load_image_with_alpha.py new file mode 100644 index 0000000..60704e7 --- /dev/null +++ b/blender/load_image_with_alpha.py @@ -0,0 +1,55 @@ +import folder_paths +import os +from PIL import Image, ImageOps +from PIL.PngImagePlugin import PngInfo +import numpy as np +import torch +import json + +class LoadImageWithAlpha: + @classmethod + def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] + return {"required": + {"image": (sorted(files), {"image_upload": True})}, + } + + CATEGORY = "image" + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "load_image" + def load_image(self, image): + image_path = folder_paths.get_annotated_filepath(image) + i = Image.open(image_path) + i = ImageOps.exif_transpose(i) + image = i.convert("RGBA") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + + print(image.shape) + + return (image, ) + + @classmethod + def IS_CHANGED(s, image): + image_path = folder_paths.get_annotated_filepath(image) + m = hashlib.sha256() + with open(image_path, 'rb') as f: + m.update(f.read()) + return m.digest().hex() + + @classmethod + def VALIDATE_INPUTS(s, image): + if not folder_paths.exists_annotated_filepath(image): + return "Invalid image file: {}".format(image) + + return True + +NODE_CLASS_MAPPINGS = { + "LoadImageWithAlpha": LoadImageWithAlpha, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LoadImageWithAlpha": "LoadImageWithAlpha" +} \ No newline at end of file diff --git a/blender/mesh_utils.py b/blender/mesh_utils.py index 1a5d3cb..2849eed 100644 --- a/blender/mesh_utils.py +++ b/blender/mesh_utils.py @@ -86,7 +86,7 @@ def assign_texture(bpy, BPY_OBJ, texture, texture_name): # Create an image with the required dimensions img = bpy.data.images.new( - texture_name, width=texture.shape[1], height=texture.shape[0]) + texture_name, width=texture.shape[1], height=texture.shape[0], alpha = True) # If there is no alpha channel, append one full of 1's if texture.shape[2] == 3: @@ -118,6 +118,7 @@ def assign_texture(bpy, BPY_OBJ, texture, texture_name): # Create a material mat = bpy.data.materials.new("MaterialName") mat.use_nodes = True + mat.blend_method = 'BLEND' nodes = mat.node_tree.nodes for node in nodes: nodes.remove(node) @@ -138,6 +139,8 @@ def assign_texture(bpy, BPY_OBJ, texture, texture_name): texture_node.outputs['Color']) links.new(output_node.inputs['Surface'], bsdf_node.outputs['BSDF']) + links.new(bsdf_node.inputs['Alpha'], texture_node.outputs['Alpha']) + # Assign the material to the active object if obj.data.materials: obj.data.materials[0] = mat diff --git a/blender/save_image_exclude.py b/blender/save_image_exclude.py index 7709061..862caed 100644 --- a/blender/save_image_exclude.py +++ b/blender/save_image_exclude.py @@ -33,7 +33,7 @@ class SaveImageWithWorkflow: image_path = folder_paths.get_annotated_filepath(image) i = Image.open(image_path) i = ImageOps.exif_transpose(i) - image = i.convert("RGB") + image = i.convert("RGBA") image = np.array(image).astype(np.float32) / 255.0 image = torch.from_numpy(image)[None,] images=(image)