[
  {
    "path": ".gitignore",
    "content": ".idea/\n*.pyc\n"
  },
  {
    "path": "README.md",
    "content": "# CannyEdgePytorch\n\nUses PyTorch 0.4.1 and Python 3.7 (but probably works with 2.7 also).\n\nA simple implementation of the Canny Edge Detection Algorithm (currently without hysteresis).\n\nThis project was implemented with PyTorch to take advantage of the parallelization of convolutions.\n\nThe original image:\n\n<img src=\"https://github.com/DCurro/CannyEdgePytorch/blob/master/fb_profile.jpg\" width=\"400\">\n\nFinding the gradient magnitude:\n\n<img src=\"https://github.com/DCurro/CannyEdgePytorch/blob/master/gradient_magnitude.png\" width=\"400\">\n\nEarly thresholding (to show that edge thingging matters):\n\n<img src=\"https://github.com/DCurro/CannyEdgePytorch/blob/master/thresholded.png\" width=\"400\">\n\nAnd finally, the image after non-maximum supressions:\n\n<img src=\"https://github.com/DCurro/CannyEdgePytorch/blob/master/final.png\" width=\"400\">\n"
  },
  {
    "path": "canny.py",
    "content": "from scipy.misc import imread, imsave\nimport torch\nfrom torch.autograd import Variable\nfrom net_canny import Net\n\n\ndef canny(raw_img, use_cuda=False):\n    img = torch.from_numpy(raw_img.transpose((2, 0, 1)))\n    batch = torch.stack([img]).float()\n\n    net = Net(threshold=3.0, use_cuda=use_cuda)\n    if use_cuda:\n        net.cuda()\n    net.eval()\n\n    data = Variable(batch)\n    if use_cuda:\n        data = Variable(batch).cuda()\n\n    blurred_img, grad_mag, grad_orientation, thin_edges, thresholded, early_threshold = net(data)\n\n    imsave('gradient_magnitude.png',grad_mag.data.cpu().numpy()[0,0])\n    imsave('thin_edges.png', thresholded.data.cpu().numpy()[0, 0])\n    imsave('final.png', (thresholded.data.cpu().numpy()[0, 0] > 0.0).astype(float))\n    imsave('thresholded.png', early_threshold.data.cpu().numpy()[0, 0])\n\n\nif __name__ == '__main__':\n    img = imread('fb_profile.jpg') / 255.0\n\n    # canny(img, use_cuda=False)\n    canny(img, use_cuda=True)\n"
  },
  {
    "path": "net_canny.py",
    "content": "import torch\nimport torch.nn as nn\nimport numpy as np\nfrom scipy.signal import gaussian\n\n\nclass Net(nn.Module):\n    def __init__(self, threshold=10.0, use_cuda=False):\n        super(Net, self).__init__()\n\n        self.threshold = threshold\n        self.use_cuda = use_cuda\n\n        filter_size = 5\n        generated_filters = gaussian(filter_size,std=1.0).reshape([1,filter_size])\n\n        self.gaussian_filter_horizontal = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=(1,filter_size), padding=(0,filter_size//2))\n        self.gaussian_filter_horizontal.weight.data.copy_(torch.from_numpy(generated_filters))\n        self.gaussian_filter_horizontal.bias.data.copy_(torch.from_numpy(np.array([0.0])))\n        self.gaussian_filter_vertical = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=(filter_size,1), padding=(filter_size//2,0))\n        self.gaussian_filter_vertical.weight.data.copy_(torch.from_numpy(generated_filters.T))\n        self.gaussian_filter_vertical.bias.data.copy_(torch.from_numpy(np.array([0.0])))\n\n        sobel_filter = np.array([[1, 0, -1],\n                                 [2, 0, -2],\n                                 [1, 0, -1]])\n\n        self.sobel_filter_horizontal = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=sobel_filter.shape, padding=sobel_filter.shape[0]//2)\n        self.sobel_filter_horizontal.weight.data.copy_(torch.from_numpy(sobel_filter))\n        self.sobel_filter_horizontal.bias.data.copy_(torch.from_numpy(np.array([0.0])))\n        self.sobel_filter_vertical = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=sobel_filter.shape, padding=sobel_filter.shape[0]//2)\n        self.sobel_filter_vertical.weight.data.copy_(torch.from_numpy(sobel_filter.T))\n        self.sobel_filter_vertical.bias.data.copy_(torch.from_numpy(np.array([0.0])))\n\n        # filters were flipped manually\n        filter_0 = np.array([   [ 0, 0, 0],\n                                [ 0, 1, -1],\n                                [ 0, 0, 0]])\n\n        filter_45 = np.array([  [0, 0, 0],\n                                [ 0, 1, 0],\n                                [ 0, 0, -1]])\n\n        filter_90 = np.array([  [ 0, 0, 0],\n                                [ 0, 1, 0],\n                                [ 0,-1, 0]])\n\n        filter_135 = np.array([ [ 0, 0, 0],\n                                [ 0, 1, 0],\n                                [-1, 0, 0]])\n\n        filter_180 = np.array([ [ 0, 0, 0],\n                                [-1, 1, 0],\n                                [ 0, 0, 0]])\n\n        filter_225 = np.array([ [-1, 0, 0],\n                                [ 0, 1, 0],\n                                [ 0, 0, 0]])\n\n        filter_270 = np.array([ [ 0,-1, 0],\n                                [ 0, 1, 0],\n                                [ 0, 0, 0]])\n\n        filter_315 = np.array([ [ 0, 0, -1],\n                                [ 0, 1, 0],\n                                [ 0, 0, 0]])\n\n        all_filters = np.stack([filter_0, filter_45, filter_90, filter_135, filter_180, filter_225, filter_270, filter_315])\n\n        self.directional_filter = nn.Conv2d(in_channels=1, out_channels=8, kernel_size=filter_0.shape, padding=filter_0.shape[-1] // 2)\n        self.directional_filter.weight.data.copy_(torch.from_numpy(all_filters[:, None, ...]))\n        self.directional_filter.bias.data.copy_(torch.from_numpy(np.zeros(shape=(all_filters.shape[0],))))\n\n    def forward(self, img):\n        img_r = img[:,0:1]\n        img_g = img[:,1:2]\n        img_b = img[:,2:3]\n\n        blur_horizontal = self.gaussian_filter_horizontal(img_r)\n        blurred_img_r = self.gaussian_filter_vertical(blur_horizontal)\n        blur_horizontal = self.gaussian_filter_horizontal(img_g)\n        blurred_img_g = self.gaussian_filter_vertical(blur_horizontal)\n        blur_horizontal = self.gaussian_filter_horizontal(img_b)\n        blurred_img_b = self.gaussian_filter_vertical(blur_horizontal)\n\n        blurred_img = torch.stack([blurred_img_r,blurred_img_g,blurred_img_b],dim=1)\n        blurred_img = torch.stack([torch.squeeze(blurred_img)])\n\n        grad_x_r = self.sobel_filter_horizontal(blurred_img_r)\n        grad_y_r = self.sobel_filter_vertical(blurred_img_r)\n        grad_x_g = self.sobel_filter_horizontal(blurred_img_g)\n        grad_y_g = self.sobel_filter_vertical(blurred_img_g)\n        grad_x_b = self.sobel_filter_horizontal(blurred_img_b)\n        grad_y_b = self.sobel_filter_vertical(blurred_img_b)\n\n        # COMPUTE THICK EDGES\n\n        grad_mag = torch.sqrt(grad_x_r**2 + grad_y_r**2)\n        grad_mag += torch.sqrt(grad_x_g**2 + grad_y_g**2)\n        grad_mag += torch.sqrt(grad_x_b**2 + grad_y_b**2)\n        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))\n        grad_orientation += 180.0\n        grad_orientation =  torch.round( grad_orientation / 45.0 ) * 45.0\n\n        # THIN EDGES (NON-MAX SUPPRESSION)\n\n        all_filtered = self.directional_filter(grad_mag)\n\n        inidices_positive = (grad_orientation / 45) % 8\n        inidices_negative = ((grad_orientation / 45) + 4) % 8\n\n        height = inidices_positive.size()[2]\n        width = inidices_positive.size()[3]\n        pixel_count = height * width\n        pixel_range = torch.FloatTensor([range(pixel_count)])\n        if self.use_cuda:\n            pixel_range = torch.cuda.FloatTensor([range(pixel_count)])\n\n        indices = (inidices_positive.view(-1).data * pixel_count + pixel_range).squeeze()\n        channel_select_filtered_positive = all_filtered.view(-1)[indices.long()].view(1,height,width)\n\n        indices = (inidices_negative.view(-1).data * pixel_count + pixel_range).squeeze()\n        channel_select_filtered_negative = all_filtered.view(-1)[indices.long()].view(1,height,width)\n\n        channel_select_filtered = torch.stack([channel_select_filtered_positive,channel_select_filtered_negative])\n\n        is_max = channel_select_filtered.min(dim=0)[0] > 0.0\n        is_max = torch.unsqueeze(is_max, dim=0)\n\n        thin_edges = grad_mag.clone()\n        thin_edges[is_max==0] = 0.0\n\n        # THRESHOLD\n\n        thresholded = thin_edges.clone()\n        thresholded[thin_edges<self.threshold] = 0.0\n\n        early_threshold = grad_mag.clone()\n        early_threshold[grad_mag<self.threshold] = 0.0\n\n        assert grad_mag.size() == grad_orientation.size() == thin_edges.size() == thresholded.size() == early_threshold.size()\n\n        return blurred_img, grad_mag, grad_orientation, thin_edges, thresholded, early_threshold\n\n\nif __name__ == '__main__':\n    Net()\n"
  },
  {
    "path": "requirements.txt",
    "content": "numpy==1.15.3\nPillow==5.3.0\nscipy==1.1.0\nsix==1.11.0\ntorch==0.4.1\ntorchvision==0.2.1\n"
  }
]