-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdriving_stereo_test.py
51 lines (39 loc) · 1.62 KB
/
driving_stereo_test.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
import cv2
import numpy as np
import glob
from mobilestereonet import MobileStereoNet, CameraConfig
from mobilestereonet.utils import draw_depth
# out = cv2.VideoWriter('outpy.avi',cv2.VideoWriter_fourcc('M','J','P','G'), 20, (1762*2,800))
# Get image list
left_images = glob.glob('DrivingStereo images/left/*.png')
left_images.sort()
right_images = glob.glob('DrivingStereo images/right/*.png')
right_images.sort()
depth_images = glob.glob('DrivingStereo images/depth/*.png')
depth_images.sort()
model_path = "models/model_float32.tflite"
input_width = 320
camera_config = CameraConfig(0.546, 2000/1920*input_width) # rough estimate from the original calibration
max_distance = 30
# Initialize model
mobile_depth_estimator = MobileStereoNet(model_path, camera_config)
cv2.namedWindow("Estimated depth", cv2.WINDOW_NORMAL)
for left_path, right_path, depth_path in zip(left_images[:], right_images[:], depth_images[:]):
# Read frame from the video
left_img = cv2.imread(left_path)
right_img = cv2.imread(right_path)
depth_img = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED).astype(np.float32)/256
# Estimate the depthq
disparity_map = mobile_depth_estimator(left_img, right_img)
depth_map = mobile_depth_estimator.get_depth()
color_depth = draw_depth(depth_map, max_distance)
color_real_depth = draw_depth(depth_img, max_distance)
color_depth = cv2.resize(color_depth, (left_img.shape[1],left_img.shape[0]))
combined_image = np.hstack((left_img, color_depth))
# out.write(combined_image)
cv2.imshow("Estimated depth", combined_image)
# Press key q to stop
if cv2.waitKey(1) == ord('q'):
break
# out.release()
cv2.destroyAllWindows()