86 lines
2.2 KiB
Python
86 lines
2.2 KiB
Python
import numpy as np
|
|
import cv2
|
|
import math
|
|
import sys
|
|
|
|
# Init
|
|
target = sys.argv[1]
|
|
cap = cv2.VideoCapture("data/{}.mp4".format(target))
|
|
cap.set(cv2.CAP_PROP_FPS,240)
|
|
out = open("{}.tsv".format(target),"w+")
|
|
out.write("Time\tX\tY\tCount\tVelocity\n")
|
|
|
|
# Image Processing
|
|
scale = 16.0
|
|
diff_thresh = 50
|
|
blur_kernel = (7,7)
|
|
|
|
# Background
|
|
bg_alpha = 0.99
|
|
bg = None
|
|
bg_float = None
|
|
|
|
# Position
|
|
pos_alpha = 0.99
|
|
posx = 0.0
|
|
posy = 0.0
|
|
|
|
# Count
|
|
running_count = 0.0
|
|
|
|
# Main Loop
|
|
frame_number = 0
|
|
|
|
while True:
|
|
# Grab Frame
|
|
ret, src = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
# Ensmallen and grayen
|
|
blur = cv2.GaussianBlur(src, blur_kernel,0)
|
|
smol = cv2.resize(blur, (int(src.shape[1]/scale), int(src.shape[0]/scale)))
|
|
gray = cv2.cvtColor(smol, cv2.COLOR_BGR2GRAY)
|
|
if bg is None:
|
|
bg = gray
|
|
bg_float = np.zeros(gray.shape,dtype=np.float32)
|
|
|
|
# Get Difference
|
|
diff = cv2.absdiff(gray, bg)
|
|
retval, thresh = cv2.threshold(diff, diff_thresh, 255, cv2.THRESH_TOZERO)
|
|
|
|
# Normalize Count
|
|
count = np.sum(thresh)/(smol.shape[0]*smol.shape[1]*255)
|
|
|
|
# Process Count
|
|
if count > 0.0:
|
|
moments = cv2.moments(thresh)
|
|
cX = (moments["m10"] / moments["m00"])/smol.shape[1]
|
|
cY = (moments["m01"] / moments["m00"])/smol.shape[0]
|
|
prevx = posx
|
|
prevy = posy
|
|
posx = posx * pos_alpha + cX * (1.0-pos_alpha)
|
|
posy = posy * pos_alpha + cY * (1.0-pos_alpha)
|
|
velocity = math.sqrt((prevx-posx) * (prevx-posx) + (prevy - posy) * (prevy - posy))
|
|
running_count = running_count * pos_alpha + count * (1.0 - pos_alpha)
|
|
out.write("{}\t{}\t{}\t{}\t{}\n".format(frame_number/30.0,posx,posy,running_count,velocity))
|
|
|
|
# Calculate new BG
|
|
bg_float = bg_float * bg_alpha + gray * (1.0 - bg_alpha)
|
|
bg = bg_float.astype(np.uint8)
|
|
|
|
# Draw
|
|
#cv2.circle(smol, (int(posx*smol.shape[1]), int(posy*smol.shape[0])), 3, (255,0,0),6)
|
|
|
|
# Show Results
|
|
# cv2.imshow("thresh", cv2.resize(thresh, (int(src.shape[1]/2.0), int(src.shape[0]/2.0))))
|
|
# cv2.imshow("smol", cv2.resize(smol, (int(src.shape[1]/2.0), int(src.shape[0]/2.0))))
|
|
# key = cv2.waitKey(1)
|
|
# if key == 27:
|
|
# break
|
|
|
|
frame_number += 1
|
|
print(frame_number)
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows() |