depth-pro / src /depth_pro /eval /boundary_metrics.py
akhaliq's picture
akhaliq HF staff
Upload folder using huggingface_hub
de1b1de verified
from typing import List, Tuple
import numpy as np
def connected_component(r: np.ndarray, c: np.ndarray) -> List[List[int]]:
"""Find connected components in the given row and column indices.
Args:
----
r (np.ndarray): Row indices.
c (np.ndarray): Column indices.
Yields:
------
List[int]: Indices of connected components.
"""
indices = [0]
for i in range(1, r.size):
if r[i] == r[indices[-1]] and c[i] == c[indices[-1]] + 1:
indices.append(i)
else:
yield indices
indices = [i]
yield indices
def nms_horizontal(ratio: np.ndarray, threshold: float) -> np.ndarray:
"""Apply Non-Maximum Suppression (NMS) horizontally on the given ratio matrix.
Args:
----
ratio (np.ndarray): Input ratio matrix.
threshold (float): Threshold for NMS.
Returns:
-------
np.ndarray: Binary mask after applying NMS.
"""
mask = np.zeros_like(ratio, dtype=bool)
r, c = np.nonzero(ratio > threshold)
if len(r) == 0:
return mask
for ids in connected_component(r, c):
values = [ratio[r[i], c[i]] for i in ids]
mi = np.argmax(values)
mask[r[ids[mi]], c[ids[mi]]] = True
return mask
def nms_vertical(ratio: np.ndarray, threshold: float) -> np.ndarray:
"""Apply Non-Maximum Suppression (NMS) vertically on the given ratio matrix.
Args:
----
ratio (np.ndarray): Input ratio matrix.
threshold (float): Threshold for NMS.
Returns:
-------
np.ndarray: Binary mask after applying NMS.
"""
return np.transpose(nms_horizontal(np.transpose(ratio), threshold))
def fgbg_depth(
d: np.ndarray, t: float
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Find foreground-background relations between neighboring pixels.
Args:
----
d (np.ndarray): Depth matrix.
t (float): Threshold for comparison.
Returns:
-------
Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: Four matrices indicating
left, top, right, and bottom foreground-background relations.
"""
right_is_big_enough = (d[..., :, 1:] / d[..., :, :-1]) > t
left_is_big_enough = (d[..., :, :-1] / d[..., :, 1:]) > t
bottom_is_big_enough = (d[..., 1:, :] / d[..., :-1, :]) > t
top_is_big_enough = (d[..., :-1, :] / d[..., 1:, :]) > t
return (
left_is_big_enough,
top_is_big_enough,
right_is_big_enough,
bottom_is_big_enough,
)
def fgbg_depth_thinned(
d: np.ndarray, t: float
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Find foreground-background relations between neighboring pixels with Non-Maximum Suppression.
Args:
----
d (np.ndarray): Depth matrix.
t (float): Threshold for NMS.
Returns:
-------
Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: Four matrices indicating
left, top, right, and bottom foreground-background relations with NMS applied.
"""
right_is_big_enough = nms_horizontal(d[..., :, 1:] / d[..., :, :-1], t)
left_is_big_enough = nms_horizontal(d[..., :, :-1] / d[..., :, 1:], t)
bottom_is_big_enough = nms_vertical(d[..., 1:, :] / d[..., :-1, :], t)
top_is_big_enough = nms_vertical(d[..., :-1, :] / d[..., 1:, :], t)
return (
left_is_big_enough,
top_is_big_enough,
right_is_big_enough,
bottom_is_big_enough,
)
def fgbg_binary_mask(
d: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Find foreground-background relations between neighboring pixels in binary masks.
Args:
----
d (np.ndarray): Binary depth matrix.
Returns:
-------
Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: Four matrices indicating
left, top, right, and bottom foreground-background relations in binary masks.
"""
assert d.dtype == bool
right_is_big_enough = d[..., :, 1:] & ~d[..., :, :-1]
left_is_big_enough = d[..., :, :-1] & ~d[..., :, 1:]
bottom_is_big_enough = d[..., 1:, :] & ~d[..., :-1, :]
top_is_big_enough = d[..., :-1, :] & ~d[..., 1:, :]
return (
left_is_big_enough,
top_is_big_enough,
right_is_big_enough,
bottom_is_big_enough,
)
def edge_recall_matting(pr: np.ndarray, gt: np.ndarray, t: float) -> float:
"""Calculate edge recall for image matting.
Args:
----
pr (np.ndarray): Predicted depth matrix.
gt (np.ndarray): Ground truth binary mask.
t (float): Threshold for NMS.
Returns:
-------
float: Edge recall value.
"""
assert gt.dtype == bool
ap, bp, cp, dp = fgbg_depth_thinned(pr, t)
ag, bg, cg, dg = fgbg_binary_mask(gt)
return 0.25 * (
np.count_nonzero(ap & ag) / max(np.count_nonzero(ag), 1)
+ np.count_nonzero(bp & bg) / max(np.count_nonzero(bg), 1)
+ np.count_nonzero(cp & cg) / max(np.count_nonzero(cg), 1)
+ np.count_nonzero(dp & dg) / max(np.count_nonzero(dg), 1)
)
def boundary_f1(
pr: np.ndarray,
gt: np.ndarray,
t: float,
return_p: bool = False,
return_r: bool = False,
) -> float:
"""Calculate Boundary F1 score.
Args:
----
pr (np.ndarray): Predicted depth matrix.
gt (np.ndarray): Ground truth depth matrix.
t (float): Threshold for comparison.
return_p (bool, optional): If True, return precision. Defaults to False.
return_r (bool, optional): If True, return recall. Defaults to False.
Returns:
-------
float: Boundary F1 score, or precision, or recall depending on the flags.
"""
ap, bp, cp, dp = fgbg_depth(pr, t)
ag, bg, cg, dg = fgbg_depth(gt, t)
r = 0.25 * (
np.count_nonzero(ap & ag) / max(np.count_nonzero(ag), 1)
+ np.count_nonzero(bp & bg) / max(np.count_nonzero(bg), 1)
+ np.count_nonzero(cp & cg) / max(np.count_nonzero(cg), 1)
+ np.count_nonzero(dp & dg) / max(np.count_nonzero(dg), 1)
)
p = 0.25 * (
np.count_nonzero(ap & ag) / max(np.count_nonzero(ap), 1)
+ np.count_nonzero(bp & bg) / max(np.count_nonzero(bp), 1)
+ np.count_nonzero(cp & cg) / max(np.count_nonzero(cp), 1)
+ np.count_nonzero(dp & dg) / max(np.count_nonzero(dp), 1)
)
if r + p == 0:
return 0.0
if return_p:
return p
if return_r:
return r
return 2 * (r * p) / (r + p)
def get_thresholds_and_weights(
t_min: float, t_max: float, N: int
) -> Tuple[np.ndarray, np.ndarray]:
"""Generate thresholds and weights for the given range.
Args:
----
t_min (float): Minimum threshold.
t_max (float): Maximum threshold.
N (int): Number of thresholds.
Returns:
-------
Tuple[np.ndarray, np.ndarray]: Array of thresholds and corresponding weights.
"""
thresholds = np.linspace(t_min, t_max, N)
weights = thresholds / thresholds.sum()
return thresholds, weights
def invert_depth(depth: np.ndarray, eps: float = 1e-6) -> np.ndarray:
"""Inverts a depth map with numerical stability.
Args:
----
depth (np.ndarray): Depth map to be inverted.
eps (float): Minimum value to avoid division by zero (default is 1e-6).
Returns:
-------
np.ndarray: Inverted depth map.
"""
inverse_depth = 1.0 / depth.clip(min=eps)
return inverse_depth
def SI_boundary_F1(
predicted_depth: np.ndarray,
target_depth: np.ndarray,
t_min: float = 1.05,
t_max: float = 1.25,
N: int = 10,
) -> float:
"""Calculate Scale-Invariant Boundary F1 Score for depth-based ground-truth.
Args:
----
predicted_depth (np.ndarray): Predicted depth matrix.
target_depth (np.ndarray): Ground truth depth matrix.
t_min (float, optional): Minimum threshold. Defaults to 1.05.
t_max (float, optional): Maximum threshold. Defaults to 1.25.
N (int, optional): Number of thresholds. Defaults to 10.
Returns:
-------
float: Scale-Invariant Boundary F1 Score.
"""
assert predicted_depth.ndim == target_depth.ndim == 2
thresholds, weights = get_thresholds_and_weights(t_min, t_max, N)
f1_scores = np.array(
[
boundary_f1(invert_depth(predicted_depth), invert_depth(target_depth), t)
for t in thresholds
]
)
return np.sum(f1_scores * weights)
def SI_boundary_Recall(
predicted_depth: np.ndarray,
target_mask: np.ndarray,
t_min: float = 1.05,
t_max: float = 1.25,
N: int = 10,
alpha_threshold: float = 0.1,
) -> float:
"""Calculate Scale-Invariant Boundary Recall Score for mask-based ground-truth.
Args:
----
predicted_depth (np.ndarray): Predicted depth matrix.
target_mask (np.ndarray): Ground truth binary mask.
t_min (float, optional): Minimum threshold. Defaults to 1.05.
t_max (float, optional): Maximum threshold. Defaults to 1.25.
N (int, optional): Number of thresholds. Defaults to 10.
alpha_threshold (float, optional): Threshold for alpha masking. Defaults to 0.1.
Returns:
-------
float: Scale-Invariant Boundary Recall Score.
"""
assert predicted_depth.ndim == target_mask.ndim == 2
thresholds, weights = get_thresholds_and_weights(t_min, t_max, N)
thresholded_target = target_mask > alpha_threshold
recall_scores = np.array(
[
edge_recall_matting(
invert_depth(predicted_depth), thresholded_target, t=float(t)
)
for t in thresholds
]
)
weighted_recall = np.sum(recall_scores * weights)
return weighted_recall