66 lines
2.8 KiB
Python
66 lines
2.8 KiB
Python
|
#!/usr/bin/env python
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
@project:
|
|||
|
|
@File : edge_loss
|
|||
|
|
@Author : qiqq
|
|||
|
|
@create_time : 2023/1/5 22:06
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import logging
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
|
|||
|
|
|
|||
|
|
class EdgeSegLoss(nn.Module):
|
|||
|
|
def __init__(self, classes, ignore_index=255,
|
|||
|
|
edge_weight=1, seg_weight=1, seg_body_weight=1, att_weight=1):
|
|||
|
|
super(EdgeSegLoss, self).__init__()
|
|||
|
|
self.num_classes = classes
|
|||
|
|
|
|||
|
|
|
|||
|
|
self.ignore_index = ignore_index
|
|||
|
|
self.edge_weight = edge_weight
|
|||
|
|
self.seg_weight = seg_weight
|
|||
|
|
self.att_weight = att_weight
|
|||
|
|
self.seg_body_weight = seg_body_weight
|
|||
|
|
|
|||
|
|
def bce2d(self, input, target):
|
|||
|
|
'''input是网络的输出，没有经过sigmode的。target的标签'''
|
|||
|
|
n, c, h, w = input.size() #c=1 代表是边缘
|
|||
|
|
|
|||
|
|
log_p = input.transpose(1, 2).transpose(2, 3).contiguous().view(1, -1) #.view(1, -1)把原来的tensor reshape成两个维度1是第一个维度是1，-1是第二给维度由估计得来，也就是把原来的tensor拉成一个一行b*c*h*w的行向量 #以（2，1，512，512）---》（1，524288）
|
|||
|
|
target_t = target.transpose(1, 2).transpose(2, 3).contiguous().view(1, -1)
|
|||
|
|
target_trans = target_t.clone()
|
|||
|
|
|
|||
|
|
pos_index = (target_t == 1) #边缘的 (1,524288) 是boole类型
|
|||
|
|
neg_index = (target_t == 0) #非边缘的 (1,524288) 是boole类型
|
|||
|
|
ignore_index = (target_t > 1) #那些既不是边缘也不是背景的的 要被和忽略的
|
|||
|
|
|
|||
|
|
target_trans[pos_index] = 1 #（1，524288）去掉了ignor的并且把boole类型转成0 1类型的
|
|||
|
|
target_trans[neg_index] = 0
|
|||
|
|
|
|||
|
|
pos_index = pos_index.data.cpu().numpy().astype(bool)
|
|||
|
|
neg_index = neg_index.data.cpu().numpy().astype(bool)
|
|||
|
|
ignore_index = ignore_index.data.cpu().numpy().astype(bool) #以上三个搞到cpu numpy上
|
|||
|
|
|
|||
|
|
weight = torch.Tensor(log_p.size()).fill_(0) #一个全为0的 （1，524288）目的是好像是为每个piex生成一个权重
|
|||
|
|
weight = weight.numpy()
|
|||
|
|
pos_num = pos_index.sum() #所有边缘的pix数量 （在）[0,0,:,:]这张图中共61754
|
|||
|
|
neg_num = neg_index.sum() #非边缘的像素数量
|
|||
|
|
sum_num = pos_num + neg_num #有效的像素数量 #462534
|
|||
|
|
weight[pos_index] = neg_num * 1.0 / sum_num #以边缘和背景的频率作为weight （注意是数量多的权重大，数量少的权重小）
|
|||
|
|
weight[neg_index] = pos_num * 1.0 / sum_num
|
|||
|
|
|
|||
|
|
weight[ignore_index] = 0
|
|||
|
|
|
|||
|
|
|
|||
|
|
weight = torch.from_numpy(weight).cuda()
|
|||
|
|
log_p = log_p.cuda() #原始的被seshape成了（1，bchw）的行向量
|
|||
|
|
target_t = target_t.cuda()
|
|||
|
|
|
|||
|
|
loss = F.binary_cross_entropy_with_logits(log_p, target_t, weight, size_average=True)
|
|||
|
|
|
|||
|
|
return loss
|