File size: 2,942 Bytes
6cd90b7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
# coding=utf-8
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Metrics for evaluating the performance of the model."""
import torch
def IoU(mask1, mask2, threshold=0.5):
"""Calculate Intersection over Union (IoU) between prediction and GT masks.
Args:
mask1: A torch.Tensor denoting the prediction, shape (N, H, W), where N is
the number of masks.
mask2: A torch.Tensor denoting the ground truth, shape (N, H, W), where N
is the number of masks.
threshold: The threshold to binarize masks.
Returns:
IoU of `mask1` and `mask2`.
"""
if threshold > 0:
mask1, mask2 = (mask1 > threshold).to(torch.bool), (mask2 > threshold).to(
torch.bool
)
intersection = torch.sum(mask1 * (mask1 == mask2), dim=[-1, -2]).squeeze()
union = torch.sum(mask1 + mask2, dim=[-1, -2]).squeeze()
if union.sum() == 0:
return 0
return (intersection.to(torch.float) / union).mean().item()
def IoM(pred, target, min_pred_threshold=0.2):
"""Calculate Intersection over the area of gt Mask and pred Mask (IoM).
between prediction and each ground truth masks.
Precaution:
this function works for prediction and target that are binary masks,
where 1 represents the mask and 0 represents the background.
Args:
pred: A torch.Tensor denoting the prediction, shape (N, H, W), where N is
the number of masks.
target: A torch.Tensor denoting the ground truth, shape (N, H, W), where N
is the number of masks.
min_pred_threshold: prediction threshold.
Returns:
ious: A torch.Tensor denoting the IoU, shape (N,).
"""
# calculate the intersection over all masks
intersection = torch.einsum("mij,nij->mn", pred.to(target.device), target)
area_pred = torch.einsum("mij->m", pred)
area_target = torch.einsum("nij->n", target)
# we calculate the IoM by dividing the intersection over the minimum area.
iom_target = torch.einsum("mn,n->mn", intersection, 1 / area_target)
iom_pred = torch.einsum("mn,m->mn", intersection, 1 / area_pred)
# if the intersection is smaller than a certain percentage of the area of
# the pred mask, we consider it as background.
iom_target[iom_pred < min_pred_threshold] = 0
# we consider the IoM as the maximum IoM between the pred mask and
# the target mask.
iom = torch.max(iom_target, iom_pred)
iom = iom.max(dim=0)[0]
return iom
|