MiDaSDepthEstimation/main.py at master · NNEngine/MiDaSDepthEstimation · GitHub

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
import torch
import cv2
import numpy as np
from torchvision.transforms import Compose

from MiDaS.midas.dpt_depth import DPTDepthModel
from MiDaS.midas.transforms import Resize, NormalizeImage, PrepareForNet

from distance_estimation import estimate_distance_range

# ----- CONFIGURATION -----
model_type = "dpt_hybrid"  # or "dpt_large"
model_path = f"MiDaS/weights/{model_type}_384.pt"
use_camera_id = 0  # Use 0 for default webcam

# ----- DEVICE SETUP -----
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ----- LOAD MODEL -----
print("Loading model...")
model = DPTDepthModel(
    path=model_path,
    backbone="vitb_rn50_384" if model_type == "dpt_hybrid" else "vitl16_384",
    non_negative=True
)
model.eval()
model.to(device)

# ----- TRANSFORM -----
transform = Compose([
    Resize(
        384, 384,
        resize_target=None,
        keep_aspect_ratio=True,
        ensure_multiple_of=32,
        resize_method="minimal",
        image_interpolation_method=cv2.INTER_CUBIC,
    ),
    NormalizeImage(mean=[0.485, 0.456, 0.406],
                   std=[0.229, 0.224, 0.225]),
    PrepareForNet(),
])


# ----- START CAMERA -----
cap = cv2.VideoCapture(use_camera_id)
if not cap.isOpened():
    print("Error: Cannot open webcam")
    exit()

print("Press ESC to exit")

window_name = "MiDaS Depth Estimation (Left: Original, Right: Depth)"

cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.setWindowProperty("MiDaS Depth Estimation (Left: Original, Right: Depth)", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)

while True:
    ret, frame = cap.read()
    if not ret:
        print("Failed to grab frame")
        break

    # Convert BGR to RGB
    img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # Transform input
    sample = transform({"image": img_rgb})
    image = sample["image"]
    image = torch.from_numpy(image).unsqueeze(0).to(device)

    # Predict depth
    with torch.no_grad():
        prediction = model(image)
        prediction = torch.nn.functional.interpolate(
            prediction.unsqueeze(1),
            size=img_rgb.shape[:2],
            mode="bicubic",
            align_corners=False,
        ).squeeze().cpu().numpy()

    # Normalize depth for display
    depth_min = prediction.min()
    depth_max = prediction.max()
    if depth_max - depth_min > 0:
        prediction = (prediction - depth_min) / (depth_max - depth_min)
    else:
        prediction = 0

    depth_visual = (255 * prediction).astype(np.uint8)
    depth_colormap = cv2.applyColorMap(depth_visual, cv2.COLORMAP_INFERNO)

    # Stack both original and depth maps
    combined = cv2.hconcat([frame, depth_colormap])

    # Get depth at center
    h, w = prediction.shape
    center_depth = prediction[h//2, w//2]
    distance_range = estimate_distance_range(center_depth)

    text1 = f"Depth value: {center_depth:.3f}"
    text2 = f"Estimated Distance: {distance_range}"
    cv2.putText(combined, text1, (50, 50),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
    cv2.putText(combined, text2, (50, 90),
            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)


    cv2.imshow(window_name, combined)

    if cv2.waitKey(1) == 27:  # ESC key
        break

cap.release()
cv2.destroyAllWindows()