File size: 4,627 Bytes
56bd2b5 |
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 |
import numpy as np
import torchvision.transforms as transforms
import torch
from PIL import Image
from depth.metric_depth.zoedepth.models.builder import build_model
from depth.metric_depth.zoedepth.utils.config import get_config
def depth_of_images(image, model):
"""
This function takes in a list of images and returns the depth of the images"""
# Born out of Issue 36.
# Allows the user to set up own test files to infer on (Create a folder my_test and add subfolder input and output in the metric_depth directory before running this script.)
# Make sure you have the necessary libraries
# Code by @1ssb
# Global settings
DATASET = 'nyu' # Lets not pick a fight with the model's dataloader
color_image = Image.fromarray(image).convert('RGB')
original_width, original_height = color_image.size
image_tensor = transforms.ToTensor()(color_image).unsqueeze(0).to('cuda' if torch.cuda.is_available() else 'cpu')
# input as bx3xhxw (unnormalized image)
pred_o = model(image_tensor, dataset=DATASET)
if isinstance(pred_o, dict):
pred = pred_o.get('metric_depth', pred_o.get('out'))
features = pred_o.get('depth_features', None)
elif isinstance(pred_o, (list, tuple)):
pred = pred[-1]
pred = pred.squeeze().detach().cpu().numpy()
# Resize color image and depth to final size
resized_pred = Image.fromarray(pred).resize((original_width, original_height), Image.NEAREST)
# resized_pred is the image shaped to the original image size, depth is in meters
return np.array(resized_pred)
def setup_depth_model(model_name, pretrained_resource):
DATASET = 'nyu' # Lets not pick a fight with the model's dataloader
config = get_config(model_name, "eval", DATASET)
config.pretrained_resource = pretrained_resource
model = build_model(config).to('cuda' if torch.cuda.is_available() else 'cpu')
model.eval()
return model
def init_dataset():
''' dataloader stuff.
I'm not sure what the difference between the omni3d dataset and load omni3D json functions are. this is a 3rd alternative to this. The train script calls something similar to this.'''
cfg, filter_settings = get_config_and_filter_settings()
dataset_names = ['SUNRGBD_train','SUNRGBD_val','SUNRGBD_test']
dataset_paths_to_json = ['datasets/Omni3D/'+dataset_name+'.json' for dataset_name in dataset_names]
# for dataset_name in dataset_names:
# simple_register(dataset_name, filter_settings, filter_empty=True)
# Get Image and annotations
datasets = data.Omni3D(dataset_paths_to_json, filter_settings=filter_settings)
data.register_and_store_model_metadata(datasets, cfg.OUTPUT_DIR, filter_settings)
thing_classes = MetadataCatalog.get('omni3d_model').thing_classes
dataset_id_to_contiguous_id = MetadataCatalog.get('omni3d_model').thing_dataset_id_to_contiguous_id
infos = datasets.dataset['info']
dataset_id_to_unknown_cats = {}
possible_categories = set(i for i in range(cfg.MODEL.ROI_HEADS.NUM_CLASSES + 1))
dataset_id_to_src = {}
for info in infos:
dataset_id = info['id']
known_category_training_ids = set()
if not dataset_id in dataset_id_to_src:
dataset_id_to_src[dataset_id] = info['source']
for id in info['known_category_ids']:
if id in dataset_id_to_contiguous_id:
known_category_training_ids.add(dataset_id_to_contiguous_id[id])
# determine and store the unknown categories.
unknown_categories = possible_categories - known_category_training_ids
dataset_id_to_unknown_cats[dataset_id] = unknown_categories
return datasets
if __name__ == '__main__':
import os
from detectron2.data.catalog import MetadataCatalog
from cubercnn import data
from priors import get_config_and_filter_settings
import torch.nn.functional as F
from tqdm import tqdm
# datasets = init_dataset()
os.makedirs('datasets/depth_maps', exist_ok=True)
depth_model = 'zoedepth'
pretrained_resource = 'local::depth/checkpoints/depth_anything_metric_depth_indoor.pt'
model = setup_depth_model(depth_model, pretrained_resource)
for img_id in tqdm(os.listdir('datasets/coco_examples')):
file_path = 'coco_examples/'+img_id
# for img_id, img_info in track(datasets.imgs.items()):
# file_path = img_info['file_path']
img = np.array(Image.open('datasets/'+file_path))
depth = depth_of_images(img, model)
np.savez_compressed(f'datasets/depth_maps/{img_id}.npz', depth=depth) |