具有加权类别的PyTorch语义分割

时间:2020-10-21 13:08:25

标签: pytorch cnn unity3d-unet semantic-segmentation

我正在研究如何对图像进行语义分割,但是我想为这些类分配权重。我正在使用此Semantic Segmentation PyTorch,在实际的实现中,无法将权重应用于类。 我试图更改util / functional.py,如下面报告的代码一样,但是培训阶段并未提高性能。特别是IOU_Score在0.26的值上振荡,而没有任何良好的性能。

我正在使用example文件夹中报告的笔记本电脑进行多类配置。

任何人都可以帮助我了解如何更改指标或如何使用分配了课程权重的其他指标。

谢谢您的帮助。

def iou(pr, gt, eps=1e-7, threshold=None, ignore_channels=None, weights=None):
    """Calculate Intersection over Union between ground truth and prediction
    Args:
        pr (torch.Tensor): predicted tensor
        gt (torch.Tensor):  ground truth tensor
        eps (float): epsilon to avoid zero division
        threshold: threshold for outputs binarization
    Returns:
        float: IoU (Jaccard) score
    """

    pr = _threshold(pr, threshold=threshold)
    pr, gt = _take_channels(pr, gt, ignore_channels=ignore_channels)

    #print('PR', pr.shape, type(pr))
    #print('GT', gt.shape, type(gt))

    pr1 = torch.zeros(pr.shape[0], pr.shape[1], pr.shape[2], pr.shape[3])
    gt1 = torch.zeros(gt.shape[0], gt.shape[1], gt.shape[2], gt.shape[3])

    pr1[:, 0, :, :] = pr[:, 0, :, :] / 100
    pr1[:, 1, :, :] = pr[:, 1, :, :] / 2
    pr1[:, 2, :, :] = pr[:, 2, :, :] / 2
    pr1[:, 3, :, :] = pr[:, 3, :, :] / 100
    pr1[:, 4, :, :] = pr[:, 4, :, :] / 100

    gt1[:, 0, :, :] = gt[:, 0, :, :] / 100
    gt1[:, 1, :, :] = gt[:, 1, :, :] / 2
    gt1[:, 2, :, :] = gt[:, 2, :, :] / 2
    gt1[:, 3, :, :] = gt[:, 3, :, :] / 100
    gt1[:, 4, :, :] = gt[:, 4, :, :] / 100

    intersection = torch.sum(gt1 * pr1)
    union = torch.sum(gt1) + torch.sum(pr1) - intersection + eps
    return (intersection + eps) / union


jaccard = iou


def f_score(pr, gt, beta=1, eps=1e-7, threshold=None, ignore_channels=None, weights=None):
    """Calculate F-score between ground truth and prediction
    Args:
        pr (torch.Tensor): predicted tensor
        gt (torch.Tensor):  ground truth tensor
        beta (float): positive constant
        eps (float): epsilon to avoid zero division
        threshold: threshold for outputs binarization
    Returns:
        float: F score
    """

    pr = _threshold(pr, threshold=threshold)
    pr, gt = _take_channels(pr, gt, ignore_channels=ignore_channels)

    #print('PR', pr.shape, type(pr))
    #print('GT', gt.shape, type(gt))

    pr1 = torch.zeros(pr.shape[0], pr.shape[1], pr.shape[2], pr.shape[3])
    gt1 = torch.zeros(gt.shape[0], gt.shape[1], gt.shape[2], gt.shape[3])

    pr1[:, 0, :, :] = pr[:, 0, :, :] / 100
    pr1[:, 1, :, :] = pr[:, 1, :, :] / 2
    pr1[:, 2, :, :] = pr[:, 2, :, :] / 2
    pr1[:, 3, :, :] = pr[:, 3, :, :] / 100
    pr1[:, 4, :, :] = pr[:, 4, :, :] / 100

    gt1[:, 0, :, :] = gt[:, 0, :, :] / 100
    gt1[:, 1, :, :] = gt[:, 1, :, :] / 2
    gt1[:, 2, :, :] = gt[:, 2, :, :] / 2
    gt1[:, 3, :, :] = gt[:, 3, :, :] / 100
    gt1[:, 4, :, :] = gt[:, 4, :, :] / 100

    tp = torch.sum(gt1 * pr1)
    fp = torch.sum(pr1) - tp
    fn = torch.sum(gt1) - tp

    score = ((1 + beta ** 2) * tp + eps) \
            / ((1 + beta ** 2) * tp + beta ** 2 * fn + fp + eps)

    return score

0 个答案:

没有答案