Spaces:
Sleeping
Sleeping
File size: 2,995 Bytes
d4b77ac |
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 |
# under development
import cv2
import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
class EdgeLoss(nn.Module):
def __init__(self, device, reduction='mean'):
super().__init__()
self.loss_fn = nn.L1Loss(reduction=reduction)
self.gaussianKernel = self._gaussianKernel2dOpencv(kernel_size=5, sigma=1)
self.gaussianKernel = np.reshape(self.gaussianKernel, (1, 1, 5, 5)) # 5 is kernel size
self.gaussianKernel = torch.from_numpy(self.gaussianKernel).float().to(device)
def forward(self, outputs, GTs, masks, cannyEdges):
"""
Calculate the L1 loss in the edge regions
edges are detected by canny operator
Args:
outputs: torch tensor, shape [b, c, h, w]
GTs: torch tensor, shape [b, c, h, w]
masks: torch tensor, shape [b, 1, h, w]
cannyEdges: shape [b, c, h, w], 1 indicates edge regions, while 0 indicates nonedge regions
cannyEdges should be provided by the dataloader
Returns: edge loss between outputs and GTs
"""
cannyEdges = self.gaussianBlur(cannyEdges)
loss = self.loss_fn(outputs * cannyEdges * masks, GTs * cannyEdges * masks) / torch.mean(masks)
return loss
def gaussianBlur(self, cannyEdges, iteration=2):
for i in range(iteration):
cannyEdges = F.conv2d(cannyEdges, self.gaussianKernel, stride=1, padding=2)
return cannyEdges
def _gaussianKernel2dOpencv(self, kernel_size=5, sigma=1):
kx = cv2.getGaussianKernel(kernel_size, sigma)
ky = cv2.getGaussianKernel(kernel_size, sigma)
return np.multiply(kx, np.transpose(ky))
if __name__ == '__main__':
from PIL import Image
from torchvision.transforms import ToTensor
from skimage.feature import canny
from torchvision.utils import save_image
from skimage.color import rgb2gray
output = Image.open('images/00001_res.jpg')
GT = Image.open('images/img2.jpg')
# mask = Image.open('images/mask.png')
cannyMap = canny(rgb2gray(np.array(GT)), sigma=2).astype(np.float32)
output = ToTensor()(output).unsqueeze(0)
GT = ToTensor()(GT).unsqueeze(0)
# mask = ToTensor()(mask).unsqueeze(0)
mask = torch.ones_like(GT)
cannyMap = ToTensor()(cannyMap).unsqueeze(0)
output = output * 2 - 1
GT = GT * 2 - 1
EdgeLossLayer = EdgeLoss(2, 'cpu')
edgeLoss, cannyPriority, edgeMap_output, edgeMap_GT, errorMap, errorMap2 = EdgeLossLayer(output, GT, mask, cannyMap)
print(edgeLoss)
save_image(cannyPriority, 'images/cannyPriority.jpg')
save_image(edgeMap_output, 'images/edgeMap_output.jpg')
save_image(edgeMap_GT, 'images/edgeMap_GT.jpg')
save_image(edgeMap_output * cannyPriority, 'images/edgeMap_output_canny.jpg')
save_image(edgeMap_GT * cannyPriority, 'images/edgeMap_GT_canny.jpg')
save_image(errorMap, 'images/errorMap.jpg')
save_image(errorMap2, 'images/errorMap2.jpg')
|