def main():
# Pretty videos:
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
kyleOutputVideo = cv2.VideoWriter('../UntrackedFiles/BallOutputKyle.mp4',
fourcc, 60.0, (1920, 1080))
meganOutputVideo = cv2.VideoWriter('../UntrackedFiles/BallOutputMegan.mp4',
fourcc, 60.0, (1920, 1080))
meganFilename = '../UntrackedFiles/stereoClip5_Megan.mov'
kyleFilename = '../UntrackedFiles/stereoClip5_Kyle.mov'
vrMegan1 = VideoReader(meganFilename)
vrMegan2 = VideoReader(meganFilename)
vrKyle1 = VideoReader(kyleFilename)
vrKyle2 = VideoReader(kyleFilename)
numFrameForward = 5
vrMegan2.setNextFrame(numFrameForward)
vrKyle2.setNextFrame(numFrameForward)
# find court corners:
cfKyle = CourtFinder()
cfMegan = CourtFinder()
numFrames = int(vrMegan1.getNumFrames())
vrMegan1.setNextFrame(int(numFrames / 2))
ret, frame = vrMegan1.readFrame()
cfMegan.FindCourtCorners(frame, 0)
vrKyle1.setNextFrame(int(numFrames / 2))
ret, frame = vrKyle1.readFrame()
cfKyle.FindCourtCorners(frame, 0)
if (not cfMegan.found_corners) or not (cfKyle.found_corners):
print "Couldn't find the court. Exiting."
return
# reset frame index to beginning
vrMegan1.setNextFrame(0)
vrKyle1.setNextFrame(0)
frameNum = 1
meganCam = Camera("megan", cfMegan.corners_sort)
kyleCam = Camera("kyle", cfKyle.corners_sort)
# make a ball finder
bf = BallFinder()
kf = KalmanFilter(vrMegan1.framerate)
csvData = []
while (True):
ret1, kyleFrame1 = vrKyle1.readFrame()
ret2, kyleFrame2, = vrKyle2.readFrame()
ret3, meganFrame1 = vrMegan1.readFrame()
ret4, meganFrame2, = vrMegan2.readFrame()
if not (ret1) or not (ret2) or not (ret3) or not (ret4):
print 'Ending after', frameNum - 1, 'frames.'
break
kyleMask, kyleRays = getBallCandidateRays(bf, kyleCam, kyleFrame1,
kyleFrame2)
meganMask, meganRays = getBallCandidateRays(bf, meganCam, meganFrame1,
meganFrame2)
# check all candidate rays for candidate balls
minDist = 1000000
# TODO set inf
ballPt = []
# all ball points and distances
threshDist = 0.2
# rays must be within .1 m of each other for intersect
ballCandidates = []
candidateCertainty = []
for kyleRay in kyleRays:
for meganRay in meganRays:
pt, dist, _D, _E = IntersectRays(kyleRay, meganRay)
if dist < threshDist and pt[1] < 3.5:
# don't include candidates clearly not valid intersect points
# also don't include candidates that are clearly too high to be the ball
courtBuffer = 2
if pt[0] < Camera.HALF_COURT_X + courtBuffer and pt[
0] > -Camera.HALF_COURT_X - courtBuffer:
if pt[2] < Camera.HALF_COURT_Z + 0.6: # and pt[2] > -Camera.HALF_COURT_Z - 0:
ballCandidates.append(pt)
candidateCertainty.append(dist)
if dist < minDist:
minDist = dist
ballPt = pt
kf.processMeas(ballCandidates, candidateCertainty)
# ========== CSV and VIDEO output ===========
csvTuple = list()
csvTuple.append(frameNum)
if np.linalg.norm(kf.sigma_k, 'fro') < 100: # valid result
# Format the tuple for successful reading
csvTuple.append(1)
posVel = np.reshape(kf.mu_k, (1, 6))[0]
posVel = np.round(posVel, 3)
for val in posVel:
csvTuple.append(val)
for val in kyleCam.ConvertWorldToImagePosition(posVel[0:3]):
csvTuple.append(val)
for val in meganCam.ConvertWorldToImagePosition(posVel[0:3]):
csvTuple.append(val)
# Videos
kyleOutFrame = kyleFrame1.copy()
bf.ballPixelLoc = kyleCam.ConvertWorldToImagePosition(posVel[0:3])
kyleOutFrame = cfKyle.drawCornersOnFrame(kyleOutFrame)
kyleOutFrame = bf.drawBallOnFrame(kyleOutFrame)
kyleOutFrame |= kyleMask
kyleOutputVideo.write(kyleOutFrame)
meganOutFrame = meganFrame1.copy()
meganOutFrame |= meganMask
bf.ballPixelLoc = meganCam.ConvertWorldToImagePosition(posVel[0:3])
meganOutFrame = cfMegan.drawCornersOnFrame(meganOutFrame)
meganOutFrame = bf.drawBallOnFrame(meganOutFrame)
meganOutputVideo.write(meganOutFrame)
else:
# Format the tuple for unsuccessful reading
csvTuple.append(0)
csvData.append(list(csvTuple))
print csvTuple
frameNum += 1
# <END WHILE LOOP>
with open('../UntrackedFiles/ballEst.csv', 'wb') as csvfile:
filewriter = csv.writer(csvfile, delimiter=',')
filewriter.writerows(csvData)
vrKyle1.close()
vrKyle2.close()
vrMegan1.close()
vrMegan2.close()
kyleOutputVideo.release()
meganOutputVideo.release()
