import cv2
import numpy as np
from skimage import morphology
[docs]def seg2edge(seg, thin=False):
grad = np.gradient(seg.astype(np.float32))
grad = np.sqrt(grad[0]**2 + grad[1]**2)
edge = grad != 0
if thin:
edge = morphology.thin(edge)
return edge
[docs]def sobel(x, kernel_size=3):
sobelx = cv2.Sobel(x, cv2.CV_64F, 1, 0, ksize=kernel_size)
sobely = cv2.Sobel(x, cv2.CV_64F, 0, 1, ksize=kernel_size)
sobel = np.sqrt(sobelx**2 + sobely**2)
return sobel, sobely, sobelx
[docs]def flux2angle(flux):
"""
computing orientation angle for each pixel given a flux
input:
flux: a [2, H, W] tensor represents the flux vector of each position
return:
angle: a [H, W] matrix representing angle of each location
"""
C, H, W = flux.shape
assert C == 2
top_half = flux[0,...] >= 0 # y >= 0, \theta <= \pi
bottom_half = flux[0,...] < 0 # y < 0, \theta > \pi
unit = np.zeros((2, H, W), dtype=np.float32)
unit[1,...] = 1 # unit vector: (y=0, x=1)
cos = (flux * unit).sum(axis=0)
acos = np.arccos(cos)
angle = np.zeros((H, W), dtype=np.float32)
angle[top_half] = acos[top_half]
angle[bottom_half] = 2*np.pi - acos[bottom_half]
return angle
[docs]def angle2flux(angle):
"""
The inverse of `flux2angle`
"""
return np.dstack((np.sin(angle), np.cos(angle))).transpose((2,0,1))
[docs]def dense2flux(density, kernel_size=9, normalize=True):
"""
Compute flux field of a density map
input:
density: density map with shape [H, W]
output:
flux: flux field of shape [2, H, W]
"""
sobelx = cv2.Sobel(density, cv2.CV_64F, 1, 0, ksize=kernel_size)
sobely = cv2.Sobel(density, cv2.CV_64F, 0, 1, ksize=kernel_size)
flux = np.dstack((sobely, sobelx)).transpose((2,0,1))
if normalize:
norm = np.expand_dims(np.sqrt((flux**2).sum(axis=0)), axis=0)
norm[norm==0] = 1
flux /= norm
return flux
[docs]def quantize_angle(angle, num_bins=8):
"""
angle: angle map with shape [H, W]
num_bins: number of quantization bins
"""
# clamp angle
angle[angle>=np.pi*2] = np.pi*2 - 1e-5
q = np.round(angle / (np.pi*2/num_bins)).astype(np.uint8)
q[q==num_bins] = 0
return q
[docs]def dequantize_angle(q, num_bins=8):
"""
q: quantized angles (0~num_bins-1)
num_bins: number of quantized levels
"""
assert q.min() >= 0 and q.max() < num_bins
angle = q * (np.pi*2 / num_bins)
return angle