File size: 2,321 Bytes
f53b39e |
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 |
from PIL import Image
import depth_pro
from tqdm import tqdm
import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '1'
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import argparse
from transformers import pipeline
def find_images(directory):
image_paths = []
for root, dirs, files in os.walk(directory):
for file in files:
if file.lower().endswith(('rgb.jpg', 'rgb.png')):
image_paths.append(os.path.join(root, file))
return image_paths
# Load model and preprocessing transform
parser = argparse.ArgumentParser('ml infer', add_help=False)
parser.add_argument('--a', default=0,
type=int)
parser.add_argument('--b', default=1500,
type=int)
parser.add_argument('--dataset_name', default=None,
type=str)
args = parser.parse_args()
# print(args.a)
model, transform = depth_pro.create_model_and_transforms(device='cuda')
model.eval()
pipe = pipeline(task="depth-estimation", model="depth-anything/Depth-Anything-V2-Large-hf",device='cuda')
if args.dataset_name == "Tartanair":
dir = '../../data/Tartanair_proc/'
elif args.dataset_name == "spring":
dir = '../../data/spring_proc/train/'
elif args.dataset_name == "SceneFlow":
dir = '../../data/SceneFlow/'
elif args.dataset_name == "Vkitti":
dir = '../../data/vkitti_2.0.3_proc/'
elif args.dataset_name == "PointOdyssey":
dir = '../../data/PointOdyssey_proc/'
image_paths = find_images(dir)
for image_path in tqdm(sorted(image_paths)[int(args.a):int(args.b)]):
# depthanything v2
image = Image.open(image_path)
depth = pipe(image)["predicted_depth"].numpy()
#depth = prediction["depth"].cpu() # Depth in [m].
metadata = np.load(image_path.replace('_rgb.jpg', '_metadata.npz'))
intrinsics = np.float32(metadata['camera_intrinsics'])
focallength_px = intrinsics[0][0]
np.savez_compressed(image_path[:-4]+'_pred_depth_depthanything', depth=depth,focallength_px=focallength_px)
# depthpro
image, _, f_px = depth_pro.load_rgb(image_path)
image = transform(image)
# Run inference.
prediction = model.infer(image, f_px=f_px)
depth = prediction["depth"].cpu() # Depth in [m].
np.savez_compressed(image_path[:-4]+'_pred_depth_depthpro', depth=depth, focallength_px=prediction["focallength_px"].cpu()) |