Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0caf7f7177 | ||
|
|
7e5f7c3c0a | ||
|
|
9ccef62bc6 | ||
|
|
0b8bf65b57 | ||
|
|
e72ea99a6e | ||
|
|
7cfe734e54 | ||
|
|
fecff541e6 | ||
|
|
4d23ae490a | ||
|
|
415513c1a7 | ||
|
|
dbb20b3713 | ||
|
|
e105455bc7 | ||
|
|
fb99aa0a65 | ||
|
|
726cbdba51 | ||
|
|
70e753333f | ||
|
|
a33906ede8 | ||
|
|
b24e9687c0 | ||
|
|
d5ecc8b32d | ||
|
|
23831bc8af | ||
|
|
45cb1339ce | ||
|
|
d02c46e15c | ||
|
|
aa545d16f9 | ||
|
|
31debcce47 | ||
|
|
8cc5abf0fd | ||
|
|
2393f347d4 | ||
|
|
f3a1105d8d | ||
|
|
4b5e5ff766 | ||
|
|
636a5e8bbd | ||
|
|
eb69bacdea | ||
|
|
8b1b6431ef | ||
|
|
6bc0a90de4 | ||
|
|
3476ef8d14 | ||
|
|
202fe0d7ab | ||
|
|
3053c89b25 | ||
|
|
4cb9bf3d81 | ||
|
|
d4ba71ae5a | ||
|
|
5c2230b078 | ||
|
|
f59f7b7d9f | ||
|
|
a34689b27c | ||
|
|
4ab32569c6 | ||
|
|
83264cccab | ||
|
|
64eafc6402 | ||
|
|
295681fbcb | ||
|
|
650646a5f2 | ||
|
|
a57fe0cba3 | ||
|
|
a2ca2fe876 | ||
|
|
4d03617fc6 | ||
|
|
50cd6c67e9 | ||
|
|
31577cb58c | ||
|
|
afeda6db47 | ||
|
|
a58773147c | ||
|
|
e40b1f76a6 | ||
|
|
81bf74ec1f | ||
|
|
9d411bd074 | ||
|
|
5265d27b93 | ||
|
|
a3a9872bf3 |
@@ -0,0 +1,90 @@
|
||||
### Documentation of Saliency Evaluation Metrics
|
||||
|
||||
This document details the quantitative metrics used to evaluate the performance of Salient Object Detection (SOD) and Dichotomous Image Segmentation (DIS) models. These metrics compare the model's continuous prediction map (`prediction`) with a binary ground truth map (`ground_truth`).
|
||||
|
||||
---
|
||||
|
||||
### 1. Mean Absolute Error (MAE)
|
||||
|
||||
* **What it Measures**: MAE computes the average pixel-wise absolute difference between the normalized prediction mask and the ground truth mask. It provides a direct, straightforward measure of the overall error.
|
||||
|
||||
* **Formula**:
|
||||
$$ \text{MAE} = \frac{1}{W \times H} \sum_{x=1}^{W} \sum_{y=1}^{H} | P(x,y) - G(x,y) | $$
|
||||
where `P` is the prediction, `G` is the ground truth, and `W, H` are the dimensions of the mask.
|
||||
|
||||
* **Interpretation**:
|
||||
* **Lower is better**. A score of **0** indicates a perfect pixel-for-pixel match. A higher score indicates a greater average error across the entire image.
|
||||
|
||||
* **Relevance and Justification**: MAE is a fundamental metric that penalizes all errors equally, regardless of their location (e.g., at the center of an object vs. its boundary). While its simplicity is a strength, this lack of perceptual awareness is also its primary limitation. It serves as an excellent baseline for measuring raw prediction accuracy. Its use in saliency evaluation was solidified in early benchmark papers that sought a simple way to quantify prediction error.
|
||||
|
||||
* **Bibliographic Citation**:
|
||||
> Perazzi, F., Krähenbühl, P., Pritch, Y., & Hornung, A. (2012). "Saliency filters: Contrast based filtering for salient region detection." In *Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR)*, pp. 733-740.
|
||||
|
||||
---
|
||||
|
||||
### 2. Maximum F-measure (Max Fβ)
|
||||
|
||||
* **What it Measures**: The F-measure is the harmonic mean of Precision and Recall, providing a score that balances the two. In saliency evaluation, the continuous prediction map is converted to a binary map using a series of thresholds (from 0 to 255). The F-measure is calculated for each threshold, and the **maximum** value obtained across all thresholds is reported. This adaptive thresholding makes the metric robust to models that produce well-shaped but poorly-calibrated (e.g., generally too dark or bright) saliency maps.
|
||||
|
||||
* **Interpretation**:
|
||||
* **Higher is better**. A score of **1** represents a perfect balance of precision and recall at the optimal threshold.
|
||||
|
||||
* **Relevance and Justification**: Unlike the pixel-level MAE, the F-measure is region-based. It evaluates how well the *shape* of the predicted salient region aligns with the ground truth. By finding the optimal threshold for a given prediction, it fairly assesses the quality of the saliency map's structure, forgiving issues with overall intensity. The standard beta-squared value ($β^2$) is set to **0.3** to weigh precision more heavily than recall, as proposed by the authors of the foundational paper below.
|
||||
|
||||
* **Bibliographic Citation**:
|
||||
> Achanta, R., Hemami, S., Estrada, F., & Süsstrunk, S. (2009). "Frequency-tuned salient region detection." In *Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR)*, pp. 1597-1604.
|
||||
|
||||
---
|
||||
|
||||
### 3. S-measure (Structure-measure, Sα)
|
||||
|
||||
* **What it Measures**: The S-measure is a more advanced metric designed to overcome the limitations of pixel- and region-based metrics by evaluating structural similarity. It combines two components: an **object-aware structural similarity** (evaluating the similarity between the predicted foreground and the ground truth object) and a **region-aware structural similarity** (evaluating the similarity between sub-regions of the maps).
|
||||
|
||||
* **Interpretation**:
|
||||
* **Range**:.
|
||||
* **Higher is better**. A score of **1** indicates a perfect structural match, meaning the prediction captures the true form of the salient object much better than a simple "blob" would.
|
||||
|
||||
* **Relevance and Justification**: The S-measure was introduced because previous metrics could give high scores to predictions that correctly located a salient object but failed to capture its detailed structure. This metric is less sensitive to minor pixel errors and more attuned to the geometric and structural correctness of the prediction, making it highly aligned with human visual perception.
|
||||
|
||||
* **Bibliographic Citation**:
|
||||
> Fan, D. P., Cheng, M. M., Liu, Y., Li, T., & Borji, A. (2017). "Structure-measure: A new way to evaluate saliency maps." In *Proceedings of the IEEE international conference on computer vision (ICCV)*, pp. 4548-4557.
|
||||
|
||||
---
|
||||
|
||||
### 4. E-measure (Enhanced-alignment measure, Eξ)
|
||||
|
||||
* **What it Measures**: The E-measure simultaneously captures both image-level statistics (the global mean) and local pixel-matching information. It computes an "enhanced alignment matrix" that quantifies the relationship between the prediction and ground truth, considering both their values and their statistical context.
|
||||
|
||||
* **Interpretation**:
|
||||
* **Range**:.
|
||||
* **Higher is better**. A high score indicates strong alignment in both the fine-grained pixel details and the overall statistical distribution of saliency values.
|
||||
|
||||
* **Relevance and Justification**: The E-measure provides a more holistic evaluation than F-measure or MAE. It was designed to be robust and to capture a wider range of similarities and differences that might be missed by other metrics. Its ability to combine local and global aspects makes it a powerful tool for comprehensively assessing a model's performance.
|
||||
|
||||
* **Bibliographic Citation**:
|
||||
> Fan, D. P., Gong, C., Cao, Y., Ren, B., Cheng, M. M., & Borji, A. (2018). "Enhanced-alignment measure for binary foreground map evaluation." In *Proceedings of the International Joint Conference on Artificial Intelligence (IJCAI)*, pp. 698-704.
|
||||
|
||||
---
|
||||
|
||||
### 5. Weighted F-measure (Fβw)
|
||||
|
||||
* **What it Measures**: This is a modification of the F-measure that introduces a weighting scheme to address a key flaw: not all pixels are equally important. It assigns higher importance to pixels near the center of a salient region and less importance to those near the boundary. Errors made on perceptually important pixels are therefore penalized more heavily.
|
||||
|
||||
* **Interpretation**:
|
||||
* **Range**:.
|
||||
* **Higher is better**. A high score indicates high precision and recall, especially in the most critical areas of the salient object.
|
||||
|
||||
* **Relevance and Justification**: The motivation behind the Weighted F-measure is to create a metric that better correlates with human perception of error. A small error in the middle of an object is often more jarring than an error along its noisy edge. By using a weighting map (often based on centrality or distance from the boundary), this metric provides a more perceptually accurate evaluation of a saliency map's quality.
|
||||
|
||||
* **Bibliographic Citation**:
|
||||
> Margolin, R., Zelnik-Manor, L., & Tal, A. (2014). "How to evaluate foreground maps?" In *Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR)*, pp. 248-255.
|
||||
|
||||
### Summary Table
|
||||
|
||||
| Metric | What it Measures | Range | Goal | Key Strength |
|
||||
| :--- | :--- | :--- | :--- | :--- |
|
||||
| **MAE** | Average pixel-wise error | | **Minimize** | Simple, direct error measurement. |
|
||||
| **Max F-measure** | Optimal balance of precision and recall | | **Maximize** | Robust to intensity calibration issues; region-based. |
|
||||
| **S-measure** | Structural and object-level similarity | | **Maximize** | Evaluates the geometric correctness of the prediction. |
|
||||
| **E-measure** | Global and local pixel alignment | | **Maximize** | Provides a holistic score combining local and global stats. |
|
||||
| **Weighted F-measure** | Perceptually weighted precision and recall | | **Maximize** | Aligns better with human perception by weighting errors. |
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
+2
-1
@@ -1,3 +1,4 @@
|
||||
Pillow
|
||||
numpy
|
||||
seconohe
|
||||
scipy
|
||||
seconohe >= 1.0.7
|
||||
|
||||
@@ -7,3 +7,7 @@ from seconohe.logger import initialize_logger
|
||||
__version__ = "1.4.0"
|
||||
NODES_NAME = "ImageMisc"
|
||||
main_logger = initialize_logger(NODES_NAME)
|
||||
|
||||
# Epsilon: small value to avoid "divide by 0" errors
|
||||
EPS = 1e-8
|
||||
F_POINTS = 255
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
import torch
|
||||
from typing import Tuple
|
||||
from . import EPS, F_POINTS
|
||||
|
||||
|
||||
def get_e_measure(
|
||||
pred: torch.Tensor,
|
||||
gt: torch.Tensor,
|
||||
num_thresholds: int = F_POINTS,
|
||||
chunk_size: int = -1
|
||||
) -> Tuple[float, float, float, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Calculates the E-measure scores using a memory-efficient chunking strategy.
|
||||
|
||||
This implementation is fully vectorized within chunks to maintain high performance
|
||||
while ensuring that memory usage remains low and predictable, making it suitable
|
||||
for high-resolution images.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): The continuous prediction mask (normalized to [0, 1]).
|
||||
gt (torch.Tensor): The binary ground truth mask (values are 0 or 1).
|
||||
num_thresholds (int): The number of thresholds to evaluate. Defaults to F_POINTS.
|
||||
chunk_size (int): The number of thresholds to process in a single batch.
|
||||
Lower this value if you encounter VRAM issues. Defaults to 16.
|
||||
|
||||
Returns:
|
||||
A tuple containing:
|
||||
- float: The mean E-measure score across all thresholds.
|
||||
- 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 ---
|
||||
|
||||
# Create a 1D tensor of thresholds from 0 to almost 1
|
||||
thlist = torch.linspace(0, 1 - 1e-10, num_thresholds, device=pred.device)
|
||||
all_scores = []
|
||||
|
||||
if chunk_size == -1:
|
||||
chunk_size = round(64 / (pred.shape[0] * pred.shape[1] / (1 << 20)))
|
||||
chunk_size = min(max(1, chunk_size), num_thresholds)
|
||||
|
||||
# Process thresholds in memory-efficient chunks
|
||||
for i in range(0, num_thresholds, chunk_size):
|
||||
# Get the current chunk of thresholds
|
||||
th_chunk = thlist[i:i + chunk_size]
|
||||
|
||||
# Vectorized operation on the smaller chunk.
|
||||
# Reshape pred to (1, H, W) and thlist to (N, 1, 1).
|
||||
# Broadcasting (>=) creates a binarized prediction for each threshold.
|
||||
# The result `binarized_preds` has a shape of (chunk_size, H, W).
|
||||
binarized_preds_chunk = (pred.unsqueeze(0) >= th_chunk.view(-1, 1, 1)).to(pred.dtype)
|
||||
|
||||
# Calculate scores for the current chunk
|
||||
scores_chunk = e_calculate_enhanced_scores(binarized_preds_chunk, gt)
|
||||
all_scores.append(scores_chunk)
|
||||
|
||||
# Combine the scores from all chunks
|
||||
scores = torch.cat(all_scores)
|
||||
|
||||
# 2. --- Calculate the single adaptive score ---
|
||||
|
||||
# Calculate the adaptive threshold, clamping at 1.0 to be safe.
|
||||
adaptive_th = torch.clamp(2 * pred.mean(), max=1.0)
|
||||
# Binarize the prediction with this single threshold
|
||||
adaptive_pred_binarized = (pred >= adaptive_th).to(pred.dtype)
|
||||
# Reuse the same helper function by adding a temporary batch dimension
|
||||
adaptive_score_tensor = e_calculate_enhanced_scores(adaptive_pred_binarized.unsqueeze(0), gt)
|
||||
|
||||
# 3. --- Return the final results ---
|
||||
|
||||
return (
|
||||
scores.mean().item(),
|
||||
scores.max().item(),
|
||||
adaptive_score_tensor.item(),
|
||||
scores.cpu(), thlist
|
||||
)
|
||||
|
||||
|
||||
def e_calculate_enhanced_scores(binarized_preds: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Helper function to compute E-measure scores for a batch of binarized predictions.
|
||||
This function is fully vectorized.
|
||||
|
||||
Args:
|
||||
binarized_preds (torch.Tensor): A tensor of binarized predictions, with shape
|
||||
(N, H, W), where N is the number of thresholds.
|
||||
gt (torch.Tensor): The single binary ground truth mask, with shape (H, W).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: A 1D tensor of shape (N,) containing the E-measure score
|
||||
for each binarized prediction.
|
||||
"""
|
||||
# Handle the edge case where the ground truth is all black
|
||||
if torch.mean(gt) == 0.0:
|
||||
# The score is based on how much of the prediction is also black.
|
||||
# Original formula: sum(1 - y_pred_th) / (numel - 1)
|
||||
# This is equivalent to numel * (1 - mean) / (numel - 1)
|
||||
enhanced = 1 - binarized_preds
|
||||
|
||||
# Handle the edge case where the ground truth is all white
|
||||
elif torch.mean(gt) == 1.0:
|
||||
# The score is based on how much of the prediction is also white.
|
||||
# Original formula: sum(y_pred_th) / (numel - 1)
|
||||
enhanced = binarized_preds
|
||||
|
||||
# Normal case with a mixed ground truth
|
||||
else:
|
||||
# Demean the ground truth. `gt_demeaned` has shape (H, W).
|
||||
gt_demeaned = gt - gt.mean()
|
||||
|
||||
# Demean the binarized predictions.
|
||||
# `mean` is calculated over spatial dims (H, W), keeping the threshold dim.
|
||||
# `fm` (foreground map) will have shape (N, H, W).
|
||||
fm = binarized_preds - binarized_preds.mean(dim=[-2, -1], keepdim=True)
|
||||
|
||||
# The demeaned GT will be broadcasted to match the shape of `fm`.
|
||||
align_matrix = 2 * gt_demeaned * fm / (gt_demeaned.square() + fm.square() + EPS)
|
||||
enhanced = (align_matrix + 1).square() / 4
|
||||
|
||||
# Calculate the final score for each threshold by summing over the spatial dimensions.
|
||||
# The denominator (y.numel() - 1) is a quirk from the original paper's code.
|
||||
scores = torch.sum(enhanced, dim=[-2, -1]) / (gt.numel() - 1 + EPS)
|
||||
|
||||
return scores
|
||||
@@ -0,0 +1,195 @@
|
||||
import numpy as np
|
||||
import scipy
|
||||
import torch
|
||||
from typing import Tuple
|
||||
from . import EPS, F_POINTS
|
||||
|
||||
|
||||
def _f_measure(pred: torch.Tensor, gt: torch.Tensor, threshold: torch.Tensor, beta2: float = 0.3) -> float:
|
||||
# Binarize the continuous prediction map using the current threshold.
|
||||
# Pixels >= threshold become 1.0 (positive), and others become 0.0 (negative).
|
||||
pred_binary = (pred >= threshold).float()
|
||||
|
||||
# Calculate True Positives (TP): pixels that are positive in both the prediction and ground truth.
|
||||
# Element-wise multiplication results in 1 only where both are 1.
|
||||
tp = (pred_binary * gt).sum()
|
||||
|
||||
# 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:
|
||||
return 0.0, 0.0, 0.0
|
||||
|
||||
# Calculate Precision = TP / (TP + FP).
|
||||
# The sum of `pred_binary` gives the total number of predicted positives (TP + FP).
|
||||
precision = tp / (pred_binary.sum() + EPS)
|
||||
|
||||
# Calculate Recall = TP / (TP + FN).
|
||||
# The sum of `gt` gives the total number of actual positives (TP + FN).
|
||||
recall = tp / (gt.sum() + EPS)
|
||||
|
||||
# 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)
|
||||
|
||||
return f_beta.item(), precision.item(), recall.item()
|
||||
|
||||
|
||||
def get_f_measure(pred: torch.Tensor, gt: torch.Tensor, beta2: float = 0.3) -> Tuple[float, Tuple[float, float], float]:
|
||||
"""
|
||||
Calculates the maximum F-measure score for a continuous prediction against a binary ground truth.
|
||||
|
||||
The F-measure evaluates the balance between precision and recall. Since the prediction
|
||||
is a continuous map (0.0 to 1.0), this function iterates through F_POINTS possible
|
||||
thresholds to binarize the prediction. It calculates the F-measure for each
|
||||
threshold and returns the highest (max) score found. This provides a fair
|
||||
evaluation of the prediction's structural quality, independent of its overall brightness.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): The continuous prediction mask (normalized to [0, 1]).
|
||||
gt (torch.Tensor): The binary ground truth mask (values are 0 or 1).
|
||||
beta2 (float): The beta-squared value for the F-measure. The standard value of 0.3
|
||||
is used to weigh precision more heavily than recall. Defaults to 0.3.
|
||||
|
||||
Returns:
|
||||
float: The maximum F-measure score found across all thresholds.
|
||||
tuple: The F-measure for all the 255 thresholds we tried, (threshold, F)
|
||||
float: The F-measure for an adaptative threshold
|
||||
|
||||
"""
|
||||
# Initialize f_max to store the highest F-measure score found so far.
|
||||
f_max = 0.0
|
||||
f4 = torch.Tensor(size=(4, F_POINTS), device="cpu")
|
||||
f = f4.movedim(1, 0) # Easier here, not for computation
|
||||
|
||||
# Iterate through F_POINTS evenly spaced thresholds from 0.0 to 1.0.
|
||||
# This corresponds to testing every possible 8-bit grayscale value as the cutoff.
|
||||
# The thresholds tensor is created on the same device as the input for efficiency.
|
||||
for c, threshold in enumerate(torch.linspace(0, 1, F_POINTS, device=gt.device)):
|
||||
f_beta, prec, rec = _f_measure(pred, gt, threshold, beta2)
|
||||
f[c] = torch.Tensor((threshold.item(), f_beta, prec, rec))
|
||||
|
||||
# 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.
|
||||
if f_beta > f_max:
|
||||
f_max = f_beta
|
||||
|
||||
# Also compute F for an "adaptative" threshold
|
||||
threshold = min(pred.mean() * 2, 1)
|
||||
f_adp, _, _ = _f_measure(pred, gt, threshold, beta2)
|
||||
|
||||
# After checking all thresholds, return the maximum score found.
|
||||
return f_max, f4, f_adp
|
||||
|
||||
|
||||
def get_weighted_f_measure(pred: torch.Tensor, gt: torch.Tensor, beta2: float = 0.3) -> float:
|
||||
"""
|
||||
Calculates the true research-grade Weighted F-measure (F_beta^w).
|
||||
|
||||
This is a faithful and optimized port of the reference algorithm from the
|
||||
"How to Evaluate Foreground Maps?" paper, ensuring verifiable results.
|
||||
The implementation performs pre-checks on the GPU for efficiency before
|
||||
transferring data to the CPU for SciPy-based calculations.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): The continuous prediction mask (normalized to [0, 1]).
|
||||
gt (torch.Tensor): The binary ground truth mask (values are 0 or 1).
|
||||
beta2 (float): The beta-squared value for the F-measure. Defaults to 0.3.
|
||||
|
||||
Returns:
|
||||
float: The final Weighted F-measure score.
|
||||
"""
|
||||
# --- Step 1: GPU-side Pre-checks for Efficiency ---
|
||||
|
||||
# Handle the edge case of an all-black ground truth on the GPU.
|
||||
# If the ground truth is empty, the score is 1 minus the mean of the prediction.
|
||||
# A perfect prediction (all black) would yield a score of 1.
|
||||
# This avoids the expensive CPU transfer and SciPy calculations entirely.
|
||||
if torch.mean(gt) == 0.0:
|
||||
return (1.0 - pred.mean()).item()
|
||||
|
||||
# --- Step 2: Data Transfer to CPU for NumPy/SciPy Processing ---
|
||||
|
||||
# Move tensors to the CPU and convert to NumPy arrays. SciPy functions require this.
|
||||
# Squeezing removes any singleton channel dimensions (e.g., from [1, H, W] to [H, W]).
|
||||
gt_np = gt.squeeze().cpu().numpy()
|
||||
pred_np = pred.squeeze().cpu().numpy()
|
||||
|
||||
# --- Step 3: Creation of Core Error and Dependency Maps ---
|
||||
|
||||
# Create boolean masks for foreground (gt_mask) and background (not_gt_mask) regions.
|
||||
# np.isclose is used for safe floating-point comparison.
|
||||
gt_mask = np.isclose(gt_np, 1)
|
||||
not_gt_mask = np.logical_not(gt_mask)
|
||||
|
||||
# Calculate the initial, simple absolute pixel-wise error map.
|
||||
E = np.abs(pred_np - gt_np)
|
||||
|
||||
# Calculate the Euclidean Distance Transform on the INVERTED mask.
|
||||
# For each background pixel, `dist` will be its distance to the nearest foreground pixel.
|
||||
# `idx` will store the coordinates of that nearest foreground pixel. This is key for the next step.
|
||||
# Note: The original implementation uses scipy.ndimage.morphology.distance_transform_edt
|
||||
# which is an alias for scipy.ndimage.distance_transform_edt.
|
||||
dist, idx = scipy.ndimage.morphology.distance_transform_edt(not_gt_mask, return_indices=True)
|
||||
|
||||
# --- Step 4: Pixel Dependency Map (Et -> EA -> min_E_EA) ---
|
||||
|
||||
# Create the "Pixel Dependency" map, starting with a copy of the original error.
|
||||
Et = np.array(E)
|
||||
|
||||
# This is the crucial step for pixel dependency. For every background pixel,
|
||||
# its error value is replaced with the error value of its NEAREST foreground pixel.
|
||||
# This propagates the error from the foreground edges into the background.
|
||||
Et[not_gt_mask] = E[idx[0, not_gt_mask], idx[1, not_gt_mask]]
|
||||
|
||||
# Smooth the dependency-aware error map with a Gaussian filter.
|
||||
# This creates a soft, blurred error field around the object.
|
||||
sigma = 5.0
|
||||
EA = scipy.ndimage.gaussian_filter(Et, sigma=sigma, truncate=3 / sigma, mode='constant', cval=0.0)
|
||||
|
||||
# The final error for any FOREGROUND pixel is the MINIMUM of its original error (E)
|
||||
# and the new smoothed, dependency-aware error (EA).
|
||||
# This prevents unfairly penalizing small errors right at the boundary of a correct prediction.
|
||||
# The `where=gt_mask` argument ensures this operation only applies to the foreground.
|
||||
min_E_EA = np.minimum(E, EA, where=gt_mask, out=np.array(E))
|
||||
|
||||
# --- Step 5: Pixel Importance Map (B) ---
|
||||
|
||||
# Create the "Pixel Importance" map, which starts as a uniform map of ones.
|
||||
B = np.ones(gt_np.shape)
|
||||
|
||||
# For each BACKGROUND pixel, assign an importance weight based on its distance
|
||||
# from the foreground. The weight is calculated using a non-linear function
|
||||
# that decreases as the distance increases. This makes errors near the object more important.
|
||||
B[not_gt_mask] = 2 - np.exp(np.log(1 - 0.5) / 5 * dist[not_gt_mask])
|
||||
|
||||
# The final Weighted Error Map is the element-wise product of the dependency-aware
|
||||
# error map and the pixel importance map.
|
||||
Ew = min_E_EA * B
|
||||
|
||||
# --- Step 6: Final Metric Computation ---
|
||||
|
||||
# Get a small machine epsilon value for numerically stable division.
|
||||
eps = np.spacing(1)
|
||||
|
||||
# Calculate Weighted True Positives (TPw) and False Positives (FPw) from the Ew map.
|
||||
# TPw = Sum of foreground weights - Sum of weighted errors in the foreground.
|
||||
# FPw = Sum of weighted errors in the background.
|
||||
TPw = np.sum(gt_np) - np.sum(Ew[gt_mask])
|
||||
FPw = np.sum(Ew[not_gt_mask])
|
||||
|
||||
# Calculate Weighted Recall (R) and Weighted Precision (P).
|
||||
# These definitions are specific to this metric's formulation.
|
||||
R = 1 - np.mean(Ew[gt_mask]) # Weighed Recall
|
||||
P = TPw / (eps + TPw + FPw) # Weighted Precision
|
||||
|
||||
# The final Weighted F-measure (Q) is calculated using the standard formula
|
||||
# with the newly computed weighted Precision and Recall.
|
||||
# beta2 = 0.3 is standard, weighing precision more heavily than recall.
|
||||
# Q = 2 * (R * P) / (eps + R + P) # Beta=1
|
||||
Q = (1 + beta2) * (R * P) / (eps + R + (beta2 * P))
|
||||
|
||||
# Raise an error if the result is Not a Number (NaN), indicating a potential issue.
|
||||
if np.isnan(Q):
|
||||
raise ValueError("Weighted F-measure resulted in NaN")
|
||||
|
||||
return Q
|
||||
@@ -0,0 +1,374 @@
|
||||
import json
|
||||
import numpy as np
|
||||
import os
|
||||
from PIL import Image, ImageOps, ImageSequence, ImageFile, UnidentifiedImageError
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
import torch
|
||||
from . import main_logger
|
||||
|
||||
try:
|
||||
# We need to import the built-in LoadImage class for ImageDownload
|
||||
from nodes import LoadImage, LoadImageMask
|
||||
from folder_paths import get_input_directory, get_output_directory
|
||||
has_load_image = hasattr(LoadImage, "load_image") and hasattr(LoadImageMask, "load_image")
|
||||
from comfy.utils import common_upscale
|
||||
except Exception:
|
||||
has_load_image = False
|
||||
|
||||
logger = main_logger
|
||||
|
||||
|
||||
def empty_image(b=0, h=64, w=64, c=1):
|
||||
if b == 0 and c == 1:
|
||||
dims = (h, w)
|
||||
elif c == 1:
|
||||
dims = (b, h, w)
|
||||
else:
|
||||
dims = (b, h, w, c)
|
||||
|
||||
return torch.zeros(dims, dtype=torch.float32, device="cpu")
|
||||
|
||||
|
||||
def upscale(image, width, height, upscale_method, crop="disabled"):
|
||||
# return F.interpolate(image, size=(height, width), mode=upscale_method)
|
||||
return common_upscale(image, width, height, upscale_method, crop=crop)
|
||||
|
||||
|
||||
def upscale_comfy(image, width, height, upscale_method, crop="disabled"):
|
||||
if image.dim() == 3:
|
||||
# A mask
|
||||
if upscale_method == "lanczos":
|
||||
# Lanczos needs an RGB image
|
||||
return common_upscale(image.unsqueeze(1).repeat(1, 3, 1, 1), width, height, upscale_method,
|
||||
crop=crop).movedim(1, -1)[:, :, :, 0]
|
||||
else:
|
||||
return common_upscale(image.unsqueeze(1), width, height, upscale_method, crop=crop).squeeze(1)
|
||||
return common_upscale(image.movedim(-1, 1), width, height, upscale_method, crop=crop).movedim(1, -1)
|
||||
|
||||
|
||||
def pillow(fn, arg):
|
||||
prev_value = None
|
||||
try:
|
||||
x = fn(arg)
|
||||
except (OSError, UnidentifiedImageError, ValueError): # PIL issues #4472 and #2445, also fixes ComfyUI issue #3416
|
||||
prev_value = ImageFile.LOAD_TRUNCATED_IMAGES
|
||||
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
||||
x = fn(arg)
|
||||
finally:
|
||||
if prev_value is not None:
|
||||
ImageFile.LOAD_TRUNCATED_IMAGES = prev_value
|
||||
return x
|
||||
|
||||
|
||||
class CustomLoadImage(object):
|
||||
def load_image(self, image):
|
||||
""" ComfyUI 0.3.59 loader """
|
||||
image_path = image
|
||||
|
||||
img = pillow(Image.open, image_path)
|
||||
|
||||
output_images = []
|
||||
output_masks = []
|
||||
w, h = None, None
|
||||
|
||||
excluded_formats = ['MPO']
|
||||
|
||||
for i in ImageSequence.Iterator(img):
|
||||
i = pillow(ImageOps.exif_transpose, i)
|
||||
|
||||
if i.mode == 'I':
|
||||
i = i.point(lambda i: i * (1 / 255))
|
||||
image = i.convert("RGB")
|
||||
|
||||
if len(output_images) == 0:
|
||||
w = image.size[0]
|
||||
h = image.size[1]
|
||||
|
||||
if image.size[0] != w or image.size[1] != h:
|
||||
continue
|
||||
|
||||
image = np.array(image).astype(np.float32) / 255.0
|
||||
image = torch.from_numpy(image)[None,]
|
||||
if 'A' in i.getbands():
|
||||
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
|
||||
mask = 1. - torch.from_numpy(mask)
|
||||
elif i.mode == 'P' and 'transparency' in i.info:
|
||||
mask = np.array(i.convert('RGBA').getchannel('A')).astype(np.float32) / 255.0
|
||||
mask = 1. - torch.from_numpy(mask)
|
||||
else:
|
||||
mask = empty_image()
|
||||
output_images.append(image)
|
||||
output_masks.append(mask.unsqueeze(0))
|
||||
|
||||
if len(output_images) > 1 and img.format not in excluded_formats:
|
||||
output_image = torch.cat(output_images, dim=0)
|
||||
output_mask = torch.cat(output_masks, dim=0)
|
||||
else:
|
||||
output_image = output_images[0]
|
||||
output_mask = output_masks[0]
|
||||
|
||||
return (output_image, output_mask)
|
||||
|
||||
|
||||
class CustomLoadMask(object):
|
||||
def load_image(self, image, channel):
|
||||
image_path = image
|
||||
i = pillow(Image.open, image_path)
|
||||
i = pillow(ImageOps.exif_transpose, i)
|
||||
if i.getbands() != ("R", "G", "B", "A"):
|
||||
if i.mode == 'I':
|
||||
i = i.point(lambda i: i * (1 / 255))
|
||||
i = i.convert("RGBA")
|
||||
mask = None
|
||||
c = channel[0].upper()
|
||||
if c in i.getbands():
|
||||
mask = np.array(i.getchannel(c)).astype(np.float32) / 255.0
|
||||
mask = torch.from_numpy(mask)
|
||||
if c == 'A':
|
||||
mask = 1. - mask
|
||||
else:
|
||||
mask = empty_image()
|
||||
return (mask.unsqueeze(0),)
|
||||
|
||||
|
||||
def get_image_preview_info(file_name, where="input"):
|
||||
if file_name is None:
|
||||
return {}
|
||||
# This information is for the preview, as we are an output node and we return images
|
||||
# they will be displayed in our node. Quite simple.
|
||||
if os.path.isabs(file_name):
|
||||
base = get_input_directory() if where == "input" else get_output_directory()
|
||||
ff_name = os.path.relpath(file_name, base)
|
||||
fname = os.path.basename(ff_name)
|
||||
dname = os.path.dirname(ff_name)
|
||||
else:
|
||||
fname = file_name
|
||||
dname = ""
|
||||
return {"filename": fname, "subfolder": dname, "type": where}
|
||||
|
||||
|
||||
def load_one_image(file_name, disp_name, embed_transparency, known=None):
|
||||
if file_name is None:
|
||||
return (empty_image(b=1, c=3), empty_image(b=1), file_name)
|
||||
if not os.path.isabs(file_name):
|
||||
file_name = os.path.join(get_input_directory(), file_name)
|
||||
if not os.path.exists(file_name):
|
||||
raise ValueError(f"File '{file_name}' not found")
|
||||
if known is not None and file_name in known:
|
||||
logger.debug('Recycling '+file_name)
|
||||
return (known[file_name][0], known[file_name][1], file_name)
|
||||
|
||||
try:
|
||||
if has_load_image:
|
||||
# --- REUSE ComfyUI's LoadImage LOGIC ---
|
||||
# Instantiate the built-in LoadImage node
|
||||
loader_instance = LoadImage()
|
||||
|
||||
# The LoadImage node's `load_image` method expects the filename as passed
|
||||
# by the ComfyUI widget, which is just the filename. It internally
|
||||
# resolves the path using folder_paths.
|
||||
|
||||
logger.debug(f"Calling built-in LoadImage.load_image() with filename: '{file_name}'")
|
||||
else:
|
||||
# Instantiate the built-in LoadImage node
|
||||
loader_instance = CustomLoadImage()
|
||||
logger.debug(f"Calling our CustomLoadImage.load_image() with filename: '{file_name}'")
|
||||
|
||||
# Call the method and return its result directly
|
||||
result = loader_instance.load_image(file_name)
|
||||
result = (result[0], result[1], file_name)
|
||||
# Create an RGBA image if needed
|
||||
if embed_transparency:
|
||||
image, mask = result
|
||||
# Expand the mask to (b, h, w, 1)
|
||||
mask = mask[..., None]
|
||||
# Concatenate image and mask into (b, h, w, 4)
|
||||
image_with_alpha = torch.cat([image, 1.0 - mask], dim=-1)
|
||||
result = (image_with_alpha, mask, file_name)
|
||||
if known is not None:
|
||||
known[file_name] = (result[0], result[1])
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load image '{disp_name}': {e}", exc_info=True)
|
||||
# Re-raise to make the error visible in ComfyUI
|
||||
raise IOError(f"Could not load the image file '{disp_name}'. "
|
||||
"It may be corrupt or in an unsupported format.") from e
|
||||
|
||||
|
||||
def load_one_mask(file_name, disp_name, channel='red', known=None):
|
||||
if file_name is None:
|
||||
return (empty_image(b=1), file_name)
|
||||
if not os.path.isabs(file_name):
|
||||
file_name = os.path.join(get_input_directory(), file_name)
|
||||
if not os.path.exists(file_name):
|
||||
raise ValueError(f"File '{file_name}' not found")
|
||||
if known is not None and file_name in known:
|
||||
logger.debug('Recycling '+file_name)
|
||||
return (known[file_name], file_name)
|
||||
|
||||
try:
|
||||
if has_load_image:
|
||||
loader_instance = LoadImageMask()
|
||||
logger.debug(f"Calling built-in LoadImageMask.load_image() with filename: '{file_name}'")
|
||||
else:
|
||||
# Instantiate the built-in LoadImage node
|
||||
loader_instance = CustomLoadMask()
|
||||
logger.debug(f"Calling our CustomLoadMask.load_image() with filename: '{file_name}'")
|
||||
|
||||
# Call the method and return its result directly
|
||||
result = loader_instance.load_image(file_name, channel)
|
||||
if known is not None:
|
||||
known[file_name] = result[0]
|
||||
return (result[0], file_name)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load image '{disp_name}': {e}", exc_info=True)
|
||||
# Re-raise to make the error visible in ComfyUI
|
||||
raise IOError(f"Could not load the image file '{disp_name}'. "
|
||||
"It may be corrupt or in an unsupported format.") from e
|
||||
|
||||
|
||||
def load_image_wrapper(file_name, embed_transparency=False, disp_name=None, show_preview=True, channel=None):
|
||||
disp_name = disp_name or file_name
|
||||
if channel is not None:
|
||||
result = load_one_mask(file_name, disp_name, channel)
|
||||
else:
|
||||
result = load_one_image(file_name, disp_name, embed_transparency)
|
||||
if not show_preview:
|
||||
return result
|
||||
return {"ui": {"images": [get_image_preview_info(file_name)]}, "result": result}
|
||||
|
||||
|
||||
def load_images_wrapper(file_names, embed_transparency=False, disp_names=None, show_preview=True, batch_size=1,
|
||||
channel=None):
|
||||
# We work with lists
|
||||
if isinstance(file_names, str):
|
||||
file_names = [file_names]
|
||||
disp_names = disp_names or file_names
|
||||
if isinstance(disp_names, str):
|
||||
disp_names = [disp_names]
|
||||
|
||||
imgs = []
|
||||
masks = []
|
||||
known_imgs = {}
|
||||
known_masks = {}
|
||||
is_mask = channel is not None
|
||||
used_file_names = []
|
||||
all_preview_imgs = []
|
||||
total = len(file_names)
|
||||
for i in range(0, total, batch_size):
|
||||
if batch_size > 1:
|
||||
# Add upto batch_size images in a batch
|
||||
imgs_batch = []
|
||||
masks_batch = []
|
||||
max_w = max_h = 0
|
||||
max_mw = max_mh = 0
|
||||
for j in range(batch_size):
|
||||
index = i+j
|
||||
if index >= total:
|
||||
continue
|
||||
file_name = file_names[i+j]
|
||||
disp_name = disp_names[i+j]
|
||||
|
||||
if is_mask:
|
||||
# A mask
|
||||
mask, file_name = load_one_mask(file_name, disp_name, channel, known=known_masks)
|
||||
else:
|
||||
img, mask, file_name = load_one_image(file_name, disp_name, embed_transparency, known=known_imgs)
|
||||
|
||||
max_w = max(max_w, img.shape[2])
|
||||
max_h = max(max_h, img.shape[1])
|
||||
imgs_batch.append(img)
|
||||
|
||||
max_mw = max(max_mw, mask.shape[2])
|
||||
max_mh = max(max_mh, mask.shape[1])
|
||||
masks_batch.append(mask)
|
||||
|
||||
used_file_names.append(file_name)
|
||||
|
||||
if show_preview:
|
||||
all_preview_imgs.append(get_image_preview_info(file_name))
|
||||
|
||||
for j in range(len(imgs_batch)):
|
||||
if not is_mask:
|
||||
img = imgs_batch[j]
|
||||
H, W = img.shape[1:3]
|
||||
if H != max_h or W != max_w:
|
||||
logger.debug(f"Upscaling image to fit batch: {W}x{H} -> {max_w}x{max_h}")
|
||||
imgs_batch[j] = upscale_comfy(img, max_w, max_h, "bicubic")
|
||||
|
||||
mask = masks_batch[j]
|
||||
H, W = mask.shape[1:3]
|
||||
if H != max_mh or W != max_mw:
|
||||
logger.debug(f"Upscaling mask to fit batch: {W}x{H} -> {max_mw}x{max_mh}")
|
||||
masks_batch[j] = upscale_comfy(mask, max_mw, max_mh, "bicubic")
|
||||
|
||||
if not is_mask:
|
||||
imgs.append(torch.cat(imgs_batch))
|
||||
masks.append(torch.cat(masks_batch))
|
||||
else:
|
||||
# Add a single image
|
||||
file_name = file_names[i]
|
||||
|
||||
if is_mask:
|
||||
# A mask
|
||||
mask, file_name = load_one_mask(file_name, disp_names[i], channel, known=known_masks)
|
||||
else:
|
||||
img, mask, file_name = load_one_image(file_name, disp_names[i], embed_transparency, known=known_imgs)
|
||||
imgs.append(img)
|
||||
|
||||
masks.append(mask)
|
||||
used_file_names.append(file_name)
|
||||
|
||||
if show_preview:
|
||||
all_preview_imgs.append(get_image_preview_info(file_name))
|
||||
|
||||
logger.debug(f"Loaded {len(masks)} batches:")
|
||||
items = masks if is_mask else imgs
|
||||
for n, i in enumerate(items):
|
||||
logger.debug(f"{n}) {i.shape}")
|
||||
|
||||
result = (imgs, masks, used_file_names) if channel is None else (masks, used_file_names)
|
||||
if not show_preview:
|
||||
return result
|
||||
|
||||
return {"ui": {"images": all_preview_imgs}, "result": result}
|
||||
|
||||
|
||||
def save_image(images, filenames, prompt=None, extra_pnginfo=None, compress_level=4, show_preview=True):
|
||||
if isinstance(filenames, str):
|
||||
if not filenames:
|
||||
raise ValueError("You must provide a file name")
|
||||
filenames = [filenames]
|
||||
|
||||
# Make a list with all the images, unroll the batches
|
||||
imgs = []
|
||||
for img in images:
|
||||
imgs.extend([i for i in img])
|
||||
|
||||
if len(filenames) != len(imgs):
|
||||
raise ValueError(f"{len(imgs)} images provided but only {len(filenames)} file names")
|
||||
|
||||
all_preview_imgs = []
|
||||
for index, (image, filename) in enumerate(zip(imgs, filenames)):
|
||||
if filename is None:
|
||||
continue
|
||||
i = 255. * image.cpu().numpy()
|
||||
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
||||
metadata = PngInfo()
|
||||
if prompt is not None:
|
||||
metadata.add_text("prompt", json.dumps(prompt))
|
||||
if extra_pnginfo is not None:
|
||||
for x in extra_pnginfo:
|
||||
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
|
||||
|
||||
logger.debug(f"Saving {image.shape[1]}x{image.shape[0]} image to {filename}")
|
||||
img.save(filename, pnginfo=metadata, compress_level=compress_level)
|
||||
|
||||
if show_preview:
|
||||
all_preview_imgs.append(get_image_preview_info(filename, where="output"))
|
||||
|
||||
if not show_preview:
|
||||
return ()
|
||||
return {"ui": {"images": all_preview_imgs}, "result": ()}
|
||||
+2086
-316
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,195 @@
|
||||
import torch
|
||||
from . import EPS
|
||||
|
||||
|
||||
def get_s_measure(pred: torch.Tensor, gt: torch.Tensor, alpha: float = 0.5) -> float:
|
||||
"""
|
||||
Calculates the S-measure for a given prediction and ground truth mask.
|
||||
The S-measure evaluates structural similarity, combining object-aware and
|
||||
region-aware metrics.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): The continuous prediction mask (normalized to [0, 1]).
|
||||
gt (torch.Tensor): The binary ground truth mask (values are 0 or 1).
|
||||
alpha (float): The weight for balancing object-aware vs. region-aware scores.
|
||||
Defaults to 0.5.
|
||||
|
||||
Returns:
|
||||
float: The final S-measure score.
|
||||
"""
|
||||
# gt is assumed to be binary
|
||||
y = gt.mean()
|
||||
if y == 0:
|
||||
# If the ground truth is all black, the score is the inverse of the prediction's average.
|
||||
# A perfect prediction would also be all black (mean=0), yielding a score of 1.
|
||||
x = pred.mean()
|
||||
q = 1.0 - x
|
||||
elif y == 1:
|
||||
# If the ground truth is all white, the score is simply the prediction's average.
|
||||
# A perfect prediction would be all white (mean=1), yielding a score of 1.
|
||||
x = pred.mean()
|
||||
q = x
|
||||
else:
|
||||
# For a mixed ground truth, balance the object and region scores.
|
||||
q = alpha * s_object(pred, gt) + (1 - alpha) * s_region(pred, gt)
|
||||
# Ensure the score is non-negative.
|
||||
if q < 0:
|
||||
return 0
|
||||
|
||||
return q.item()
|
||||
|
||||
|
||||
def s_object(pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Calculates the object-aware structural similarity score.
|
||||
This measures the similarity for foreground and background regions separately
|
||||
and combines them based on the foreground's size.
|
||||
"""
|
||||
# Create a map of foreground-only predictions
|
||||
fg = torch.where(gt == 0, torch.zeros_like(pred), pred)
|
||||
# Create a map of background-only predictions (inverted)
|
||||
bg = torch.where(gt == 1, torch.zeros_like(pred), 1 - pred)
|
||||
|
||||
# Calculate scores for each region
|
||||
o_fg = s_object_calc(fg, gt)
|
||||
o_bg = s_object_calc(bg, 1 - gt)
|
||||
|
||||
# Combine scores based on foreground area
|
||||
u = gt.mean() # The ratio of foreground pixels
|
||||
q = u * o_fg + (1 - u) * o_bg
|
||||
return q
|
||||
|
||||
|
||||
def s_object_calc(pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Helper function to compute the similarity score for a given region (FG or BG).
|
||||
"""
|
||||
# Select the prediction pixels corresponding to the region of interest in the ground truth
|
||||
region_pixels = pred[gt == 1]
|
||||
|
||||
# If the region is empty, the score is undefined, but 0 is a safe return.
|
||||
if region_pixels.numel() == 0:
|
||||
return torch.tensor(0.0, device=pred.device)
|
||||
|
||||
x = region_pixels.mean()
|
||||
|
||||
# Note: The original paper's formula uses variance (sigma**2), but many public
|
||||
# implementations use std dev (sigma). We follow this implementation's logic.
|
||||
sigma_x = region_pixels.std()
|
||||
|
||||
# A score that rewards high mean (x) but penalizes high variance (sigma_x).
|
||||
# It is maximized when x=1 and sigma_x=0.
|
||||
score = 2.0 * x / (x**2 + 1.0 + sigma_x + EPS)
|
||||
return score
|
||||
|
||||
|
||||
def s_region(pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Calculates the region-aware structural similarity score.
|
||||
This divides the image into four quadrants based on the ground truth's
|
||||
centroid and computes a weighted SSIM score.
|
||||
"""
|
||||
# Find the centroid (center of mass) of the ground truth mask
|
||||
X, Y = s_centroid(gt)
|
||||
|
||||
# Divide the ground truth and prediction into four quadrants based on the centroid
|
||||
gt_parts = s_divide_tensor(gt, X, Y)
|
||||
pred_parts = s_divide_tensor(pred, X, Y)
|
||||
|
||||
# Calculate the area weights for each quadrant
|
||||
w = s_calculate_weights(gt.shape, X, Y)
|
||||
|
||||
# Calculate SSIM for each quadrant and combine the quadrant scores using their area weights
|
||||
return sum([w[i] * s_ssim(pred_parts[i], gt_parts[i]) for i in range(4)])
|
||||
|
||||
|
||||
def s_centroid(gt: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Calculates the centroid (center of mass) of a binary mask.
|
||||
"""
|
||||
rows, cols = gt.shape[-2:]
|
||||
total = gt.sum()
|
||||
|
||||
# Handle the edge case of an all-black mask
|
||||
if total == 0:
|
||||
X = torch.tensor(round(cols / 2), device=gt.device)
|
||||
Y = torch.tensor(round(rows / 2), device=gt.device)
|
||||
else:
|
||||
# Create coordinate ranges directly on the target device, avoiding CPU-GPU transfer.
|
||||
i_coords = torch.arange(cols, device=gt.device, dtype=gt.dtype)
|
||||
j_coords = torch.arange(rows, device=gt.device, dtype=gt.dtype)
|
||||
|
||||
# Calculate weighted average of coordinates
|
||||
X = torch.round((gt.sum(dim=0) * i_coords).sum() / total)
|
||||
Y = torch.round((gt.sum(dim=1) * j_coords).sum() / total)
|
||||
|
||||
return X.long(), Y.long()
|
||||
|
||||
|
||||
def s_divide_tensor(tensor: torch.Tensor, X: torch.Tensor, Y: torch.Tensor) -> tuple:
|
||||
"""
|
||||
Divides a tensor into four quadrants based on a pivot point (X, Y).
|
||||
This is memory-efficient as it returns views, not copies.
|
||||
"""
|
||||
h, w = tensor.shape[-2:]
|
||||
LT = tensor[..., :Y, :X]
|
||||
RT = tensor[..., :Y, X:]
|
||||
LB = tensor[..., Y:, :X]
|
||||
RB = tensor[..., Y:, X:]
|
||||
return LT, RT, LB, RB
|
||||
|
||||
|
||||
def s_calculate_weights(shape: tuple, X: torch.Tensor, Y: torch.Tensor) -> tuple:
|
||||
"""Calculates the proportional area of the four quadrants."""
|
||||
h, w = shape[-2:]
|
||||
area = h * w
|
||||
|
||||
# Ensure coordinates are float for division
|
||||
Xf = X.float()
|
||||
Yf = Y.float()
|
||||
|
||||
w1 = Xf * Yf / area
|
||||
w2 = (w - Xf) * Yf / area
|
||||
w3 = Xf * (h - Yf) / area
|
||||
w4 = 1.0 - w1 - w2 - w3 # More stable calculation for the last weight
|
||||
|
||||
return (w1, w2, w3, w4)
|
||||
|
||||
|
||||
def s_ssim(pred: torch.Tensor, gt: torch.Tensor) -> float:
|
||||
"""
|
||||
Computes a custom structural similarity (SSIM-like) score between two tensors.
|
||||
"""
|
||||
# If a quadrant is empty, its contribution to similarity is ambiguous.
|
||||
# Returning 0 is a safe choice, but 1 could also be argued if gt is also empty.
|
||||
if pred.numel() == 0 or gt.numel() == 0:
|
||||
return 0.0
|
||||
|
||||
h, w = pred.shape[-2:]
|
||||
N = h * w
|
||||
|
||||
# Means
|
||||
x = pred.mean()
|
||||
y = gt.mean()
|
||||
|
||||
# Variances and Covariance (using unbiased estimator N-1)
|
||||
# This is numerically safer than calculating std dev separately.
|
||||
sigma_x2 = ((pred - x) * (pred - x)).sum() / (N - 1 + EPS)
|
||||
sigma_y2 = ((gt - y) * (gt - y)).sum() / (N - 1 + EPS)
|
||||
sigma_xy = ((pred - x) * (gt - y)).sum() / (N - 1 + EPS)
|
||||
|
||||
# Numerator and denominator of the SSIM formula
|
||||
alpha = 4 * x * y * sigma_xy
|
||||
beta = (x*x + y*y) * (sigma_x2 + sigma_y2)
|
||||
|
||||
# Handle special cases for stability, as defined in the original code
|
||||
if alpha != 0:
|
||||
Q = alpha / (beta + EPS)
|
||||
elif alpha == 0 and beta == 0:
|
||||
# If both inputs are flat and identical, similarity is perfect.
|
||||
Q = 1.0
|
||||
else:
|
||||
# If numerator is 0 but denominator isn't, similarity is 0.
|
||||
Q = 0
|
||||
|
||||
return Q
|
||||
Reference in New Issue
Block a user