|
import os |
|
import torch |
|
import argparse |
|
import numpy as np |
|
import cv2 |
|
from pathlib import Path |
|
from tqdm.auto import tqdm |
|
|
|
|
|
def parse_arguments(): |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument( |
|
"--img_path", type=str, required=True, help="path to the file or the folder" |
|
) |
|
parser.add_argument( |
|
"--save_folder", type=str, default="./depth/", help="path to the folder" |
|
) |
|
parser.add_argument( |
|
"--midas_model", type=str, default="DPT_Large", help="Midas model name" |
|
) |
|
parser.add_argument("--use_cuda", action="store_true") |
|
parser.add_argument("--baseline", type=float, default=0.54) |
|
parser.add_argument("--focal", type=float, default=721.09) |
|
parser.add_argument("--img_scale", type=float, default=1) |
|
return parser.parse_args() |
|
|
|
|
|
def get_depth_estimation_model(model_name: str, device="cpu"): |
|
assert model_name in ["DPT_Large", "DPT_Hybrid", "MiDaS_small"] |
|
|
|
midas = torch.hub.load("intel-isl/MiDaS", model_name) |
|
midas.eval() |
|
midas.to(device) |
|
|
|
midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms") |
|
if model_name in ["DPT_Large", "DPT_Hybrid"]: |
|
transform = midas_transforms.dpt_transform |
|
else: |
|
transform = midas_transforms.small_transform |
|
return midas, transform |
|
|
|
|
|
def getDisparityMap(model, transform, img_path): |
|
img = cv2.imread(str(img_path)) |
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
|
input_batch = transform(img) |
|
|
|
with torch.no_grad(): |
|
prediction = model(input_batch.cuda()) |
|
|
|
prediction = torch.nn.functional.interpolate( |
|
prediction.unsqueeze(1), |
|
size=img.shape[:2], |
|
mode="bicubic", |
|
align_corners=False, |
|
).squeeze() |
|
return prediction.cpu().numpy() |
|
|
|
|
|
def get_depth_map( |
|
midas, midas_transform, imgp, baseline=0.54, focal=721.09, img_scale=1 |
|
): |
|
disp = getDisparityMap(midas, midas_transform, imgp) |
|
disp[disp < 0] = 0 |
|
disp = disp + 1e-3 |
|
depth = baseline * focal / (disp * img_scale) |
|
|
|
return depth |
|
|
|
|
|
def get_depth_map_new(midas, midas_transform, imgp): |
|
depth = getDisparityMap(midas, midas_transform, imgp) |
|
depth[depth < 0] = 0 |
|
depth = depth + 1e-3 |
|
depth = depth |
|
return depth.max() - depth |
|
|
|
|
|
def main(): |
|
args = parse_arguments() |
|
|
|
device = torch.device("cpu") |
|
if args.use_cuda: |
|
device = torch.device("cuda") |
|
|
|
|
|
baseline = args.baseline |
|
focal = args.focal |
|
img_scale = args.img_scale |
|
|
|
imgP = Path(args.img_path) |
|
save_folder = Path(args.save_folder) |
|
if not save_folder.exists(): |
|
os.makedirs(str(save_folder)) |
|
|
|
midas, midas_transform = get_depth_estimation_model( |
|
model_name=args.midas_model, device=device |
|
) |
|
|
|
if imgP.is_dir(): |
|
image_files = sorted(Path(imgP).glob("*")) |
|
for imgp in tqdm(image_files): |
|
depth = get_depth_map( |
|
midas, midas_transform, imgp, baseline, focal, img_scale |
|
) |
|
np.save(save_folder / imgp.stem, depth) |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
