[Metrics][Added] Support for lists and save individual results

- We can process a list of batches
- We can now save the metrics for each image
This commit is contained in:
Salvador E. Tropea
2025-11-09 16:20:24 -03:00
parent 650646a5f2
commit 295681fbcb
3 changed files with 191 additions and 109 deletions
+3 -2
View File
@@ -9,7 +9,7 @@ def get_e_measure(
gt: torch.Tensor,
num_thresholds: int = 255,
chunk_size: int = 16
) -> Tuple[float, float, float, torch.Tensor]:
) -> Tuple[float, float, float, torch.Tensor, torch.Tensor]:
"""
Calculates the E-measure scores using a memory-efficient chunking strategy.
@@ -30,6 +30,7 @@ def get_e_measure(
- float: The maximum E-measure score across all thresholds.
- float: The adaptive E-measure score.
- torch.Tensor: A 1D tensor with the E-measure score for each threshold.
- torch.Tensor: A 1D tensor with the thresholds.
"""
# 1. --- Calculate scores for all thresholds using chunking ---
@@ -71,7 +72,7 @@ def get_e_measure(
scores.mean().item(),
scores.max().item(),
adaptive_score_tensor.item(),
scores
scores, thlist
)
+4 -1
View File
@@ -26,6 +26,7 @@ def get_f_measure(pred: torch.Tensor, gt: torch.Tensor, beta2: float = 0.3) -> f
"""
# Initialize f_max to store the highest F-measure score found so far.
f_max = 0.0
f = []
# Iterate through 256 evenly spaced thresholds from 0.0 to 1.0.
# This corresponds to testing every possible 8-bit grayscale value as the cutoff.
@@ -42,6 +43,7 @@ def get_f_measure(pred: torch.Tensor, gt: torch.Tensor, beta2: float = 0.3) -> f
# Optimization: If there are no true positives, the F-measure will be 0.
# We can skip the rest of the calculations for this threshold.
if tp == 0:
f.append((threshold, 0))
continue
# Calculate Precision = TP / (TP + FP).
@@ -55,6 +57,7 @@ def get_f_measure(pred: torch.Tensor, gt: torch.Tensor, beta2: float = 0.3) -> f
# Calculate the F-beta score using the computed precision and recall.
# The beta^2=0.3 value is standard in saliency detection literature.
f_beta = (1 + beta2) * precision * recall / (beta2 * precision + recall + EPS)
f.append((threshold, f_beta))
# Update f_max if the F-beta score for the current threshold is the highest yet.
# .item() extracts the single float value from the 0-dimensional tensor.
@@ -62,7 +65,7 @@ def get_f_measure(pred: torch.Tensor, gt: torch.Tensor, beta2: float = 0.3) -> f
f_max = f_beta.item()
# After checking all thresholds, return the maximum score found.
return f_max
return f_max, f
def get_weighted_f_measure(pred: torch.Tensor, gt: torch.Tensor, beta2: float = 0.3) -> float:
+184 -106
View File
@@ -112,6 +112,8 @@ SAVE_WORKFLOW = ("BOOLEAN", {
SHOW_PREVIEW = ("BOOLEAN", {
"default": True,
"tooltip": "Show a preview of the images"})
SOD_NAMES = {'mae': "MAE", 'max_f_mes': "Max F-measure", 's_mes': "S-measure", 'e_mes': "E-measure mean",
'wf_mes': "Weighted F-measure"}
# A dictionary to cache loaded fonts
font_cache = {}
@@ -148,6 +150,16 @@ def parse_size(size_str, reference_dim):
return 0
def send_progress_text(unique_id, msg):
if unique_id and PromptServer is not None:
try:
PromptServer.instance.send_progress_text(msg, unique_id)
except Exception:
pass
else:
logger.info(msg)
# Define sort methods for the node input
sort_methods = [
"None",
@@ -670,107 +682,193 @@ class SaliencyEvaluationMetrics:
"prediction": ("MASK",),
"ground_truth": ("MASK",),
},
"hidden": {
"unique_id": "UNIQUE_ID",
},
"optional": {
"mae_enable": ("BOOLEAN", {
"default": True,
"tooltip": "Compute the MAE"}),
"max_f_mes_enable": ("BOOLEAN", {
"default": True,
"tooltip": "Compute the Max_F-measure"}),
"s_mes_enable": ("BOOLEAN", {
"default": True,
"tooltip": "Compute the S-measure"}),
"e_mes_enable": ("BOOLEAN", {
"default": True,
"tooltip": "Compute the E-measure"}),
"wf_mes_enable": ("BOOLEAN", {
"default": True,
"tooltip": "Compute the Weighted F-measure"}),
"img_name": ("STRING", {"forceInput" = True, "tooltip": "Name used as base to save the parameters"}),
"result_save": ("BOOLEAN", {"default": False, "tooltip": "Save computed values to IMG_NAME.csv"}),
"mae_enable": ("BOOLEAN", {"default": True, "tooltip": "Compute the MAE"}),
"mae_save": ("BOOLEAN", {"default": False, "tooltip": "Save the MAE using IMG_NAME_MAE.csv"}),
"max_f_mes_enable": ("BOOLEAN", {"default": True, "tooltip": "Compute the Max_F-measure"}),
"max_f_mes_save": ("BOOLEAN", {"default": False, "tooltip": "Save the F-measure using IMG_NAME_F.csv"}),
"s_mes_enable": ("BOOLEAN", {"default": True, "tooltip": "Compute the S-measure"}),
"s_mes_save": ("BOOLEAN", {"default": False, "tooltip": "Save the S-measure using IMG_NAME_S.csv"}),
"e_mes_enable": ("BOOLEAN", {"default": True, "tooltip": "Compute the E-measure"}),
"e_mes_save": ("BOOLEAN", {"default": False, "tooltip": "Save the E-measure using IMG_NAME_E.csv"}),
"wf_mes_enable": ("BOOLEAN", {"default": True, "tooltip": "Compute the Weighted F-measure"}),
"wf_mes_save": ("BOOLEAN", {"default": False, "tooltip": "Save the Weighted F-measure using IMG_NAME_wF.csv"}),
},
}
RETURN_TYPES = ("DICT", "FLOAT", "FLOAT", "FLOAT", "FLOAT", "FLOAT")
RETURN_NAMES = ("all", "MAE", "Max_F-measure", "S-measure", "E-measure", "Weighted_F-measure")
OUTPUT_IS_LIST = (True, False, False, False, False, False)
RETURN_TYPES = ("DICT", "STRING", "FLOAT", "FLOAT", "FLOAT", "FLOAT", "FLOAT")
RETURN_NAMES = ("all", "img_name", "MAE", "Max_F-measure", "S-measure", "E-measure", "Weighted_F-measure")
OUTPUT_IS_LIST = (True, True, False, False, False, False, False)
INPUT_IS_LIST = True
FUNCTION = "evaluate"
CATEGORY = BASE_CATEGORY + "/" + "Analysis"
UNIQUE_NAME = "SET_SaliencyEvaluationMetrics"
DISPLAY_NAME = "Saliency Evaluation Metrics"
def evaluate(self, prediction: torch.Tensor, ground_truth: torch.Tensor, mae_enable: bool = True,
max_f_mes_enable: bool = True, s_mes_enable: bool = True, e_mes_enable: bool = True,
wf_mes_enable: bool = True):
# Ensure tensors are on the same device and float type
def evaluate(self, prediction: torch.Tensor, ground_truth: torch.Tensor, unique_id,
img_name, result_save: bool = False, mae_enable: bool = True, mae_save: bool = False,
max_f_mes_enable: bool = True, max_f_mes_save: bool = False, s_mes_enable: bool = True,
s_mes_save: bool = False, e_mes_enable: bool = True, e_mes_save: bool = False,
wf_mes_enable: bool = True, wf_mes_save: bool = True):
# Flatten arguments that aren't really expected to be lists
mae_enable = mae_enable[0]
mae_save = mae_save[0]
max_f_mes_enable = max_f_mes_enable[0]
max_f_mes_save = max_f_mes_save[0]
s_mes_enable = s_mes_enable[0]
s_mes_save = s_mes_save[0]
e_mes_enable = e_mes_enable[0]
e_mes_save = e_mes_save[0]
wf_mes_enable = wf_mes_enable[0]
wf_mes_save = wf_mes_save[0]
unique_id = unique_id[0]
device = get_default_comfy_device()
inputs_are_copies = get_canonical_device(prediction.device) != device
gt = ground_truth.to(device)
pred = prediction.to(device)
# Match batch sizes
batch_size = min(pred.shape[0], gt.shape[0])
pred = pred[:batch_size]
gt = gt[:batch_size]
# Ensure we have lists of the same length
gt_len = len(ground_truth)
pred_len = len(prediction)
if gt_len != pred_len:
raise ValueError(f"Got {pred_len} predictions and {gt_len} ground thruths, they must match")
# Ensure masks are normalized to [0, 1] range
pred = batched_min_max_norm(pred, in_place=inputs_are_copies)
gt = batched_min_max_norm(gt, in_place=inputs_are_copies)
# --- Initialize accumulators for metrics ---
mae_total, f_measure_max_total, s_measure_total, e_measure_total, weighted_f_total = 0, 0, 0, 0, 0
# Find how many images we have
imgs_len = sum((i.shape[0] for i in prediction))
# Initialize accumulators for metrics
mae_total = f_measure_max_total = s_measure_total = e_measure_total = weighted_f_total = 0
e_measure_max_total = e_measure_adp_total = 0
all = []
for i in range(batch_size):
pred_i = pred[i]
gt_i = gt[i]
res = {}
# 1. Mean Absolute Error (MAE)
if mae_enable:
mae = torch.mean(torch.abs(pred_i - gt_i)).item()
logger.debug(f"MAE: {mae}")
mae_total += mae
res['mae'] = mae
# Names counter
index_name = 0
names_len = len(img_name)
if names_len != imgs_len:
raise ValueError(f"Got {imgs_len} images and {names_len} names, they must match")
# --- Metrics requiring binary ground truth ---
if max_f_mes_enable or s_mes_enable or e_mes_enable or wf_mes_enable:
gt_binary = (gt_i >= 0.5).float()
for index_img in range(pred_len):
# Ensure tensors are on the same device
inputs_are_copies = get_canonical_device(prediction[index_img].device) != device
gt = ground_truth[index_img].to(device)
pred = prediction[index_img].to(device)
# 2. Max F-measure
if max_f_mes_enable:
f_max = get_f_measure(pred_i, gt_binary)
f_measure_max_total += f_max
logger.debug(f"F_max: {f_max}")
res['max_f_mes'] = f_max
# Ensure masks are normalized to [0, 1] range
gt = batched_min_max_norm(gt, in_place=inputs_are_copies)
pred = batched_min_max_norm(pred, in_place=inputs_are_copies)
# 3. S-measure
if s_mes_enable:
s_measure = get_s_measure(pred_i, gt_binary)
s_measure_total += s_measure
logger.debug(f"S: {s_measure}")
res['s_mes'] = s_measure
for i in range(gt.shape[0]):
# Get the next name
imgp = Path(img_name[index_name])
index_name += 1
logger.debug(f"{index_name}) {imgp.name}")
# 4. E-measure
if e_mes_enable:
e_mean, e_max, e_adp, _ = get_e_measure(pred_i, gt_binary)
e_measure_total += e_mean
logger.debug(f"E: {e_mean} {e_max} {e_adp}")
res['e_mes'] = e_mean
pred_i = pred[i]
gt_i = gt[i]
res = {}
# 5. Weighted F-measure
if wf_mes_enable:
wf = get_weighted_f_measure(pred_i, gt_binary)
weighted_f_total += wf
logger.debug(f"wF: {wf}")
res['wf_mes'] = wf
# 1. Mean Absolute Error (MAE)
if mae_enable:
mae = torch.mean(torch.abs(pred_i - gt_i)).item()
logger.debug(f"MAE: {mae}")
mae_total += mae
res['mae'] = mae
if mae_save:
with open(Path(imgp.parent, imgp.stem+"_mae.csv"), "wt") as f:
f.write(f"MAE\n{mae}")
all.append(res)
# --- Metrics requiring binary ground truth ---
if max_f_mes_enable or s_mes_enable or e_mes_enable or wf_mes_enable:
gt_binary = (gt_i >= 0.5).float()
# --- Average metrics over the batch ---
mae_avg = mae_total / batch_size
f_measure_avg = f_measure_max_total / batch_size
s_measure_avg = s_measure_total / batch_size
e_measure_avg = e_measure_total / batch_size
weighted_f_avg = weighted_f_total / batch_size
# 2. Max F-measure
if max_f_mes_enable:
f_max, all_f = get_f_measure(pred_i, gt_binary)
f_measure_max_total += f_max
logger.debug(f"Fβmax: {f_max}")
res['max_f_mes'] = f_max
if max_f_mes_save:
with open(Path(imgp.parent, imgp.stem+"_F.csv"), "wt") as f:
f.write("Threshold, F-measure\n")
for fn in all_f:
f.write(f"{fn[0]}, {fn[1]}\n")
f.write(f"\nMax, {f_max}\n")
# 3. S-measure
if s_mes_enable:
s_measure = get_s_measure(pred_i, gt_binary)
s_measure_total += s_measure
logger.debug(f"Sα: {s_measure}")
res['s_mes'] = s_measure
if s_mes_save:
with open(Path(imgp.parent, imgp.stem+"_S.csv"), "wt") as f:
f.write(f"S-measure\n{s_measure}")
# 4. E-measure
if e_mes_enable:
e_mean, e_max, e_adp, all_e, thres = get_e_measure(pred_i, gt_binary)
e_measure_total += e_mean
e_measure_max_total += e_max
e_measure_adp_total += e_adp
logger.debug(f"Eϕ: {e_mean} {e_max} {e_adp}")
res['e_mes'] = e_mean
if e_mes_save:
with open(Path(imgp.parent, imgp.stem+"_E.csv"), "wt") as f:
f.write("Threshold, E-measure\n")
for index, en in enumerate(all_e):
f.write(f"{thres[index]}, {en}\n")
f.write("\n")
f.write(f"Mean, {e_mean}\n")
f.write(f"Max, {e_max}\n")
f.write(f"Adaptive, {e_adp}\n")
# 5. Weighted F-measure
if wf_mes_enable:
wf = get_weighted_f_measure(pred_i, gt_binary)
weighted_f_total += wf
logger.debug(f"Fβw: {wf}")
res['wf_mes'] = wf
if wf_mes_save:
with open(Path(imgp.parent, imgp.stem+"_wF.csv"), "wt") as f:
f.write(f"Weighted F-measure\n{wf}")
if result_save and res:
with open(Path(imgp.parent, imgp.stem+".csv"), "wt") as f:
f.write(','.join([SOD_NAMES[v] for v in res.keys()])+"\n")
f.write(','.join([str(v) for v in res.values()])+"\n")
all.append(res)
# Average metrics over the batch/es
mae_avg = mae_total / gt_len
f_measure_avg = f_measure_max_total / gt_len
s_measure_avg = s_measure_total / gt_len
e_measure_avg = e_measure_total / gt_len
e_measure_max_avg = e_measure_max_total / gt_len
e_measure_adp_avg = e_measure_adp_total / gt_len
weighted_f_avg = weighted_f_total / gt_len
# Show results in the node
msg = "<table>"
if mae_enable:
msg += f"<tr><td>MAE</td><td>{mae_avg:.4f}</td></tr>"
if max_f_mes_enable:
msg += f"<tr><td>Fβmax</td><td>{f_measure_avg:.4f}</td></tr>"
if s_mes_enable:
msg += f"<tr><td>Sα</td><td>{s_measure_avg:.4f}</td></tr>"
if e_mes_enable:
msg += f"<tr><td>Eϕmean</td><td>{e_measure_avg:.4f}</td></tr>"
msg += f"<tr><td>Eϕmax</td><td>{e_measure_max_avg:.4f}</td></tr>"
msg += f"<tr><td>Eϕadp</td><td>{e_measure_adp_avg:.4f}</td></tr>"
if wf_mes_enable:
msg += f"<tr><td>Fβw</td><td>{weighted_f_avg:.4f}</td></tr>"
msg += "</table>"
send_progress_text(unique_id, msg)
logger.warning(unique_id)
logger.warning(msg)
return (all, mae_avg, f_measure_avg, s_measure_avg, e_measure_avg, weighted_f_avg)
@@ -1746,20 +1844,8 @@ class ImageResize:
if mask is not None:
mask_chunks.append(sub_out_mask.cpu() if sub_out_mask is not None else None)
# Per-batch progress update
if unique_id and PromptServer is not None:
try:
PromptServer.instance.send_progress_text(
f"<tr><td>Resize Image</td><td>batch {current_batch}/{total_batches} · images {end_idx}/{B}"
"</td></tr>",
unique_id
)
except Exception:
pass
else:
try:
logger.info(f"batch {current_batch}/{total_batches} · images {end_idx}/{B}")
except Exception:
pass
send_progress_text(unique_id, f"<tr><td>Resize Image</td><td>batch {current_batch}/{total_batches}"
" · images {end_idx}/{B}</td></tr>")
out_image = torch.cat(chunks, dim=0)
if mask is not None and any(m is not None for m in mask_chunks):
out_mask = torch.cat([m for m in mask_chunks if m is not None], dim=0)
@@ -1767,19 +1853,11 @@ class ImageResize:
out_mask = None
# Progress UI
if unique_id and PromptServer is not None:
try:
num_elements = out_image.numel()
element_size = out_image.element_size()
memory_size_mb = (num_elements * element_size) / (1024 * 1024)
PromptServer.instance.send_progress_text(
f"<tr><td>Output: </td><td><b>{out_image.shape[0]}</b> x <b>{out_image.shape[2]}</b> x <b>"
f"{out_image.shape[1]} | {memory_size_mb:.2f}MB</b></td></tr>",
unique_id
)
except Exception:
pass
num_elements = out_image.numel()
element_size = out_image.element_size()
memory_size_mb = (num_elements * element_size) / (1024 * 1024)
send_progress_text(unique_id, f"<tr><td>Output: </td><td><b>{out_image.shape[0]}</b> x <b>{out_image.shape[2]}"
f"</b> x <b>{out_image.shape[1]} | {memory_size_mb:.2f} MiB</b></td></tr>")
return (out_image.cpu(), out_image.shape[2], out_image.shape[1],
out_mask.cpu() if out_mask is not None else