# CCC Tracking With Orientation Calculation

import cv2

import numpy as np

def main():

video_capture = cv2.VideoCapture("fiveCCC.mp4")

width = int(video_capture.get(3))

height = int(video_capture.get(4))

fourcc = cv2.VideoWriter_fourcc('M', 'P', '4', '2')

videoWriter = cv2.VideoWriter("output.avi", fourcc=fourcc, fps=30.0,

frameSize=(width, height))

frame = 1

track_points = np.full([5, 1, 2], fill_value=-1)

# coordinates of CCC targets in inches

world_coord = np.array([[-3.7, -2.275, 0], [0, -2.275, 0], [3.7, -2.275, 0],

[-3.7, 2.275, 0], [3.7, 2.275, 0]]).astype("float32")

# camera properties:

fx = 531

fy = fx

cx = 320

cy = 240

K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]).astype("float32")

while True:

got_image, bgr_image = video_capture.read()

if not got_image:

break # End of video; exit the while loop

# convert to greyscale

gray_img = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2GRAY)

#threshold image using otsu's method

thresh, binary_img_white = cv2.threshold(gray_img, thresh=0, maxval=255,

type=cv2.THRESH_BINARY + cv2.THRESH_OTSU)

# create the complement of the binary image so that the outer circles become white

binary_img_black = cv2.bitwise_not(binary_img_white)

# get connected component stats for binary image and complement

connected_comp_white = cv2.connectedComponentsWithStats(binary_img_white)

connected_comp_black = cv2.connectedComponentsWithStats(binary_img_black)

# get centroid matrices from overall result of connectedComp function

centroids_white = connected_comp_white[3]

centroids_black = connected_comp_black[3]

# define threshold for acceptable distance between centroids

centroid_thresh = 1

# this value assigns numbers to the points of the CCC targets for the extra credit

point_num = 0

for i in range(connected_comp_white[0]):

for j in range(connected_comp_black[0]):

# get distance between centroids of each connected component

distance = get_distance(centroids_black[j], centroids_white[i])

if distance < centroid_thresh: # if this distance is less than the threshold:

cv2.drawMarker(bgr_image, position=(int(centroids_white[i, 0]), int(centroids_white[i, 1])),

color=(0, 255, 255), markerType=cv2.MARKER_CROSS, thickness=2,

markerSize=round(30 - frame/30))

# on first frame, fill track_points matrix with centroid values that have been identified as CCCs

if frame == 1:

track_points[point_num] = centroids_white[i]

cv2.putText(bgr_image, text=str(point_num),

org=(int(centroids_white[i, 0]), int(centroids_white[i, 1])),

fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(0, 0, 255), thickness=2)

# increment point_num so that the next entry of track_points can be filled

point_num += 1

else: # if it is not first frame:

goodPoint = True # this value is used to determine if a random point has been detected

# for one frame that is not one of the CCCs

# set min_distance to distance between current centroid and first track_point

min_distance = get_distance((track_points[0, 0, 0], track_points[0, 0, 1]),

centroids_white[i])

min_index = 0 # set min_index to first index of track_points

for k in range(5): # loop through all track_points (These will be the CCCs from frame 1)

new_distance = get_distance((track_points[k, 0, 0], track_points[k, 0, 1]),

centroids_white[i])

if new_distance < min_distance: # if a new min distance has been found, replace it

min_distance = new_distance

min_index = k # min index will be location of closest centroid from

# previous frame, as they do not move much each frame

if min_distance > 10: # sometimes random other points are identified as CCCs

goodPoint = False # this throws them out

if goodPoint: # only update CCC's previous location if it is a good point

track_points[min_index] = centroids_white[i]

cv2.putText(bgr_image, text=str(min_index), # draw numbers

org=(int(centroids_white[i, 0]), int(centroids_white[i, 1])),

fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1 - frame/600,

color=(0, 0, 255), thickness=2)

cv2.putText(bgr_image, text=str(frame), org=(10, 30), # draw frame count

fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(0, 255, 255), thickness=2)

# Dimensions: [m, n, 3]

track_points_resize = track_points.transpose()

# Dimensions: [3, n, m]

track_points_resize = track_points.reshape(track_points.shape[0],

(track_points.shape[1] * track_points.shape[2])).astype("float32")

success, rotation, translation = cv2.solvePnP(world_coord, track_points_resize, K, distCoeffs=None)

# Draw coordinate axes onto the image. Scale the length of the axes

# according to the size of the model, so that the axes are visible.

W = np.amax(world_coord, axis=0) - np.amin(world_coord, axis=0)

# Size of model in X,Y,Z

L = np.linalg.norm(W)

# Length of the diagonal of the bounding box

d = L/5

# This will be the length of the coordinate axes

pAxes = np.float32([[0, 0, 0], # origin

[d, 0, 0], # x axis

[0, d, 0], # y axis

[0, 0, d] # z axis

])

pImg, J = cv2.projectPoints(

objectPoints=pAxes, # Numpy array, size (N,3)

rvec=rotation, tvec=translation,

cameraMatrix=K, distCoeffs=None)

pImg = pImg.reshape(-1, 2) # reshape from size (N,1,2) to (N,2)

cv2.line(bgr_image, tuple(np.int32(pImg[0])), tuple(np.int32(pImg[3])), (255, 0, 0), 2) # z

cv2.line(bgr_image, tuple(np.int32(pImg[0])), tuple(np.int32(pImg[1])), (0, 0, 255), 2) # x

cv2.line(bgr_image, tuple(np.int32(pImg[0])), tuple(np.int32(pImg[2])), (0, 255, 0), 2) # y

# draw pose

translation_string = "Translation: " + str(translation[0]) + str(translation[1]) + str(translation[2])

rotation_string = "Rotation: " + str(rotation[0]) + str(rotation[1]) + str(rotation[2])

cv2.putText(bgr_image, text=translation_string, org=(10, 60),

fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=.6, color=(0, 0, 0), thickness=2)

cv2.putText(bgr_image, text=rotation_string, org=(10, 90),

fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=.6, color=(0, 0, 0), thickness=2)

cv2.imshow("CCC Tracking", bgr_image)

cv2.waitKey(30)

frame += 1

videoWriter.write(bgr_image)

videoWriter.release()

def get_distance(point1, point2): # this function handles the right triangle math for two points

distance_x = point1[1] - point2[1]

distance_y = point1[0] - point2[0]

distance = (distance_y ** 2 + distance_x ** 2) ** 0.5

return distance

if __name__ == "__main__":

main()