File size: 4,823 Bytes
57db94b |
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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 |
from pytorch_v0 import *
import kornia
###################### CANNY ######################
def canny(img, a=100, b=200):
img = I(img).convert('L')
return I(cv2.Canny(img.cv2(), a, b))
# https://www.pyimagesearch.com/2015/04/06/zero-parameter-automatic-canny-edge-detection-with-python-and-opencv/
def canny_pis(img, sigma=0.33):
# compute the median of the single channel pixel intensities
img = I(img).convert('L').uint8(ch_last=False)
v = np.median(img)
# apply automatic Canny edge detection using the computed median
lower = int(max(0, (1.0 - sigma) * v))
upper = int(min(255, (1.0 + sigma) * v))
edged = cv2.Canny(img[0], lower, upper)
# return the edged image
return I(edged)
# https://en.wikipedia.org/wiki/Otsu%27s_method
def canny_otsu(img):
img = I(img).convert('L').uint8(ch_last=False)
high, _ = cv2.threshold(img[0], 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
low = 0.5 * high
return I(cv2.Canny(img[0], low, high))
def xdog(img, t=1.0, epsilon=0.04, phi=100, sigma=3, k=1.6):
img = I(img).convert('L').uint8(ch_last=False)
grey = np.asarray(img, dtype=np.float32)
g0 = scipy.ndimage.gaussian_filter(grey, sigma)
g1 = scipy.ndimage.gaussian_filter(grey, sigma * k)
#ans = ((1+p) * g0 - p * g1) / 255
ans = (g0 - t * g1) / 255
ans = 1 + np.tanh(phi*(ans-epsilon)) * (ans<epsilon)
return ans
def dog(img, t=1.0, sigma=1.0, k=1.6, epsilon=0.01, kernel_factor=4, clip=True):
img = I(img).convert('L').tensor()[None]
kern0 = max(2*int(sigma*kernel_factor)+1, 3)
kern1 = max(2*int(sigma*k*kernel_factor)+1, 3)
g0 = kornia.filters.gaussian_blur2d(
img, (kern0,kern0), (sigma,sigma), border_type='replicate',
)
g1 = kornia.filters.gaussian_blur2d(
img, (kern1,kern1), (sigma*k,sigma*k), border_type='replicate',
)
ans = 0.5 + t*(g1-g0) - epsilon
ans = ans.clip(0,1) if clip else ans
return ans[0].numpy()
# input: (bs,rgb(a),h,w) or (bs,1,h,w)
# returns: (bs,1,h,w)
def batch_dog(img, t=1.0, sigma=1.0, k=1.6, epsilon=0.01, kernel_factor=4, clip=True):
# to grayscale if needed
bs,ch,h,w = img.shape
if ch in [3,4]:
img = kornia.color.rgb_to_grayscale(img[:,:3])
else:
assert ch==1
# calculate dog
kern0 = max(2*int(sigma*kernel_factor)+1, 3)
kern1 = max(2*int(sigma*k*kernel_factor)+1, 3)
g0 = kornia.filters.gaussian_blur2d(
img, (kern0,kern0), (sigma,sigma), border_type='replicate',
)
g1 = kornia.filters.gaussian_blur2d(
img, (kern1,kern1), (sigma*k,sigma*k), border_type='replicate',
)
ans = 0.5 + t*(g1-g0) - epsilon
ans = ans.clip(0,1) if clip else ans
return ans
############### LOSSES + METRICS ###############
class LineRatioMetric(torchmetrics.Metric):
def __init__(
self, convert_dog=True,
t=2.0, sigma=1.0, k=1.6, epsilon=0.01, kernel_factor=4, clip=False,
**kwargs,
):
super().__init__(**kwargs)
self.convert_dog = convert_dog
self.dog_params = {
't': t, 'sigma': sigma, 'k': k, 'epsilon': epsilon,
'kernel_factor': kernel_factor, 'clip': clip,
}
self.add_state('running_sum', default=torch.tensor(0.0), dist_reduce_fx='sum')
self.add_state('running_count', default=torch.tensor(0.0), dist_reduce_fx='sum')
return
def update(self, preds: torch.Tensor, target: torch.Tensor):
if self.convert_dog:
preds = (batch_dog(preds, **self.dog_params)>0.5).float()
target = (batch_dog(target, **self.dog_params)>0.5).float()
preds = preds.sum((1,2,3))
target = target.sum((1,2,3))
dist = torch.nan_to_num(preds/target, nan=1.0, posinf=1.0, neginf=1.0)
self.running_sum += dist.sum()
self.running_count += len(dist)
return
def compute(self):
return self.running_sum.float() / self.running_count
class DoGLoss(nn.Module):
def __init__(
self, convert_dog=True, mode='l1',
t=2.0, sigma=1.0, k=1.6, epsilon=0.01, kernel_factor=4, clip=False,
):
super().__init__()
assert mode in ['l1', 'l2']
self.convert_dog = convert_dog
self.mode = mode
self.dog_params = {
't': t, 'sigma': sigma, 'k': k, 'epsilon': epsilon,
'kernel_factor': kernel_factor, 'clip': clip,
}
return
def forward(self, preds, target):
if self.convert_dog:
preds = batch_dog(preds, **self.dog_params)
target = batch_dog(target, **self.dog_params)
if self.mode=='l1':
return (preds-target).abs().mean(dim=(1,2,3))
elif self.mode=='l2':
return (preds-target).pow(2).mean(dim=(1,2,3))
|