Counting gear teeth

count_gear_teeth.py

import cv2, numpy as np, math
raw_image = cv2.imread('/home/stephen/Desktop/gear6.jpg')
bilateral_filtered_image = cv2.bilateralFilter(raw_image, 5, 175, 175)
# Added median blurring to improve edge detection
median_blurred_images = cv2.medianBlur(bilateral_filtered_image, 5)
edge_detected_image = cv2.Canny(median_blurred_images, 75, 200)
# Switched from RETR_TREE to RETR_EXTERNAL to only extract most outer contours
_, contours, _ = cv2.findContours(edge_detected_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
contour_list = []
for contour in contours:
    approx = cv2.approxPolyDP(contour,0.01*cv2.arcLength(contour,True),True)
    area = cv2.contourArea(contour)
    if ((len(approx) > 5) & (len(approx) < 25) & (area > 50) ):
        contour_list.append(contour)
c = max(contours, key = cv2.contourArea)
raw_image = np.zeros_like(raw_image)
cv2.drawContours(raw_image, contour_list, -1, (255, 255, 255), 2)

# Show the user the contour image
# The parameters for blur and canny may need
# to be altered if the gear is not detected correctly 
cv2.imshow('img', raw_image)
cv2.waitKey(0)

# Get centroid
M = cv2.moments(c)
cX = int(M["m10"] / M["m00"])
cY = int(M["m01"] / M["m00"])
_, gear_radius = cv2.minEnclosingCircle(c)
# Add some padding and crop the image
gear_radius = int(gear_radius*.9)
#raw_image = raw_image[gear_radius:gear_radius*4, :]

#cY -= gear_radius
centroid = cX, cY


# Distance function
def distance(a,b): return math.sqrt((a[0]-b[0])**2 + (a[1]-b[1])**2)


vid_writer = cv2.VideoWriter('/home/stephen/Desktop/gear.avi',
                             cv2.VideoWriter_fourcc('M','J','P','G'),
                             60, (raw_image.shape[1], raw_image.shape[0]))


# Start at angle 0, and increment the angle 1/200 rad
angle = 0
increment = 1/200
# Create a list for the distances from the centroid to the edge of the gear tooth
distances = []
# Create an image for display purposes
display_image = raw_image.copy()
# Sweep around the circle (until one full revolution)
while angle < 2*math.pi:
    # Compute a ray from the center of the circle with the current angle
    img_size = max(raw_image.shape)
    ray_end = int(math.sin(angle) * img_size + cX), int(math.cos(angle) * img_size + cY)
    center = cX, cY
    # Create mask
    mask = np.zeros((raw_image.shape[0], raw_image.shape[1]), np.uint8)
    # Draw a line on the mask
    cv2.line(mask, center, ray_end, 255, 2)
    # Mask out the gear slice (this is the portion of the gear the us below the line)
    gear_slice = cv2.bitwise_and(raw_image, raw_image, mask = mask)
    # Threshold the image
    _, thresh = cv2.threshold(cv2.cvtColor(gear_slice, cv2.COLOR_BGR2GRAY), 0 , 255, 0)
    # Find the contours in the edge_slice
    _, edge_slice_contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # Get the center of the edge slice contours
    try: M = cv2.moments(max(edge_slice_contours, key = cv2.contourArea))
    except:
        print("Contours were not detected correctly. Please change parameters.")
        break
    edge_location = int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"])
    cv2.circle(display_image, edge_location, 0, (0,255,0), 4)
    # Find the distance from the center of the gear to the edge of the gear...at this specific angle
    edge_center_distance = distance(center, edge_location)
    # Find the xy coordinates for this point on the graph - draw blue circle
    graph_point = int(angle*0.5*raw_image.shape[1]/math.pi), int(edge_center_distance+ 1.5*gear_radius)
    cv2.circle(display_image, graph_point, 0, (0,255,0), 2)    
    # Add this distance to the list of distances
    distances.append(-edge_center_distance)
    # Create a temporary image and draw the ray on it
    temp = display_image.copy()
    cv2.line(temp, ray_end, (cX,cY), (0,0,255), 2)
    # Show the image and wait
    cv2.imshow('raw_image', temp)
    vid_writer.write(temp)
    k = cv2.waitKey(1)
    if k == 27: break
    # Increment the angle
    angle += increment
# Clean up
cv2.destroyAllWindows()

# Now it's time to process the data with Scipy and display it with Matplotlib
import matplotlib.pyplot as plt
plt.plot(distances)
plt.show()
import scipy.fftpack
# Calculate the Fourier transform
yf = scipy.fftpack.fft(distances)
fig, ax = plt.subplots()
# Plot the relevant part of the Fourier transform (a gear will have between 2 and 200 teeth)
ax.plot(yf[2:200])
plt.show()
# Find the peak in the Fourier transform
num_teeth = list(yf).index(max(yf[2:200])) - 1
print('Number of teeth in this gear: ' + str(num_teeth))