[ECE558] Final Project Blob Detection

import cv2
import numpy as np
import math
import matplotlib.pyplot as plt
import argparse
from collections import Counter
np.fft.restore_all()


def resize(img, newY, newX):

    res = np.zeros([newY, newX], dtype = 'uint8')
    i = 0
    factY = newY/img.shape[0]
    factX = newX/img.shape[1]
    while(i < newY):
        j = 0
        while(j < newX):
            res[i][j] = img[math.floor(i/factY)][math.floor(j/factX)]
            j += 1
        i += 1

    return np.uint8(res)

def conv2(img, kernel, padding_style):

    h, w = kernel.shape
    yO, xO = img.shape
    paddedImg = np.zeros([yO + math.floor(h/2) + math.floor(h/2), xO + math.floor(w/2) + math.floor(w/2)], dtype='float32')
    paddedImg[math.floor(h/2) : yO + math.floor(h/2),math.floor(w/2) : xO + math.floor(w/2)] = img[:,:]

    if(padding_style == 'zero-padding'):
        paddedImg = paddedImg
    elif(padding_style == 'wrap-around'):
        i = math.floor(w/2) - 1
        j = xO + math.floor(w/2) - 1
        while(i >= 0):
            paddedImg[:,i] = paddedImg[:,j]
            j -= 1
            i -= 1
        i = math.floor(h/2) - 1
        j = yO + math.floor(w/2) - 1
        while(i >= 0):
            paddedImg[i,:] = paddedImg[j,:]
            j -= 1
            i -= 1
        i = math.floor(w/2) + xO
        j = math.floor(w/2)
        while(i < paddedImg.shape[1]):
            paddedImg[:,i] = paddedImg[:,j]
            j += 1
            i += 1
        i = math.floor(h/2) + yO
        j = math.floor(h/2)
        while(i < paddedImg.shape[0]):
            paddedImg[i,:] = paddedImg[j,:]
            j += 1
            i += 1
    elif(padding_style == 'copy-edge'):
        i = math.floor(w/2) - 1
        j = math.floor(w/2)
        while(i >= 0):
            paddedImg[:,i] = paddedImg[:,j]
            i -= 1
        i = math.floor(h/2) - 1
        j = math.floor(h/2)
        while(i >= 0):
            paddedImg[i,:] = paddedImg[j,:]
            i -= 1
        i = math.floor(w/2) + xO
        j = xO + math.floor(w/2) - 1
        while(i < paddedImg.shape[1]):
            paddedImg[:,i] = paddedImg[:,j]
            i += 1
        i = math.floor(h/2) + yO
        j = yO + math.floor(h/2) - 1
        while(i < paddedImg.shape[0]):
            paddedImg[i,:] = paddedImg[j,:]
            i += 1
    elif(padding_style == 'reflect-across-edge'):
        i = math.floor(w/2) - 1
        j = math.floor(w/2)
        while(i >= 0):
            paddedImg[:,i] = paddedImg[:,j]
            j += 1
            i -= 1
        i = math.floor(h/2) - 1
        j = math.floor(h/2)
        while(i >= 0):
            paddedImg[i,:] = paddedImg[j,:]
            j += 1
            i -= 1
        i = math.floor(w/2) + xO
        j = xO + math.floor(w/2) - 1
        while(i < paddedImg.shape[1]):
            paddedImg[:,i] = paddedImg[:,j]
            j -= 1
            i += 1
        i = math.floor(h/2) + yO
        j = yO + math.floor(h/2) - 1
        while(i < paddedImg.shape[0]):
            paddedImg[i,:] = paddedImg[j,:]
            j -= 1
            i += 1

    img = paddedImg
    y, x = img.shape
    c = 1
    imgF = np.fft.fft2(img)
    imgFS = np.fft.fftshift(imgF)
    kernelF = np.fft.fft2(kernel, [y,x])
    kernelFS = np.fft.fftshift(kernelF)
    imgF = imgFS * kernelFS
    imgishift = np.fft.ifftshift(imgF)
    img = np.fft.ifft2(imgishift)
    img = np.abs(img)
    img = img[h-1:h-1+yO, w-1:w-1+xO]
    return img

