track.py

import numpy as np
import cv2
import glob, os, sys, time, datetime

# TODO corners save when offline processing
ONLINE = True
CALIBRATE = False
RELATIVE_DESTINATION_PATH = str(datetime.date.today()) + '_distance/'
FPS = 60
THRESHOLD_WALL_VS_FLOOR = 80
THRESHOLD_ANIMAL_VS_FLOOR = 70
HD = 1280, 640
BGR_COLOR = {'red': (0,0,255),
             'green': (127,255,0),
             'blue': (255,127,0),
             'yellow': (0,127,255),
             'black': (0,0,0),
             'white': (255,255,255)}
WAIT_DELAY = 1

perspectiveMatrix = dict()
croppingPolygon = np.array([[0,0]])
croppingPolygons = dict()
tetragons = []
name = ""

RENEW_TETRAGON = True

def counterclockwiseSort(tetragon):
    tetragon = sorted(tetragon, key=lambda e: e[0])
    tetragon[0:2] = sorted(tetragon[0:2], key=lambda e: e[1])
    tetragon[2:4] = sorted(tetragon[2:4], key=lambda e: e[1], reverse=True)
    return tetragon

# TODO pointlike tetragon moving instead drawing it by clicking
# mouse callback function for drawing a cropping polygon
def drawFloorCrop(event, x, y, flags, params):
    global perspectiveMatrix, name, RENEW_TETRAGON
    imgCroppingPolygon = np.zeros_like(params['imgFloorCorners'])
    if event == cv2.EVENT_RBUTTONUP:
        cv2.destroyWindow(f'Floor Corners for {name}')
    if len(params['croppingPolygons'][name]) > 4 and event == cv2.EVENT_LBUTTONUP:
        RENEW_TETRAGON = True
        h = params['imgFloorCorners'].shape[0]
        # delete 5th extra vertex of the floor cropping tetragon
        params['croppingPolygons'][name] = np.delete(params['croppingPolygons'][name], -1, 0)
        params['croppingPolygons'][name] = params['croppingPolygons'][name] - [h,0]
        
        # Sort cropping tetragon vertices counter-clockwise starting with top left
        params['croppingPolygons'][name] = counterclockwiseSort(params['croppingPolygons'][name])
        # Get the matrix of perspective transformation
        params['croppingPolygons'][name] = np.reshape(params['croppingPolygons'][name], (4,2))
        tetragonVertices = np.float32(params['croppingPolygons'][name])
        tetragonVerticesUpd = np.float32([[0,0], [0,h], [h,h], [h,0]])
        perspectiveMatrix[name] = cv2.getPerspectiveTransform(tetragonVertices, tetragonVerticesUpd)
    if event == cv2.EVENT_LBUTTONDOWN:
        if len(params['croppingPolygons'][name]) == 4 and RENEW_TETRAGON:
            params['croppingPolygons'][name] = np.array([[0,0]])
            RENEW_TETRAGON = False
        if len(params['croppingPolygons'][name]) == 1:
            params['croppingPolygons'][name][0] = [x,y]
        params['croppingPolygons'][name] = np.append(params['croppingPolygons'][name], [[x,y]], axis=0)
    if event == cv2.EVENT_MOUSEMOVE and not (len(params['croppingPolygons'][name]) == 4 and RENEW_TETRAGON):
        params['croppingPolygons'][name][-1] = [x,y]
        if len(params['croppingPolygons'][name]) > 1:
            cv2.fillPoly(
                imgCroppingPolygon,
                [np.reshape(
                    params['croppingPolygons'][name],
                    (len(params['croppingPolygons'][name]),2)
                )],
                BGR_COLOR['green'], cv2.LINE_AA)
            imgCroppingPolygon = cv2.addWeighted(params['imgFloorCorners'], 1.0, imgCroppingPolygon, 0.5, 0.)
            cv2.imshow(f'Floor Corners for {name}', imgCroppingPolygon)

def angle_cos(p0, p1, p2):
    d1, d2 = (p0 - p1).astype('float'), (p2 - p1).astype('float')
    return np.abs(np.dot(d1, d2) / np.sqrt(np.dot(d1, d1) * np.dot(d2, d2)))

