Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,686 Bytes
98844c3 94c12b7 98844c3 94c12b7 98844c3 |
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 |
import numpy as np
import imageio
import torch
import os
import matplotlib.pyplot as plt
import cv2
from promptda.utils.logger import Log
# DEVICE = 'cuda' if torch.cuda.is_available(
# ) else 'mps' if torch.backends.mps.is_available() else 'cpu'
def to_tensor_func(arr):
if arr.ndim == 2:
arr = arr[:, :, np.newaxis]
return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
def to_numpy_func(tensor):
arr = tensor.squeeze(0).permute(1, 2, 0).cpu().numpy()
if arr.shape[2] == 1:
arr = arr[:, :, 0]
return arr
def ensure_multiple_of(x, multiple_of=14):
return int(x // multiple_of * multiple_of)
def load_image(image_path, to_tensor=True, max_size=1008, multiple_of=14):
'''
Load image from path and convert to tensor
max_size // 14 = 0
'''
image = np.asarray(imageio.imread(image_path)).astype(np.float32)
image = image / 255.
max_size = max_size // multiple_of * multiple_of
if max(image.shape) > max_size:
h, w = image.shape[:2]
scale = max_size / max(h, w)
tar_h = ensure_multiple_of(h * scale)
tar_w = ensure_multiple_of(w * scale)
image = cv2.resize(image, (tar_w, tar_h), interpolation=cv2.INTER_AREA)
if to_tensor:
return to_tensor_func(image)
return image
def load_depth(depth_path, to_tensor=True):
'''
Load depth from path and convert to tensor
'''
if depth_path.endswith('.png'):
depth = np.asarray(imageio.imread(depth_path)).astype(np.float32)
depth = depth / 1000.
elif depth_path.endswith('.npz'):
depth = np.load(depth_path)['depth']
else:
raise ValueError(f"Unsupported depth format: {depth_path}")
if to_tensor:
return to_tensor_func(depth)
return depth
def save_depth(depth,
prompt_depth=None,
image=None,
output_path='data/output/depth.png',
save_vis=True):
'''
Save depth to path
'''
os.makedirs(os.path.dirname(output_path), exist_ok=True)
depth = to_numpy_func(depth)
uint16_depth = (depth * 1000.).astype(np.uint16)
imageio.imwrite(output_path, uint16_depth)
if not save_vis:
return
output_path = output_path.replace('.png', '_vis.png')
prompt_depth = to_numpy_func(prompt_depth)
image = to_numpy_func(image)
plt.subplot(1, 3, 1)
plt.imshow(image)
plt.axis('off')
plt.subplot(1, 3, 2)
plt.imshow(prompt_depth)
plt.axis('off')
plt.subplot(1, 3, 3)
plt.imshow(depth)
plt.axis('off')
plt.tight_layout()
plt.savefig(output_path)
plt.close()
Log.info(f'Saved depth to {output_path}')
|