85 lines
3.6 KiB
Python
85 lines
3.6 KiB
Python
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
import base64
|
|
import io
|
|
import folder_paths
|
|
import os
|
|
import random
|
|
|
|
class BasicDrawingBoard:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image_data": ("STRING", {"default": ""}),
|
|
"width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 64}),
|
|
"height": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 64}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",),
|
|
"auto_mask": ("MASK",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
RETURN_NAMES = ("IMAGE", "MASK")
|
|
FUNCTION = "process_drawing"
|
|
CATEGORY = "CustomNodes/Drawing"
|
|
|
|
def process_drawing(self, image_data, width, height, image=None, auto_mask=None):
|
|
temp_filename = ""
|
|
auto_mask_filename = ""
|
|
|
|
# 1. Image de fond
|
|
if image is not None:
|
|
i = 255. * image[0].cpu().numpy()
|
|
base_img = Image.fromarray(np.uint8(i)).convert("RGBA")
|
|
base_img = base_img.resize((width, height), Image.Resampling.LANCZOS)
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
temp_filename = f"drawing_bg_{random.randint(0, 1000000)}.png"
|
|
base_img.save(os.path.join(temp_dir, temp_filename))
|
|
else:
|
|
base_img = Image.new("RGBA", (width, height), (0, 0, 0, 255))
|
|
|
|
# 2. FUSION DES MASQUES (Correction pour SAM Multi-détection)
|
|
if auto_mask is not None:
|
|
# Si SAM envoie plusieurs masques, on les fusionne en un seul
|
|
# auto_mask a souvent une forme (N, H, W). On prend le max sur la dimension 0.
|
|
if auto_mask.shape[0] > 1:
|
|
combined_mask = torch.max(auto_mask, dim=0)[0].unsqueeze(0)
|
|
else:
|
|
combined_mask = auto_mask
|
|
|
|
m = 255. * combined_mask[0].cpu().numpy()
|
|
mask_img = Image.fromarray(np.uint8(m)).convert("L")
|
|
mask_img = mask_img.resize((width, height), Image.Resampling.LANCZOS)
|
|
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
auto_mask_filename = f"auto_mask_{random.randint(0, 1000000)}.png"
|
|
mask_img.save(os.path.join(temp_dir, auto_mask_filename))
|
|
|
|
# On met à jour auto_mask pour la sortie du nœud également
|
|
auto_mask = combined_mask
|
|
|
|
# 3. Fusion avec le dessin manuel
|
|
if image_data and image_data != "":
|
|
header, encoded = image_data.split(",", 1)
|
|
data = base64.b64decode(encoded)
|
|
drawing_img = Image.open(io.BytesIO(data)).convert("RGBA")
|
|
drawing_img = drawing_img.resize((width, height), Image.Resampling.LANCZOS)
|
|
base_img.alpha_composite(drawing_img)
|
|
|
|
final_mask_img = drawing_img.split()[-1]
|
|
final_mask_np = np.array(final_mask_img).astype(np.float32) / 255.0
|
|
mask_tensor = torch.from_numpy(np.where(final_mask_np > 0, 1.0, 0.0))[None,]
|
|
else:
|
|
mask_tensor = auto_mask if auto_mask is not None else torch.zeros((1, height, width))
|
|
|
|
final_img = base_img.convert("RGB")
|
|
img_tensor = torch.from_numpy(np.array(final_img).astype(np.float32) / 255.0)[None,]
|
|
|
|
return {"ui": {"bg_file": [temp_filename], "auto_mask": [auto_mask_filename]}, "result": (img_tensor, mask_tensor)}
|
|
|
|
NODE_CLASS_MAPPINGS = {"BasicDrawingBoard": BasicDrawingBoard}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"BasicDrawingBoard": "🎨 Basic Drawing Board"} |