def floorCrop(filename):
    global perspectiveMatrix, tetragons, name, croppingPolygons
    name = os.path.splitext(filename)[0]
    cap = cv2.VideoCapture(filename)
    h, w = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))

    # Take first non-null frame and find corners within it
    ret, frame = cap.read()
    while not frame.any():
        ret, frame = cap.read()

    frame = frame[:, w-h : w]

    frameGray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    kernelSize = (5,5)
    frameBlur = cv2.GaussianBlur(frameGray, kernelSize, 0)
    retval, mask = cv2.threshold(frameBlur, THRESHOLD_WALL_VS_FLOOR, 255, cv2.THRESH_BINARY_INV)
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

    tetragons = []
    HALF_AREA = 0.5 * h * h
    for contour in contours:
        contourPerimeter = cv2.arcLength(contour, True)
        hull = cv2.convexHull(contour)
        contour = cv2.approxPolyDP(hull, 0.02 * contourPerimeter, True)

        # If the contour is convex tetragon
        # and its area is above a half of total frame area,
        # then it's most likely the floor
        if len(contour) == 4 and cv2.contourArea(contour) > HALF_AREA:
            contour = contour.reshape(-1, 2)
            max_cos = np.max([angle_cos(contour[i], contour[(i + 1) % 4], contour[(i + 2) % 4]) for i in range(4)])
            if max_cos < 0.3:
                tetragons.append(contour)
    frameGray = cv2.cvtColor(frameGray, cv2.COLOR_GRAY2BGR)
    imgSquare = np.zeros_like(frameGray)
    cv2.fillPoly(imgSquare, tetragons, BGR_COLOR['red'], cv2.LINE_AA)
    # cv2.add(frameGray, imgSquare / 2, frameGray)
    cv2.drawContours(frameGray, tetragons, -1, BGR_COLOR['red'], 2, cv2.LINE_AA)

    if len(tetragons) > 0:
        tetragonVertices = tetragons[0]
    else:
        tetragonVertices = np.float32([[0,0], [0,h], [h,h], [h,0]])
    # Sort the cropping tetragon vertices according to the following order:
    # [left,top], [left,bottom], [right,bottom], [right,top]
    tetragonVertices = counterclockwiseSort(tetragonVertices)
    croppingPolygons[name] = tetragonVertices
    tetragonVertices = np.float32(tetragonVertices)
    tetragonVerticesUpd = np.float32([[0,0], [0,h], [h,h], [h,0]])
    perspectiveMatrix[name] = cv2.getPerspectiveTransform(np.float32(croppingPolygons[name]), tetragonVerticesUpd)
    frame = cv2.warpPerspective(frame, perspectiveMatrix[name], (h,h))
    
    imgFloorCorners = np.hstack([frame, frameGray])
    cv2.imshow(f'Floor Corners for {name}', imgFloorCorners)
    cv2.setMouseCallback(
        f'Floor Corners for {name}',
        drawFloorCrop,
        {'imgFloorCorners': imgFloorCorners, 'croppingPolygons': croppingPolygons},
        )
    k = cv2.waitKey(0)
    if k == 27:
        sys.exit()
    cv2.destroyWindow(f'Floor Corners for {name}')
    return tetragonVertices, perspectiveMatrix[name]

