Spaces:
Runtime error
Runtime error
File size: 6,640 Bytes
5d87992 |
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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 |
# Copyright Niantic 2021. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the ManyDepth licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.
import os
import time
import json
import argparse
import numpy as np
from PIL import Image
import matplotlib as mpl
import matplotlib.cm as cm
import torch
from torch import Tensor
import torchvision
from torchvision import transforms
import torch.nn.functional as F
from src.networks import *
from utils import transformation_from_parameters, disp_to_depth, line
def load_and_preprocess_image(image, resize_width, resize_height):
image_ori = image.convert('RGB')
W, H = image_ori.size
W_resized = W - W % 32
H_resized = H - H % 32
img_ori_npy = np.array(image_ori)[0:H_resized, 0:W_resized]
image = image_ori.resize((resize_width, resize_height), Image.Resampling.LANCZOS)
image = transforms.ToTensor()(image)
image_ori = transforms.ToTensor()(img_ori_npy).unsqueeze(0)
image = line(image).unsqueeze(0)
if torch.cuda.is_available():
return image_ori.cuda(), image.cuda(), (H, W)
return image_ori, image, (H, W)
def load_and_preprocess_intrinsics(intrinsics_path, resize_width, resize_height):
K = np.eye(4)
with open(intrinsics_path, 'r') as f:
K[:3, :3] = np.array(json.load(f))
# Convert normalised intrinsics to 1/4 size unnormalised intrinsics.
# (The cost volume construction expects the intrinsics corresponding to 1/4 size images)
K[0, :] *= resize_width // 4
K[1, :] *= resize_height // 4
invK = torch.Tensor(np.linalg.pinv(K)).unsqueeze(0)
K = torch.Tensor(K).unsqueeze(0)
if torch.cuda.is_available():
return K.cuda(), invK.cuda()
return K, invK
def tensor2img(img: Tensor) -> np.ndarray:
return (255.0 * img.permute(1, 2, 0).cpu().detach().numpy()).astype(np.uint8)
def test_simple(image: Image):
"""Function to predict for a single image or folder of images
"""
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
# Loading pretrained model
encoder_dict = torch.load("src/weights/encoder.pth", map_location=device)
encoder = ResnetEncoderMatching(18, False,
input_width=encoder_dict['width'],
input_height=encoder_dict['height'],
adaptive_bins=True,
min_depth_bin=encoder_dict['min_depth_bin'],
max_depth_bin=encoder_dict['max_depth_bin'],
depth_binning='linear',
num_depth_bins=96)
filtered_dict_enc = {k: v for k, v in encoder_dict.items() if k in encoder.state_dict()}
encoder.load_state_dict(filtered_dict_enc)
depth_decoder = DepthDecoder(num_ch_enc=encoder.num_ch_enc, scales=range(4))
loaded_dict = torch.load("src/weights/depth.pth", map_location=device)
depth_decoder.load_state_dict(loaded_dict)
pose_enc_dict = torch.load("src/weights/pose_encoder.pth", map_location=device)
pose_dec_dict = torch.load("src/weights/pose.pth", map_location=device)
pose_enc = ResnetEncoder(18, False, num_input_images=2)
pose_dec = PoseDecoder(pose_enc.num_ch_enc,
num_input_features=1,
num_frames_to_predict_for=2)
pose_enc.load_state_dict(pose_enc_dict, strict=True)
pose_dec.load_state_dict(pose_dec_dict, strict=True)
restoration_dict = torch.load("src/weights/uie_model.pth", map_location=device)
uie_model = MainModel()
uie_model.load_state_dict(restoration_dict, strict=False)
# Setting states of networks
encoder.eval()
depth_decoder.eval()
pose_enc.eval()
pose_dec.eval()
uie_model.eval()
if torch.cuda.is_available():
encoder.cuda()
depth_decoder.cuda()
pose_enc.cuda()
pose_dec.cuda()
uie_model.cuda()
# Load input data
input_image_ori, input_image, original_size = load_and_preprocess_image(image,
resize_width=encoder_dict['width'],
resize_height=encoder_dict['height'])
source_image_ori, source_image, _ = load_and_preprocess_image(image,
resize_width=encoder_dict['width'],
resize_height=encoder_dict['height'])
K, invK = load_and_preprocess_intrinsics('canyons_intrinsics.json',
resize_width=encoder_dict['width'],
resize_height=encoder_dict['height'])
with torch.no_grad():
# Estimate poses
pose_inputs = [source_image, input_image]
pose_inputs = [pose_enc(torch.cat(pose_inputs, 1))]
axisangle, translation = pose_dec(pose_inputs)
pose = transformation_from_parameters(axisangle[:, 0], translation[:, 0], invert=True)
pose *= 0 # zero poses are a signal to the encoder not to construct a cost volume
source_image *= 0
# Estimate depth
output, lowest_cost, _ = encoder(current_image=input_image,
lookup_images=source_image.unsqueeze(1),
poses=pose.unsqueeze(1),
K=K,
invK=invK,
min_depth_bin=encoder_dict['min_depth_bin'],
max_depth_bin=encoder_dict['max_depth_bin'])
output = depth_decoder(output)
sigmoid_output = output[("disp", 0)]
_, depth_output = disp_to_depth(sigmoid_output, min_depth=0.1, max_depth=20)
sigmoid_output_resized = F.interpolate(
sigmoid_output, original_size, mode="bilinear", align_corners=False)
sigmoid_output_resized = sigmoid_output_resized.cpu().numpy()[:, 0]
depth = F.interpolate(
depth_output, input_image_ori.shape[2:], mode="bilinear", align_corners=False)
beta, J, A = uie_model(input_image_ori)
beta[0] = 5.0 * beta[0]
beta[1] = 5.0 * beta[1]
t1 = torch.exp(-beta[0] * depth)
D1 = J * t1
B1 = (1 - torch.exp(-beta[1] * depth)) * A
I_rec = D1 + B1
J_out = Image.open(tensor2img(J[0]))
return J_out
|