def laplacian_of_gaussian_kernel(size_x, size_y, sigma):

    x, y = np.mgrid[-size_x//2 + 1:size_x//2 + 1,-size_y//2 + 1:size_y//2 + 1]
    kernel = (-1 / (np.pi * (sigma ** 4))) * (1 - (((x ** 2) + (y ** 2))/(2 * (sigma ** 2)))) * np.exp((-1) * (((x ** 2) + (y ** 2))/(2 * (sigma ** 2))))
    kernel = (kernel + kernel.mean()) * (sigma ** 2)
    return kernel

def generate_scale_space(img, num_of_scales, sigma, scale_multiplier):

    scale_space = []
    sigmas = []
    for i in range(0, num_of_scales):
        scaled_sigma = sigma * (scale_multiplier ** i)
        kernel_size = max(1,math.floor(6*scaled_sigma))
        if(kernel_size % 2 == 0):
            kernel_size += 1
        kernel = laplacian_of_gaussian_kernel(kernel_size, kernel_size, scaled_sigma)
        current = conv2(img, kernel, 'copy-edge')
        scale_space.append(current.astype('float64'))
        sigmas.append(scaled_sigma)
        # cv2.imshow('Scale Space', np.uint8(current))
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()

    scale_space = np.array([i for i in scale_space])
    return scale_space, sigmas

def nms_3D(scale_space_2D_NMS, threshold, sigma_scale, up_scale):

    y, x = scale_space_2D_NMS[0].shape[0], scale_space_2D_NMS[0].shape[1]
    num_of_scales = len(scale_space_2D_NMS)
    coordinates = []
    for k in range(0, num_of_scales):
        current = np.zeros([y, x], dtype = 'float64')
        for i in range(1, y):
            for j in range(1, x):
                img_slice = scale_space_2D_NMS[max(0, k-1):min(num_of_scales, k+2),i-1:i+2,j-1:j+2]
                result = np.max(img_slice.flatten())
                if(np.uint8(result) >= threshold and result == scale_space_2D_NMS[k,i,j] and Counter(img_slice.flatten())[result] == 1):
                    coordinates.append((up_scale * i, up_scale * j, up_scale * sigma_scale[k]))
                    current[i][j] = result

    coordinates = list(set(coordinates))
    return coordinates


def build_octaves(img, num_of_scales, sigma, scale_multiplier, num_of_octaves):

    octaves = []
    sigma_scales = []
    for i in range(num_of_octaves):
        scale_space, sigmas = generate_scale_space(img, num_of_scales, sigma, scale_multiplier)
        octaves.append(scale_space)
        sigma_scales.append(sigmas)
        if(i < num_of_octaves - 1):
            width = int(img.shape[1] * 50 / 100)
            height = int(img.shape[0] * 50 / 100)
            dim = (width, height)
            img = resize(img, height, width)
            sigma = 2 * sigma
    return octaves, sigma_scales


def non_max_draw(img, coordinates, label):
    y, x, c = img.shape
    done = np.zeros([y, x, c], dtype = 'uint8')
    coordinates = np.array(coordinates)
    coordinates = coordinates[coordinates[:,2].argsort()]
    sigmasss = []
    i = 0
    while(i < len(coordinates)):
        y, x, sigma_k = coordinates[i]
        if(done[int(y),int(x),0] == 255):
            i += 1
            continue
        sigmasss.append(np.float32(sigma_k * (2 ** 0.5)))
        img = cv2.circle(img, (int(x), int(y)), np.float32((sigma_k * (2 ** 0.5)) * 1.2), (0, 0, 255), 1)
        done = cv2.circle(done, (int(x), int(y)), np.float32(sigma_k * (4 ** 0.5)), (255, 255, 255), -1)
        i += 1
    cv2.imshow(label, np.uint8(img))

def detect_bobs(img, num_of_octaves, num_of_scales, sigma, scale_multiplier):
    octaves, sigma_scale = build_octaves(img, num_of_scales, sigma, scale_multiplier, num_of_octaves)
    coordinates = []
    # set particular index to change threshold for that index, eg: 3rd octave, index = 2, set thresholds[2] = 0
    thresholds = [32] * num_of_octaves
    for i, octave in enumerate(octaves):
        coordinates += nms_3D(octave, thresholds[i], sigma_scale[i], 2 ** i)

    img = cv2.cvtColor(img.astype('uint8'), cv2.COLOR_GRAY2BGR)
    img = non_max_draw(img, coordinates, 'Blobs')
    cv2.waitKey(0)
    cv2.destroyAllWindows()

# driver code
parser = argparse.ArgumentParser(description='Blob Detection')
requiredArgs = parser.add_argument_group('required arguments')
requiredArgs.add_argument('-i', action='store', dest='img_path',help='Enter image path (relative path)', required=True)
requiredArgs.add_argument('-o', action='store', dest='num_of_octaves', type=str, help='Enter number of ocatves', required=True)
requiredArgs.add_argument('-n', action='store', dest='num_of_scales', type=str, help='Enter number of layers in the octave', required=True)
requiredArgs.add_argument('-s', action='store', dest='sigma', type=str, help='Enter starting sigma', required=True)
requiredArgs.add_argument('-k', action='store', dest='scale_multiplier', type=str, help='Enter scale multiplier', required=True)

args = parser.parse_args()
img_path = args.img_path
num_of_octaves = int(args.num_of_octaves)
num_of_scales = int(args.num_of_scales)
sigma = float(args.sigma)
scale_multiplier = float(args.scale_multiplier)

img = cv2.imread(img_path, 0)
img = img.astype('float64')

detect_bobs(img, num_of_octaves, num_of_scales, sigma, scale_multiplier)


# to display log kernel
# kernel = laplacian_of_gaussian_kernel(81, 81, 13.5)
# plt.imshow(kernel, interpolation='none')
# plt.colorbar()
# plt.show()