[ Computer Vision ] Real-Time Depth Estimation CODE

Pytorch에서는 Depth Estimation과 관련된 Model을 제공해줍니다..!

https://pytorch.org/hub/intelisl_midas_v2/

 

우선 필요한 라이브러리를 전부 불러옵니다.

import cv2
import torch
import time
import numpy as np

model_type = "MiDaS_small"
midas = torch.hub.load("intel-isl/MiDaS", model_type)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
midas.to(device)
midas.eval()

midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")

transform = midas_transforms.small_transform

이후 아래의 코드를 실행하면 됩니다..

midas architecture를 통하여, 예측을 진행한 후, 이미지를 interpolate합니다. 이후, Color Map을 적용한 후 출력을 합니다.
(추후 조금 더 깊게 살펴보도록 하겠습니다 :) )

import cv2
import torch
import time
import numpy as np

model_type = "MiDaS_small"
midas = torch.hub.load("intel-isl/MiDaS", model_type)

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
midas.to(device)
midas.eval()
midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")

transform = midas_transforms.small_transform

cap = cv2.VideoCapture(1)

while cap.isOpened():
    success, img = cap.read()
    start = time.time()
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    input_batch = transform(img).to(device)

    with torch.no_grad():
        prediction = midas(input_batch)

        prediction = torch.nn.functional.interpolate(
            prediction.unsqueeze(1),
            size = img.shape[:2],
            mode = "bicubic",
            align_corners= False,
        ).squeeze()
    
    depth_map = prediction.cpu().numpy()
    depth_map = cv2.normalize(depth_map, None, 0, 1, norm_type = cv2.NORM_MINMAX)

    end = time.time()
    totalTime = end - start
    fps = 1 / totalTime

    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    depth_map = (depth_map*255).astype(np.uint8)
    depth_map = cv2.applyColorMap(depth_map, cv2.COLORMAP_MAGMA)        # cap.set(cv2.CAP_PROP_FRAME_WIDTH,640)

    cv2.imshow("mom", img)
    cv2.imshow("image", depth_map)

    if cv2.waitKey(10) == ord("x"):
        break

cap.release()
cv2.destroyAllWindows()