├── .gitignore ├── final.png ├── fb_profile.jpg ├── net_canny.pyc ├── thin_edges.png ├── thresholded.png ├── gradient_magnitude.png ├── requirements.txt ├── README.md ├── canny.py └── net_canny.py /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | *.pyc 3 | -------------------------------------------------------------------------------- /final.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCurro/CannyEdgePytorch/HEAD/final.png -------------------------------------------------------------------------------- /fb_profile.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCurro/CannyEdgePytorch/HEAD/fb_profile.jpg -------------------------------------------------------------------------------- /net_canny.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCurro/CannyEdgePytorch/HEAD/net_canny.pyc -------------------------------------------------------------------------------- /thin_edges.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCurro/CannyEdgePytorch/HEAD/thin_edges.png -------------------------------------------------------------------------------- /thresholded.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCurro/CannyEdgePytorch/HEAD/thresholded.png -------------------------------------------------------------------------------- /gradient_magnitude.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCurro/CannyEdgePytorch/HEAD/gradient_magnitude.png -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy==1.15.3 2 | Pillow==5.3.0 3 | scipy==1.1.0 4 | six==1.11.0 5 | torch==0.4.1 6 | torchvision==0.2.1 7 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # CannyEdgePytorch 2 | 3 | Uses PyTorch 0.4.1 and Python 3.7 (but probably works with 2.7 also). 4 | 5 | A simple implementation of the Canny Edge Detection Algorithm (currently without hysteresis). 6 | 7 | This project was implemented with PyTorch to take advantage of the parallelization of convolutions. 8 | 9 | The original image: 10 | 11 | 12 | 13 | Finding the gradient magnitude: 14 | 15 | 16 | 17 | Early thresholding (to show that edge thingging matters): 18 | 19 | 20 | 21 | And finally, the image after non-maximum supressions: 22 | 23 | 24 | -------------------------------------------------------------------------------- /canny.py: -------------------------------------------------------------------------------- 1 | from scipy.misc import imread, imsave 2 | import torch 3 | from torch.autograd import Variable 4 | from net_canny import Net 5 | 6 | 7 | def canny(raw_img, use_cuda=False): 8 | img = torch.from_numpy(raw_img.transpose((2, 0, 1))) 9 | batch = torch.stack([img]).float() 10 | 11 | net = Net(threshold=3.0, use_cuda=use_cuda) 12 | if use_cuda: 13 | net.cuda() 14 | net.eval() 15 | 16 | data = Variable(batch) 17 | if use_cuda: 18 | data = Variable(batch).cuda() 19 | 20 | blurred_img, grad_mag, grad_orientation, thin_edges, thresholded, early_threshold = net(data) 21 | 22 | imsave('gradient_magnitude.png',grad_mag.data.cpu().numpy()[0,0]) 23 | imsave('thin_edges.png', thresholded.data.cpu().numpy()[0, 0]) 24 | imsave('final.png', (thresholded.data.cpu().numpy()[0, 0] > 0.0).astype(float)) 25 | imsave('thresholded.png', early_threshold.data.cpu().numpy()[0, 0]) 26 | 27 | 28 | if __name__ == '__main__': 29 | img = imread('fb_profile.jpg') / 255.0 30 | 31 | # canny(img, use_cuda=False) 32 | canny(img, use_cuda=True) 33 | -------------------------------------------------------------------------------- /net_canny.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import numpy as np 4 | from scipy.signal import gaussian 5 | 6 | 7 | class Net(nn.Module): 8 | def __init__(self, threshold=10.0, use_cuda=False): 9 | super(Net, self).__init__() 10 | 11 | self.threshold = threshold 12 | self.use_cuda = use_cuda 13 | 14 | filter_size = 5 15 | generated_filters = gaussian(filter_size,std=1.0).reshape([1,filter_size]) 16 | 17 | self.gaussian_filter_horizontal = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=(1,filter_size), padding=(0,filter_size//2)) 18 | self.gaussian_filter_horizontal.weight.data.copy_(torch.from_numpy(generated_filters)) 19 | self.gaussian_filter_horizontal.bias.data.copy_(torch.from_numpy(np.array([0.0]))) 20 | self.gaussian_filter_vertical = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=(filter_size,1), padding=(filter_size//2,0)) 21 | self.gaussian_filter_vertical.weight.data.copy_(torch.from_numpy(generated_filters.T)) 22 | self.gaussian_filter_vertical.bias.data.copy_(torch.from_numpy(np.array([0.0]))) 23 | 24 | sobel_filter = np.array([[1, 0, -1], 25 | [2, 0, -2], 26 | [1, 0, -1]]) 27 | 28 | self.sobel_filter_horizontal = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=sobel_filter.shape, padding=sobel_filter.shape[0]//2) 29 | self.sobel_filter_horizontal.weight.data.copy_(torch.from_numpy(sobel_filter)) 30 | self.sobel_filter_horizontal.bias.data.copy_(torch.from_numpy(np.array([0.0]))) 31 | self.sobel_filter_vertical = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=sobel_filter.shape, padding=sobel_filter.shape[0]//2) 32 | self.sobel_filter_vertical.weight.data.copy_(torch.from_numpy(sobel_filter.T)) 33 | self.sobel_filter_vertical.bias.data.copy_(torch.from_numpy(np.array([0.0]))) 34 | 35 | # filters were flipped manually 36 | filter_0 = np.array([ [ 0, 0, 0], 37 | [ 0, 1, -1], 38 | [ 0, 0, 0]]) 39 | 40 | filter_45 = np.array([ [0, 0, 0], 41 | [ 0, 1, 0], 42 | [ 0, 0, -1]]) 43 | 44 | filter_90 = np.array([ [ 0, 0, 0], 45 | [ 0, 1, 0], 46 | [ 0,-1, 0]]) 47 | 48 | filter_135 = np.array([ [ 0, 0, 0], 49 | [ 0, 1, 0], 50 | [-1, 0, 0]]) 51 | 52 | filter_180 = np.array([ [ 0, 0, 0], 53 | [-1, 1, 0], 54 | [ 0, 0, 0]]) 55 | 56 | filter_225 = np.array([ [-1, 0, 0], 57 | [ 0, 1, 0], 58 | [ 0, 0, 0]]) 59 | 60 | filter_270 = np.array([ [ 0,-1, 0], 61 | [ 0, 1, 0], 62 | [ 0, 0, 0]]) 63 | 64 | filter_315 = np.array([ [ 0, 0, -1], 65 | [ 0, 1, 0], 66 | [ 0, 0, 0]]) 67 | 68 | all_filters = np.stack([filter_0, filter_45, filter_90, filter_135, filter_180, filter_225, filter_270, filter_315]) 69 | 70 | self.directional_filter = nn.Conv2d(in_channels=1, out_channels=8, kernel_size=filter_0.shape, padding=filter_0.shape[-1] // 2) 71 | self.directional_filter.weight.data.copy_(torch.from_numpy(all_filters[:, None, ...])) 72 | self.directional_filter.bias.data.copy_(torch.from_numpy(np.zeros(shape=(all_filters.shape[0],)))) 73 | 74 | def forward(self, img): 75 | img_r = img[:,0:1] 76 | img_g = img[:,1:2] 77 | img_b = img[:,2:3] 78 | 79 | blur_horizontal = self.gaussian_filter_horizontal(img_r) 80 | blurred_img_r = self.gaussian_filter_vertical(blur_horizontal) 81 | blur_horizontal = self.gaussian_filter_horizontal(img_g) 82 | blurred_img_g = self.gaussian_filter_vertical(blur_horizontal) 83 | blur_horizontal = self.gaussian_filter_horizontal(img_b) 84 | blurred_img_b = self.gaussian_filter_vertical(blur_horizontal) 85 | 86 | blurred_img = torch.stack([blurred_img_r,blurred_img_g,blurred_img_b],dim=1) 87 | blurred_img = torch.stack([torch.squeeze(blurred_img)]) 88 | 89 | grad_x_r = self.sobel_filter_horizontal(blurred_img_r) 90 | grad_y_r = self.sobel_filter_vertical(blurred_img_r) 91 | grad_x_g = self.sobel_filter_horizontal(blurred_img_g) 92 | grad_y_g = self.sobel_filter_vertical(blurred_img_g) 93 | grad_x_b = self.sobel_filter_horizontal(blurred_img_b) 94 | grad_y_b = self.sobel_filter_vertical(blurred_img_b) 95 | 96 | # COMPUTE THICK EDGES 97 | 98 | grad_mag = torch.sqrt(grad_x_r**2 + grad_y_r**2) 99 | grad_mag += torch.sqrt(grad_x_g**2 + grad_y_g**2) 100 | grad_mag += torch.sqrt(grad_x_b**2 + grad_y_b**2) 101 | grad_orientation = (torch.atan2(grad_y_r+grad_y_g+grad_y_b, grad_x_r+grad_x_g+grad_x_b) * (180.0/3.14159)) 102 | grad_orientation += 180.0 103 | grad_orientation = torch.round( grad_orientation / 45.0 ) * 45.0 104 | 105 | # THIN EDGES (NON-MAX SUPPRESSION) 106 | 107 | all_filtered = self.directional_filter(grad_mag) 108 | 109 | inidices_positive = (grad_orientation / 45) % 8 110 | inidices_negative = ((grad_orientation / 45) + 4) % 8 111 | 112 | height = inidices_positive.size()[2] 113 | width = inidices_positive.size()[3] 114 | pixel_count = height * width 115 | pixel_range = torch.FloatTensor([range(pixel_count)]) 116 | if self.use_cuda: 117 | pixel_range = torch.cuda.FloatTensor([range(pixel_count)]) 118 | 119 | indices = (inidices_positive.view(-1).data * pixel_count + pixel_range).squeeze() 120 | channel_select_filtered_positive = all_filtered.view(-1)[indices.long()].view(1,height,width) 121 | 122 | indices = (inidices_negative.view(-1).data * pixel_count + pixel_range).squeeze() 123 | channel_select_filtered_negative = all_filtered.view(-1)[indices.long()].view(1,height,width) 124 | 125 | channel_select_filtered = torch.stack([channel_select_filtered_positive,channel_select_filtered_negative]) 126 | 127 | is_max = channel_select_filtered.min(dim=0)[0] > 0.0 128 | is_max = torch.unsqueeze(is_max, dim=0) 129 | 130 | thin_edges = grad_mag.clone() 131 | thin_edges[is_max==0] = 0.0 132 | 133 | # THRESHOLD 134 | 135 | thresholded = thin_edges.clone() 136 | thresholded[thin_edges