125 lines
4.1 KiB
Python
125 lines
4.1 KiB
Python
"""
|
|
Region Crop Node - Nodo che taglia l'immagine secondo le coordinate selezionate
|
|
"""
|
|
|
|
import json
|
|
import torch
|
|
|
|
|
|
class RegionCropNode:
|
|
"""
|
|
Nodo che taglia un'immagine usando le coordinate fornite dal RegionSelectorNode.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Inizializza il nodo"""
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""
|
|
Definisce gli input del nodo.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"box_metadata": ("STRING", {
|
|
"default": "",
|
|
"multiline": True,
|
|
}),
|
|
},
|
|
"optional": {
|
|
"fallback_mode": (["use_full_image", "return_zero", "error"], {
|
|
"default": "use_full_image",
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("cropped_image",)
|
|
FUNCTION = "crop_image"
|
|
CATEGORY = "image/region"
|
|
|
|
def crop_image(self, image, box_metadata, fallback_mode="use_full_image"):
|
|
"""
|
|
Taglia l'immagine secondo le coordinate nel metadata.
|
|
|
|
Args:
|
|
image: Tensor immagine (B, H, W, C)
|
|
box_metadata: JSON string con coordinate della regione
|
|
fallback_mode: Cosa fare se non ci sono coordinate valide
|
|
|
|
Returns:
|
|
(cropped_image,): Immagine ritagliata
|
|
"""
|
|
|
|
try:
|
|
metadata = json.loads(box_metadata) if box_metadata.strip() else {}
|
|
except json.JSONDecodeError:
|
|
metadata = {}
|
|
|
|
# Estrai coordinate (x1, x2, y1, y2 sono i lati)
|
|
x1 = metadata.get("x1", None)
|
|
x2 = metadata.get("x2", None)
|
|
y1 = metadata.get("y1", None)
|
|
y2 = metadata.get("y2", None)
|
|
|
|
# Se non c'è selezione
|
|
if x1 is None or x2 is None or y1 is None or y2 is None:
|
|
if fallback_mode == "use_full_image":
|
|
return (image,)
|
|
elif fallback_mode == "return_zero":
|
|
return (torch.zeros_like(image),)
|
|
else:
|
|
raise ValueError("No region coordinates provided in metadata")
|
|
|
|
# Estrai il fattore di scala se la preview era stata scalata
|
|
display_scale_factor = metadata.get("displayScaleFactor", 1.0)
|
|
if display_scale_factor and display_scale_factor != 1.0:
|
|
# Se le coordinate sono state prese da una preview scalata,
|
|
# dividi per il fattore di scala per ottenere le coordinate originali
|
|
x1 = x1 / display_scale_factor
|
|
x2 = x2 / display_scale_factor
|
|
y1 = y1 / display_scale_factor
|
|
y2 = y2 / display_scale_factor
|
|
print(f"[BoxCrop] Scale factor detected: {display_scale_factor}x. Adjusted coordinates.")
|
|
|
|
# Assicura che le coordinate siano integer
|
|
x1 = int(round(x1))
|
|
x2 = int(round(x2))
|
|
y1 = int(round(y1))
|
|
y2 = int(round(y2))
|
|
|
|
# Estrai dimensioni immagine
|
|
batch_size, img_height, img_width, channels = image.shape
|
|
|
|
# Assicura che x1 < x2 e y1 < y2
|
|
x_start = min(x1, x2)
|
|
x_end = max(x1, x2)
|
|
y_start = min(y1, y2)
|
|
y_end = max(y1, y2)
|
|
|
|
# Calcola coordinate finali con clipping ai bordi
|
|
x_start = max(0, min(x_start, img_width - 1))
|
|
x_end = max(x_start + 1, min(x_end, img_width))
|
|
y_start = max(0, min(y_start, img_height - 1))
|
|
y_end = max(y_start + 1, min(y_end, img_height))
|
|
|
|
# Verifica validità della region
|
|
if x_end <= x_start or y_end <= y_start:
|
|
if fallback_mode == "use_full_image":
|
|
return (image,)
|
|
elif fallback_mode == "return_zero":
|
|
return (torch.zeros_like(image),)
|
|
else:
|
|
raise ValueError("Invalid region coordinates: no overlap with image")
|
|
|
|
# Taglia l'immagine
|
|
cropped = image[:, y_start:y_end, x_start:x_end, :]
|
|
|
|
return (cropped,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"BoxCrop": RegionCropNode}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"BoxCrop": "✂️ BoxCrop"}
|