~estraw/visual-odometry

ref: 811ef389135491e9be7387b20737d41f69ea1b10 visual-odometry/src/visual_odometry_demo.py -rw-r--r-- 5.6 KiB
811ef389Evan Straw Add pause feature to demo program 4 months ago
                                                                                
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
166
167
#
# Program to detect and match features between the two frames of a stereo
# camera, and output depth information for those feature points.
#
# Copyright 2021 Evan Straw
# With inspiration from the OpenCV tutorial code:
#   https://docs.opencv.org/3.4/dc/dc3/tutorial_py_matcher.html
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License.  You may obtain a copy
# of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
# License for the specific language governing permissions and limitations under
# the License.
#

import cv2
import numpy as np
import time
import video_input as vi
from visual_odometry import VisualOdometry

# Constants
WIN_NAME = "keypoints"
thresh = 0.5
THRESH_SCALE = 40
intrinsics_file = "calib_images/calib_02-06-21/intrinsics.yml"
extrinsics_file = "calib_images/calib_02-06-21/extrinsics.yml"
# Camera IDs for left and right cameras, respectively.
LCAM = 0
RCAM = 2
DATA_PATH = "/tmp/KITTI"
SEQ_NO = "00"


# No-op callback function for trackbars.
def noop(x):
    pass


def drawInfo(vo, width, fps, rvec, tvec, rvec_real=None, tvec_real=None):
    f_height = 26
    f_height_small = 21
    bar_height = 2 * (f_height + 4)
    c_red = (127, 127, 255)
    c_grn = (127, 255, 127)
    c_blu = (255, 190, 190)
    if rvec_real is None or tvec_real is None:
        bar_height = int(bar_height / 2)

    info_bar = np.zeros((bar_height, out_frame.shape[1], 3), np.uint8)
    cv2.putText(info_bar, ("FPS: %d" % fps),
                (0, int(1.5 * f_height)),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255))

    msgs = vo.formatRT(rvec, tvec).split("\t", 1)
    cv2.putText(info_bar, msgs[0],
                (200, f_height_small),
                cv2.FONT_HERSHEY_PLAIN, 1.5, c_grn)
    cv2.putText(info_bar, msgs[1],
                (850, f_height_small),
                cv2.FONT_HERSHEY_PLAIN, 1.5, c_blu)
    if rvec_real is not None and tvec_real is not None:
        cv2.putText(info_bar, "Err:",
                    (145, int(2 * f_height_small + 5)),
                    cv2.FONT_HERSHEY_PLAIN, 1.5, c_red)
        emsgs = vo.formatError(rvec, tvec,
                               rvec_real, tvec_real).split("\t", 1)
        cv2.putText(info_bar, emsgs[0],
                    (200, int(2 * f_height_small + 5)),
                    cv2.FONT_HERSHEY_PLAIN, 1.5, c_red)
        cv2.putText(info_bar, emsgs[1],
                    (850, int(2 * f_height_small + 5)),
                    cv2.FONT_HERSHEY_PLAIN, 1.5, c_red)

    return info_bar


# Set up windows/trackbars.
cv2.namedWindow(WIN_NAME)
cv2.createTrackbar("Thresh", WIN_NAME, 0, THRESH_SCALE, noop)
cv2.setTrackbarPos("Thresh", WIN_NAME, int(thresh * THRESH_SCALE))

# Attempt to open cameras, exiting on failure.
#   cap: vi.VideoInput = vi.StereoCaptureInput(
#       LCAM, RCAM, intrinsics_file, extrinsics_file)
cap: vi.VideoInput = vi.KITTIDatasetInput(DATA_PATH, SEQ_NO)
if not cap.isOpened():
    print("Error: could not open cameras!")
    exit(1)

vo: VisualOdometry = VisualOdometry(cap.P1, cap.P2)

done = False
first_frame = False
while not done:
    start = time.time()
    # Attempt to grab the next frame from each camera; skip this frame if
    # unsuccessful.
    if not cap.grab():
        continue
    frame: vi.StereoFrame = cap.retrieve()
    # If either frame is empty, skip this frame.
    if frame.left is None or frame.right is None:
        continue

    if not first_frame:
        vo.start(frame)
        first_frame = True
        continue

    thresh = float(cv2.getTrackbarPos("Thresh", WIN_NAME)) / THRESH_SCALE
    rvec, tvec = vo.compute(frame, thresh=thresh)
    rvec_real, tvec_real = None, None
    if rvec is not None:
        speed = np.linalg.norm(tvec) / cap.delta
        speed_mph = speed * 2.236936
        speed_kph = speed * 3.6
        print("Estimated:\t" + vo.formatRT(rvec, tvec))
        if cap.hasGroundTruth:
            rvec_real, tvec_real = cap.poseDelta
            print("Real:\t\t" + vo.formatRT(rvec_real, tvec_real))
            print("Error:\t\t" + vo.formatError(rvec, tvec,
                                                rvec_real, tvec_real))
        print(("Over time delta: %.4f; Approx. speed: %.4f m/s"
               + " (%.4f mph, %.4f km/h)") %
              (cap.delta, speed, speed_mph, speed_kph))
    else:
        print("Not enough information to determine motion")

    elapsed = time.time() - start
    fps = int(round(1 / elapsed))
    out_frame = vo.drawTimeMatches()
    info_bar = drawInfo(vo, out_frame.shape[1], fps, rvec, tvec,
                        rvec_real, tvec_real)
    out_frame = cv2.vconcat([out_frame, info_bar])

    # Show the output frame.
    cv2.imshow(WIN_NAME, out_frame)
    #   cv2.imshow("KP", out_kp)

    # Wait up to 10ms for a keypress; if 'q' is pressed, exit the loop.
    k = cv2.waitKey(10)
    if k == ord('q'):
        done = True
    elif k == ord(' '):
        pause = True
        p_frame = cv2.putText(out_frame, "PAUSED",
                              (4, 44), cv2.FONT_HERSHEY_SIMPLEX, 2,
                              (100, 100, 255), 2)
        while pause:
            cv2.imshow(WIN_NAME, p_frame)
            k = cv2.waitKey(100)
            if k == ord(' '):
                pause = False
            elif k == ord('q'):
                pause = False
                done = True

# Clean up; release cameras and close the window.
del cap
cv2.destroyWindow(WIN_NAME)