def trace(filename):
    global perspectiveMatrix, croppingPolygons, tetragons, name, WAIT_DELAY
    # croppingPolygons[name] = np.array([[0,0]])
    name = os.path.splitext(filename)[0]
    cap = cv2.VideoCapture(filename)
    h, w = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))

    # Take first non-null frame and find corners within it
    ret, frame = cap.read()
    while not frame.any():
        ret, frame = cap.read()

    background = frame.copy()
    i_frame = 1
    n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    while frame is not None:
        ret, frame = cap.read()
        if frame is None:
            break
        background = cv2.addWeighted(frame, 0.5 * (1 - i_frame / n_frames),
                                background, 0.5 * (1 + i_frame / n_frames), 0)
        i_frame += 1
    cap = cv2.VideoCapture(filename)
    ret, frame = cap.read()

    frame = frame[:, w-h : w]
    
    # floorCrop(filename)

    video = cv2.VideoWriter(f'{RELATIVE_DESTINATION_PATH}timing/{name}_trace.avi',
        cv2.VideoWriter_fourcc(*'X264'),
        FPS, HD, cv2.INTER_LINEAR)
    imgTrack = np.zeros_like(frame)

    start = time.time()
    distance = _x = _y = 0
    while frame is not None:
        ret, frame = cap.read()
        
        if frame is None:   # not logical
            break
        frameColor = frame[:, w-h : w].copy()
        frame = cv2.subtract(frame, background)

        t = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000.
        frame = frame[:, w-h : w]
        if len(croppingPolygons[name]) == 4:
            cv2.drawContours(frameColor, [np.reshape(croppingPolygons[name], (4,2))], -1, BGR_COLOR['red'], 2, cv2.LINE_AA)
        else:
            cv2.drawContours(frameColor, tetragons, -1, BGR_COLOR['red'], 2, cv2.LINE_AA)

        frame = cv2.warpPerspective(frame, perspectiveMatrix[name], (h,h))

        frameGray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        kernelSize = (25,25)
        frameBlur = cv2.GaussianBlur(frameGray, kernelSize, 0)
        _, thresh = cv2.threshold(frameBlur, THRESHOLD_ANIMAL_VS_FLOOR, 255, cv2.THRESH_BINARY)
        contours, hierarchy = cv2.findContours(thresh.copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
        if len(contours) < 1:   # TODO more pythonic way of the check
            continue

        # Find a contour with the biggest area (animal most likely)
        contour = contours[np.argmax(list(map(cv2.contourArea, contours)))]
        M = cv2.moments(contour)
        if M['m00'] == 0:
            continue
        x = int(M['m10'] / M['m00'])
        y = int(M['m01'] / M['m00'])
        if _x == 0 and _y == 0:
            _x = x
            _y = y
        distance += np.sqrt(((x - _x) / float(h))**2 + ((y - _y) / float(h))**2)

        if ONLINE:
            # Draw the most acute angles of the contour (tail/muzzle/paws of the animal)
            hull = cv2.convexHull(contour)
            imgPoints = np.zeros(frame.shape,np.uint8)
            for i in range(2, len(hull) - 2):
                if np.dot(hull[i][0] - hull[i-2][0], hull[i][0] - hull[i+2][0]) > 0:
                    imgPoints = cv2.circle(imgPoints, (hull[i][0][0],hull[i][0][1]), 5, BGR_COLOR['yellow'], -1, cv2.LINE_AA)

            # Draw a contour and a centroid of the animal
            cv2.drawContours(imgPoints, [contour], 0, BGR_COLOR['green'], 2, cv2.LINE_AA)
            imgPoints = cv2.circle(imgPoints, (x,y), 5, BGR_COLOR['black'], -1)
            
            # Draw a track of the animal
            # imgTrack = cv2.add(np.zeros_like(imgTrack), cv2.line(imgTrack, (x,y), (_x,_y),
                # (255, 127, int(cap.get(cv2.CAP_PROP_POS_AVI_RATIO) * 255)), 1, cv2.LINE_AA))
            imgTrack = cv2.addWeighted(np.zeros_like(imgTrack), 0.85, cv2.line(imgTrack, (x,y), (_x,_y),
                (255, 127, int(cap.get(cv2.CAP_PROP_POS_AVI_RATIO) * 255)), 1, cv2.LINE_AA), 0.98, 0.)
            imgContour = cv2.add(imgPoints, imgTrack)

            frame = cv2.bitwise_and(frame, frame, mask=thresh)
            frame = cv2.addWeighted(frame, 0.4, imgContour, 1.0, 0.)

            cv2.putText(frame, 'Distance ' + str('%.2f' % distance),
                (190,420), cv2.FONT_HERSHEY_DUPLEX, 1, BGR_COLOR['white'])
            cv2.putText(frame, 'Time ' + str('%.0f sec' % (cap.get(cv2.CAP_PROP_POS_MSEC) / 1000.)),
                (200,450), cv2.FONT_HERSHEY_DUPLEX, 1, BGR_COLOR['white'])
            cv2.circle(frame, (x,y), 5, BGR_COLOR['black'], -1, cv2.LINE_AA)
            layout = np.hstack((frame, frameColor))

            cv2.imshow(f'Open Field Trace of {name}', layout)
            video.write(cv2.resize(layout, HD))

            k = cv2.waitKey(WAIT_DELAY) & 0xff
            if k == 27:
                break
            if k == 32:
                if WAIT_DELAY == 1:
                    WAIT_DELAY = 0  # pause
                else:
                    WAIT_DELAY = 1  # play as fast as possible
        _x = x
        _y = y
    cv2.destroyAllWindows()
    cap.release()

    if ONLINE:
        video.release()
        cv2.imwrite(RELATIVE_DESTINATION_PATH + 'traces/' + name + '_[distance]=%.2f' % distance +
            '_[time]=%.1fs' % t + '.png', cv2.resize(imgTrack, (max(HD), max(HD))))

    print(filename + '\tdistance %.2f\t' % distance + 'processing/real time %.1f' % float(time.time() - start) + '/%.1f s' % t)
    file.write(name + ',%.2f' % distance + ',%.1f\n' % t)
    file.close()

if len(sys.argv) > 1 and '--online' in sys.argv:
    ONLINE = True
if not os.path.exists(RELATIVE_DESTINATION_PATH + 'traces'):
    os.makedirs(RELATIVE_DESTINATION_PATH + 'traces')
if not os.path.exists(RELATIVE_DESTINATION_PATH + 'timing'):
    os.makedirs(RELATIVE_DESTINATION_PATH + 'timing')
file = open(RELATIVE_DESTINATION_PATH + 'distances.csv', 'w')
file.write('animal,distance [unit of the box side],run time [seconds]\n')
file.close()

for filename in glob.glob('*.avi'):
    floorCrop(filename)
for filename in glob.glob('*.avi'):
    file = open(RELATIVE_DESTINATION_PATH + 'distances.csv', 'a')
    trace(filename)