| 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() |
|
|
