virtualPainter/paintModifier.py
2025-10-18 01:25:52 +02:00

84 lines
3.4 KiB
Python

import bpy
from .constants import image_name, image_width, image_height
import numpy as np
from .libs.PIL import Image
import mathutils
from timeit import default_timer as timer
def apply_color():
start = timer()
context = bpy.context
paint_buffer = bpy.data.images.get(image_name)
normal_buffer = context.scene.normal_texture_settings.texture
custom_collection = context.scene.texture_collection_settings
normal_color = context.scene.normal_color_global
returnValue = -1
if paint_buffer:
# apply normal color
if normal_buffer and normal_color != mathutils.Vector((1.0, 1.0, 1.0)):
if paint_buffer.size[0] == normal_buffer.size[0] and paint_buffer.size[1] == normal_buffer.size[1]:
# create arrays
paint_mask_pixels = np.empty(len(paint_buffer.pixels), dtype=np.float32)
paint_buffer.pixels.foreach_get(paint_mask_pixels)
paint_mask_pixels = (paint_mask_pixels * 255).astype(np.uint8)
target_pixels = np.empty(len(normal_buffer.pixels), dtype=np.float32)
normal_buffer.pixels.foreach_get(target_pixels)
target_pixels = (target_pixels * 255).astype(np.uint8)
new_pixels = np.tile((np.array([*normal_color, 1.0], dtype=np.float32) * 255), (image_width * image_height, 1)).astype(np.uint8)
# Convert arrays to PIL Images
paint_mask_image = Image.frombuffer('RGBA', (image_width, image_height), paint_mask_pixels, 'raw', 'RGBA', 0, 1)
target_image = Image.frombuffer('RGBA', (image_width, image_height), target_pixels, 'raw', 'RGBA', 0, 1)
new_image = Image.frombuffer('RGBA', (image_width, image_height), new_pixels, 'raw', 'RGBA', 0, 1)
r_mask, g_mask, b_mask, a_mask = paint_mask_image.split() # separates the channels
paint_mask_image = Image.merge("RGBA", (r_mask, g_mask, b_mask, r_mask)) # replace alpha with red
# Blend buffers
blended_image = Image.composite(new_image, target_image, paint_mask_image)
# Convert blended image back to numpy array and normalize
blended_pixels = np.array(blended_image, dtype=np.uint8).astype(np.float32) / 255.0
# Save texture
normal_buffer.pixels.foreach_set(blended_pixels.ravel())
normal_buffer.update()
print("Color applied")
returnValue = 0
else:
print("Texture sizes don't match")
returnValue = -1
else:
print("no normal buffer selected")
returnValue = -1
# reset paint buffer
reset_pixels = np.empty(len(paint_buffer.pixels), dtype=np.float32)
reset_pixels[0::4] = 0 # R
reset_pixels[1::4] = 0 # G
reset_pixels[2::4] = 0 # B
reset_pixels[3::4] = 1 # A
paint_buffer.pixels.foreach_set(reset_pixels)
paint_buffer.update()
else:
print("Texture 'vPainter_buffer' not found")
returnValue = -1
end = timer()
print(f"Time taken to apply color: {end - start:.6f} seconds")
return returnValue
def register():
bpy.utils.register_class(VPAINTER_OT_apply_color)
def unregister():
bpy.utils.unregister_class(VPAINTER_OT_apply_color)