HedgehogTrack/hedgehogbox.py

62 lines
No EOL
1.7 KiB
Python

from argparse import ArgumentParser
from pathlib import Path
import cv2
import imutils
import time
def main():
parser = ArgumentParser()
parser.add_argument("--output", type=Path, required=False, default="output.avi")
parser.add_argument("video_file", type=Path)
args = parser.parse_args()
cap = cv2.VideoCapture(args.video_file)
output = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*'MPEG'), 30, (720, 1792))
static_frame = None
while True:
ret, frame = cap.read()
if ret:
cropped_frame = frame[55:720, 0:540]
grey_frame = cv2.cvtColor(cropped_frame, cv2.COLOR_BGR2GRAY)
if static_frame is None:
static_frame = grey_frame
continue
_, thresh = cv2.threshold(grey_frame, 120, 255, cv2.THRESH_BINARY)
dilated_image = cv2.dilate(thresh, None, iterations=2)
cnts = cv2.findContours(dilated_image.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cnts = imutils.grab_contours(cnts)
for c in cnts:
if cv2.contourArea(c) < 700:
continue
# get the bounding box coordinates
(x, y, w, h) = cv2.boundingRect(c)
cv2.rectangle(cropped_frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
output.write(frame)
cv2.imshow("output", frame)
# exit if 's' key is pressed
if cv2.waitKey(1) & 0xFF == ord('s'):
break
time.sleep(0.03)
else:
print("Could not read frame!")
break
cv2.destroyAllWindows()
output.release()
cap.release()
if __name__ == "__main__":
main()