rainbow.py

import numpy as np
import cv2
import pandas as pd
import numpy.polynomial.polynomial as poly
import math

# Read Source Data
cap = cv2.VideoCapture('/home/stephen/Desktop/ss5_id_412.MP4')
df = pd.read_csv('/home/stephen/Desktop/ss5_id_412.csv')
#Write Video Out
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter('/home/stephen/Desktop/parabola.avi', fourcc, 120.0, (480,848))

# Define x and y as the first two columns of the spreadsheet
x = df['0']
y = df['1']
# Filter the data to smooth it out
from scipy.signal import savgol_filter
x = savgol_filter(x, 7, 3)
y = savgol_filter(y, 7, 3)

# Start at frame number 0
frameNum = 0
# Define trail length
trail = 30
# Colors of the rainbow
rainbow = [(0,0,255), (0,127,255), (0,255,0), (255,0,0), (95,43,46), (255,0,139)]

def distance(a,b): return(math.sqrt((a[0]-b[0])**2+(a[1]-b[1])**2))
    
# Read a few frames of the video
# Define distance to calculate parabola from (in frames)
interval = 6
for i in range(interval):
    _,_ = cap.read()
    frameNum+= 1

# https://stackoverflow.com/questions/57065080/draw-perpendicular-line-of-fixed-length-at-a-point-of-another-line
def perpendicular(slope, target, dist):
    dy = math.sqrt(3**2/(slope**2+1))*dist
    dx = -slope*dy
    left = target[0] - dx, target[1] - dy
    right = target[0] + dx, target[1] + dy
    return left, right

# Create a list to store past points
history = []
completeHistory = []

### This is the first loop through the video
### In this loop the rainbow-trail data is collected
while True:
    # Read Image
    _, img = cap.read()
    try: _ = img.shape
    except: break

    # Define a small distance
    smallDistance = 0.2

    # Find a parabola that fits
    x_values = x[frameNum-interval:frameNum+interval]
    y_values = y[frameNum-interval:frameNum+interval]
    coefs = poly.polyfit(x_values,y_values,2)

    # Graph x and y values
    for point in zip(x_values, y_values):
        point = tuple(np.array(point,int))
        #cv2.circle(img, point, 1, (255,0,255), 2)

    # Calculate the points on either end of the tangent line
    target = int(x[frameNum]), int(poly.polyval(x[frameNum], coefs))
    leftX, rightX = x[frameNum] - smallDistance, x[frameNum] + smallDistance
    left = leftX, poly.polyval(leftX, coefs)
    right = rightX, poly.polyval(rightX, coefs)
    # Calculate the slope of the tangent line
    slope = (left[1]-right[1])/(left[0]-right[0])
    intercept = target[1] - slope*target[0]
    # Draw line
    leftPoint = target[0]-10, (target[0]-10)*slope + intercept
    rightPoint = target[0]+10, (target[0]+10)*slope + intercept
    leftPoint = tuple(np.array(leftPoint, int))
    rightPoint= tuple(np.array(rightPoint, int))
    #cv2.line(img, leftPoint, rightPoint, (123,234,123), 3)

    # Draw the parabola
    xRange = np.arange(target[0]-9,target[0]+9,1)
    yRange = poly.polyval(xRange, coefs)
    rangePoints = zip(xRange, yRange)
    for ppp in rangePoints:
        center = tuple(np.array(ppp,int))
        #cv2.circle(img, center, 1, 123, 2)
    #print(yRange)

    # List to save data for this frame
    frameHistory = []
    
    # Find Perpendicular points
    for i in range(len(rainbow)):
        color = rainbow[i]
        left, right = perpendicular(slope, target, i-3)
        left, right = tuple(np.array(left, int)), tuple(np.array(right, int))
        point = left
        #cv2.circle(img, left, 1, color, 1)
        frameHistory.append(point)
    history.append(frameHistory)
    completeHistory.append(frameHistory)

    # Show the history too
    a,b,c,d,e,f = zip(*history)
    for pointList, color in zip([a,b,c,d,e,f], rainbow):
        for i in range(len(pointList)-1):
            pass
            cv2.line(img, pointList[i], pointList[i+1], color, 2)

    # Pop the oldest frame off the history if the history is longer than 0.25 seconds
    if len(history)>15:
        history.pop(0)

    out.write(img)
        
    # Show Image
    cv2.imshow('img', img)
    k = cv2.waitKey(1)
    if k == 27: break
    frameNum += 1
cv2.destroyAllWindows()
cap.release()

#Smooth the data
smoothed = []
a,b,c,d,e,f = zip(*completeHistory)
for i in a,b,c,d,e,f:
    x, y = zip(*i)
    x = savgol_filter(x, 27, 3)
    y = savgol_filter(y, 21, 2)
    smoothed.append(list(zip(x,y)))


# Read Source Data
cap = cv2.VideoCapture('/home/stephen/Desktop/ss5_id_412.MP4')
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter('/home/stephen/Desktop/parabola.avi', fourcc, 120.0, (480,848))
frameNum = 0
### This is the second loop through the video
### In this loop the smoothed rinbow-trail points are diplayed
while True:
    # Read Image
    _, img = cap.read()
    if frameNum>21:
        for color, line in zip(rainbow, smoothed):
            pointList = line[frameNum-20:frameNum-5]
            for i in range(len(pointList)-1):
                a,b = tuple(np.array(pointList[i], int)), tuple(np.array(pointList[i+1], int))
                cv2.line(img, a,b, color, 2)
    
    out.write(img)
        
    # Show Image
    cv2.imshow('img', img)
    k = cv2.waitKey(10)
    if k == 27: break
    frameNum += 1
cv2.destroyAllWindows()
cap.release()

https://www.youtube.com/watch?v=_xnA727D8Vc