-
Notifications
You must be signed in to change notification settings - Fork 15
/
video_capture.py
145 lines (114 loc) · 6 KB
/
video_capture.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
#!/usr/bin/env python3
"""Save video from OAK camera.
Source: https://github.com/maxsitt/insect-detect
License: GNU GPLv3 (https://choosealicense.com/licenses/gpl-3.0/)
Author: Maximilian Sittinger (https://github.com/maxsitt)
Docs: https://maxsitt.github.io/insect-detect-docs/
- save encoded HQ frames (1080p or 4K resolution) with H.265 (HEVC) compression to .mp4 video file
- optional arguments:
'-min' set recording time in minutes (default: 2 [min])
-> e.g. '-min 5' for 5 min recording time
'-4k' record video in 4K resolution (3840x2160 px) (default: 1080p)
'-fps' set camera frame rate (default: 25 fps)
-> e.g. '-fps 20' for 20 fps (less fps = smaller video file size)
'-af' set auto focus range in cm (min - max distance to camera)
-> e.g. '-af 14 20' to restrict auto focus range to 14-20 cm
'-mf' set manual focus position in cm (distance to camera)
-> e.g. '-mf 14' to set manual focus position to 14 cm
based on open source scripts available at https://github.com/luxonis
"""
import argparse
import logging
import time
from datetime import datetime
from fractions import Fraction
from pathlib import Path
import av
import depthai as dai
import psutil
from utils.oak_cam import convert_cm_lens_position
# Define optional arguments
parser = argparse.ArgumentParser()
group = parser.add_mutually_exclusive_group()
parser.add_argument("-min", "--min_rec_time", type=int, choices=range(1, 61), default=2,
help="Set recording time in minutes (default: 2 [min]).", metavar="1-60")
parser.add_argument("-4k", "--four_k_resolution", action="store_true",
help="Set camera resolution to 4K (3840x2160 px) (default: 1080p).")
parser.add_argument("-fps", "--frames_per_second", type=int, choices=range(1, 31), default=25,
help="Set camera frame rate (default: 25 fps).", metavar="1-30")
group.add_argument("-af", "--af_range", nargs=2, type=int,
help="Set auto focus range in cm (min - max distance to camera).", metavar=("CM_MIN", "CM_MAX"))
group.add_argument("-mf", "--manual_focus", type=int,
help="Set manual focus position in cm (distance to camera).", metavar="CM")
args = parser.parse_args()
# Set threshold value required to start and continue a recording
MIN_DISKSPACE = 100 # minimum free disk space (MB) (default: 100 MB)
# Set recording time (default: 2 minutes)
REC_TIME = args.min_rec_time * 60
# Set video resolution
RES = "1080p" if not args.four_k_resolution else "4K"
# Set frame rate (default: 25 fps)
FPS = args.frames_per_second
# Set logging level and format
logging.basicConfig(level=logging.INFO, format="%(message)s")
# Create directory per day (date) to save video
rec_start_str = datetime.now().strftime("%Y-%m-%d")
save_path = Path.home() / "insect-detect" / "videos" / rec_start_str
save_path.mkdir(parents=True, exist_ok=True)
# Create depthai pipeline
pipeline = dai.Pipeline()
# Create and configure color camera node
cam_rgb = pipeline.create(dai.node.ColorCamera)
#cam_rgb.setImageOrientation(dai.CameraImageOrientation.ROTATE_180_DEG) # rotate image 180°
cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_4_K)
if not args.four_k_resolution:
cam_rgb.setIspScale(1, 2) # downscale 4K to 1080p resolution -> HQ frames
cam_rgb.setInterleaved(False) # planar layout
cam_rgb.setFps(FPS) # frames per second available for auto focus/exposure
if args.af_range:
# Convert cm to lens position values and set auto focus range
lens_pos_min, lens_pos_max = convert_cm_lens_position((args.af_range[1], args.af_range[0]))
cam_rgb.initialControl.setAutoFocusLensRange(lens_pos_min, lens_pos_max)
if args.manual_focus:
# Convert cm to lens position value and set manual focus position
lens_pos = convert_cm_lens_position(args.manual_focus)
cam_rgb.initialControl.setManualFocus(lens_pos)
# Create and configure video encoder node and define input + output
video_enc = pipeline.create(dai.node.VideoEncoder)
video_enc.setDefaultProfilePreset(FPS, dai.VideoEncoderProperties.Profile.H265_MAIN)
cam_rgb.video.link(video_enc.input)
xout_vid = pipeline.create(dai.node.XLinkOut)
xout_vid.setStreamName("video")
video_enc.bitstream.link(xout_vid.input) # encoded HQ frames
# Connect to OAK device and start pipeline in USB2 mode
with dai.Device(pipeline, maxUsbSpeed=dai.UsbSpeed.HIGH) as device:
logging.info("Recording time: %s min\n", int(REC_TIME / 60))
# Get free disk space (MB)
disk_free = round(psutil.disk_usage("/").free / 1048576)
# Create output queue to get the encoded frames from the output defined above
q_video = device.getOutputQueue(name="video", maxSize=30, blocking=True)
# Create .mp4 container with H.265 (HEVC) compression
timestamp_video = datetime.now().strftime("%Y-%m-%d_%H-%M-%S-%f")
with av.open(f"{save_path}/{timestamp_video}_{FPS}fps_{RES}_video.mp4", "w") as container:
stream = container.add_stream("hevc", rate=FPS, options={"x265-params": "log_level=none"})
stream.width, stream.height = cam_rgb.getVideoSize()
stream.time_base = Fraction(1, 1000 * 1000)
# Set start time of recording
start_time = time.monotonic()
# Record until recording time is finished
# Stop recording early if free disk space drops below threshold
while time.monotonic() < start_time + REC_TIME and disk_free > MIN_DISKSPACE:
# Update free disk space (MB)
disk_free = round(psutil.disk_usage("/").free / 1048576)
# Get encoded video frames and save to packet
if q_video.has():
enc_video = q_video.get().getData()
packet = av.Packet(enc_video)
packet_timestamp = int((time.monotonic() - start_time) * 1000 * 1000)
packet.dts = packet_timestamp
packet.pts = packet_timestamp
# Mux packet into .mp4 container
container.mux_one(packet)
# Print duration, resolution, fps and directory of saved video + free disk space
logging.info("Saved %s min %s video with %s fps to %s\n", int(REC_TIME / 60), RES, FPS, save_path)
logging.info("Free disk space left: %s MB\n", disk_free)