-
Notifications
You must be signed in to change notification settings - Fork 15
/
frame_capture.py
165 lines (133 loc) · 6.85 KB
/
frame_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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
#!/usr/bin/env python3
"""Save frames 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 HQ frames (default: 1920x1080 px) to .jpg at the
specified capture frequency (default: 1 s)
-> stop recording early if free disk space drops below threshold
- optional arguments:
'-min' set recording time in minutes (default: 2 [min])
-> e.g. '-min 5' for 5 min recording time
'-4k' set camera resolution to 4K (3840x2160 px) (default: 1080p)
'-lq' additionally save downscaled LQ frames (e.g. 320x320 px)
'-af' set auto focus range in cm (min distance, max distance)
-> e.g. '-af 14 20' to restrict auto focus range to 14-20 cm
'-zip' store all captured data in an uncompressed .zip file for each day
and delete original directory
-> increases file transfer speed from microSD to computer
but also on-device processing time and power consumption
based on open source scripts available at https://github.com/luxonis
"""
import argparse
import logging
import time
from datetime import datetime
from pathlib import Path
import cv2
import depthai as dai
import psutil
from utils.general import zip_data
from utils.oak_cam import set_focus_range
# Define optional arguments
parser = argparse.ArgumentParser()
parser.add_argument("-min", "--min_rec_time", type=int, choices=range(1, 721), default=2,
help="Set recording time in minutes (default: 2 [min]).", metavar="1-720")
parser.add_argument("-4k", "--four_k_resolution", action="store_true",
help="Set camera resolution to 4K (3840x2160 px) (default: 1080p).")
parser.add_argument("-lq", "--save_lq_frames", action="store_true",
help="Additionally save downscaled LQ frames (320x320 px).")
parser.add_argument("-af", "--af_range", nargs=2, type=int,
help="Set auto focus range in cm (min distance, max distance).", metavar=("CM_MIN", "CM_MAX"))
parser.add_argument("-zip", "--zip_data", action="store_true",
help="Store data in an uncompressed .zip file for each day and delete original directory.")
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 capture frequency (default: 1 second)
# -> wait for specified amount of seconds between saving HQ frames
# 'CAPTURE_FREQ = 0.8' (0.2 for 4K) saves ~58 frames per minute to .jpg
CAPTURE_FREQ = 0.8 if not args.four_k_resolution else 0.2
# Set recording time (default: 2 minutes)
REC_TIME = args.min_rec_time * 60
# Set logging level and format
logging.basicConfig(level=logging.INFO, format="%(message)s")
# Create directory per day (date) and recording interval (date_time) to save HQ frames (+ LQ frames)
rec_start_str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
save_path = Path.home() / "insect-detect" / "frames" / rec_start_str[:10] / rec_start_str
save_path.mkdir(parents=True, exist_ok=True)
if args.save_lq_frames:
(save_path / "LQ_frames").mkdir(parents=True, exist_ok=True)
# Create depthai pipeline
pipeline = dai.Pipeline()
# Create and configure color camera node and define output(s)
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.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
cam_rgb.setFps(25) # frames per second available for auto focus/exposure
if args.save_lq_frames:
cam_rgb.setPreviewSize(320, 320) # downscale frames -> LQ frames
cam_rgb.setPreviewKeepAspectRatio(False) # stretch frames (16:9) to square (1:1)
xout_rgb = pipeline.create(dai.node.XLinkOut)
xout_rgb.setStreamName("frame")
cam_rgb.video.link(xout_rgb.input) # HQ frames
if args.save_lq_frames:
xout_lq = pipeline.create(dai.node.XLinkOut)
xout_lq.setStreamName("frame_lq")
cam_rgb.preview.link(xout_lq.input) # LQ frames
if args.af_range:
# Create XLinkIn node to send control commands to color camera node
xin_ctrl = pipeline.create(dai.node.XLinkIn)
xin_ctrl.setStreamName("control")
xin_ctrl.out.link(cam_rgb.inputControl)
# 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(s) to get the frames from the output(s) defined above
q_frame = device.getOutputQueue(name="frame", maxSize=4, blocking=False)
if args.save_lq_frames:
q_frame_lq = device.getOutputQueue(name="frame_lq", maxSize=4, blocking=False)
if args.af_range:
# Create input queue to send control commands to OAK camera
q_ctrl = device.getInputQueue(name="control", maxSize=16, blocking=False)
# Set auto focus range to specified cm values
af_ctrl = set_focus_range(args.af_range[0], args.af_range[1])
q_ctrl.send(af_ctrl)
# 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 HQ (+ LQ) frames and save to .jpg
if q_frame.has():
frame_hq = q_frame.get().getCvFrame()
timestamp_frame = datetime.now().strftime("%Y-%m-%d_%H-%M-%S-%f")
path_hq = f"{save_path}/{timestamp_frame}.jpg"
cv2.imwrite(path_hq, frame_hq)
if args.save_lq_frames:
if q_frame_lq.has():
frame_lq = q_frame_lq.get().getCvFrame()
path_lq = f"{save_path}/LQ_frames/{timestamp_frame}_LQ.jpg"
cv2.imwrite(path_lq, frame_lq)
# Wait for specified amount of seconds (default: 0.8 for 1080p; 0.2 for 4K)
time.sleep(CAPTURE_FREQ)
# Print number and directory of saved frames
num_frames_hq = len(list(save_path.glob("*.jpg")))
if not args.save_lq_frames:
logging.info("Saved %s HQ frames to %s\n", num_frames_hq, save_path)
else:
num_frames_lq = len(list((save_path / "LQ_frames").glob("*.jpg")))
logging.info("Saved %s HQ and %s LQ frames to %s\n", num_frames_hq, num_frames_lq, save_path)
if args.zip_data:
# Store frames in uncompressed .zip file and delete original folder
zip_data(save_path)
logging.info("Stored all captured images in %s.zip\n", save_path.parent)