Добавлены пропсы конвейера и стереодвижки, задействованные в прогоне
assets/conveyors (274 МБ) - ленты и угловая секция NVIDIA, на которые ссылается сцена относительным путём. Раньше исключались как перекачиваемые, но без них сцена не композится из коробки. cv/ - код стереодвижков, которые вызывает control_test, без весов: * defom-stereo - рабочий бейзлайн (DEFOM vitl, вход 480, iters 24) * crestereo - второй движок, точнее по габаритам (MAE 23.5 против 32.8 мм) * fast-foundationstereo - проверялся, в бейзлайн не вошёл * circular_section.py - показатель кругового сечения, перенесён в measure_plane.py: выравнивает облако по СОБСТВЕННЫМ главным осям и режет на пяти высотах вдоль каждой. Три самодельные версии (мировые оси, одно сечение) давали хуже; результаты проверки на эталонной геометрии - в circular_section_results.json Веса по-прежнему не в репозитории - источники в MODELS.md. Наборы кадров прежних прогонов (cv/flow_*, 1.26 ГБ) исключены: это выход, а не исходники. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
dasha_f
@@ -0,0 +1,307 @@
|
||||
import numpy as np
|
||||
import random
|
||||
import warnings
|
||||
import os
|
||||
import time
|
||||
from glob import glob
|
||||
from skimage import color, io
|
||||
from PIL import Image
|
||||
|
||||
import cv2
|
||||
cv2.setNumThreads(0)
|
||||
cv2.ocl.setUseOpenCL(False)
|
||||
|
||||
import torch
|
||||
from torchvision.transforms import ColorJitter, functional, Compose
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def get_middlebury_images():
|
||||
root = "../datasets/Middlebury/MiddEval3"
|
||||
with open(os.path.join(root, "official_train.txt"), 'r') as f:
|
||||
lines = f.read().splitlines()
|
||||
return sorted([os.path.join(root, 'trainingQ', f'{name}/im0.png') for name in lines])
|
||||
|
||||
|
||||
def get_eth3d_images():
|
||||
return sorted(glob('../datasets/ETH3D/two_view_training/*/im0.png'))
|
||||
|
||||
|
||||
def get_kitti_images():
|
||||
return sorted(glob('..datasets/KITTI/training/image_2/*_10.png'))
|
||||
|
||||
|
||||
def transfer_color(image, style_mean, style_stddev):
|
||||
reference_image_lab = color.rgb2lab(image)
|
||||
reference_stddev = np.std(reference_image_lab, axis=(0, 1), keepdims=True)# + 1
|
||||
reference_mean = np.mean(reference_image_lab, axis=(0, 1), keepdims=True)
|
||||
|
||||
reference_image_lab = reference_image_lab - reference_mean
|
||||
lamb = style_stddev/reference_stddev
|
||||
style_image_lab = lamb * reference_image_lab
|
||||
output_image_lab = style_image_lab + style_mean
|
||||
l, a, b = np.split(output_image_lab, 3, axis=2)
|
||||
l = l.clip(0, 100)
|
||||
output_image_lab = np.concatenate((l, a, b), axis=2)
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", category=UserWarning)
|
||||
output_image_rgb = color.lab2rgb(output_image_lab) * 255
|
||||
return output_image_rgb
|
||||
|
||||
|
||||
class AdjustGamma(object):
|
||||
|
||||
def __init__(self, gamma_min, gamma_max, gain_min=1.0, gain_max=1.0):
|
||||
self.gamma_min, self.gamma_max, self.gain_min, self.gain_max = gamma_min, gamma_max, gain_min, gain_max
|
||||
|
||||
def __call__(self, sample):
|
||||
gain = random.uniform(self.gain_min, self.gain_max)
|
||||
gamma = random.uniform(self.gamma_min, self.gamma_max)
|
||||
return functional.adjust_gamma(sample, gamma, gain)
|
||||
|
||||
def __repr__(self):
|
||||
return f"Adjust Gamma {self.gamma_min}, ({self.gamma_max}) and Gain ({self.gain_min}, {self.gain_max})"
|
||||
|
||||
|
||||
class DispAugmentor:
|
||||
def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=True, yjitter=False,
|
||||
saturation_range=[0.6, 1.4], gamma=[1, 1, 1, 1]):
|
||||
|
||||
# spatial augmentation params
|
||||
self.crop_size = crop_size
|
||||
self.min_scale = min_scale
|
||||
self.max_scale = max_scale
|
||||
self.spatial_aug_prob = 1.0
|
||||
self.stretch_prob = 0.8
|
||||
self.max_stretch = 0.2
|
||||
|
||||
# flip augmentation params
|
||||
self.yjitter = yjitter
|
||||
self.do_flip = do_flip
|
||||
self.v_flip_prob = 0.1
|
||||
|
||||
# photometric augmentation params
|
||||
self.photo_aug = Compose([ColorJitter(brightness=0.4, contrast=0.4, saturation=saturation_range, hue=0.5/3.14), AdjustGamma(*gamma)])
|
||||
self.asymmetric_color_aug_prob = 0.2
|
||||
self.eraser_aug_prob = 0.5
|
||||
|
||||
def color_transform(self, img1, img2):
|
||||
""" Photometric augmentation """
|
||||
|
||||
# asymmetric
|
||||
if np.random.rand() < self.asymmetric_color_aug_prob:
|
||||
img1 = np.array(self.photo_aug(Image.fromarray(img1)), dtype=np.uint8)
|
||||
img2 = np.array(self.photo_aug(Image.fromarray(img2)), dtype=np.uint8)
|
||||
|
||||
# symmetric
|
||||
else:
|
||||
image_stack = np.concatenate([img1, img2], axis=0)
|
||||
image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8)
|
||||
img1, img2 = np.split(image_stack, 2, axis=0)
|
||||
|
||||
return img1, img2
|
||||
|
||||
def eraser_transform(self, img1, img2, bounds=[50, 100]):
|
||||
""" Occlusion augmentation """
|
||||
|
||||
ht, wd = img1.shape[:2]
|
||||
if np.random.rand() < self.eraser_aug_prob:
|
||||
mean_color = np.mean(img2.reshape(-1, 3), axis=0)
|
||||
for _ in range(np.random.randint(1, 3)):
|
||||
x0 = np.random.randint(0, wd)
|
||||
y0 = np.random.randint(0, ht)
|
||||
dx = np.random.randint(bounds[0], bounds[1])
|
||||
dy = np.random.randint(bounds[0], bounds[1])
|
||||
img2[y0:y0 + dy, x0:x0 + dx, :] = mean_color
|
||||
|
||||
return img1, img2
|
||||
|
||||
def spatial_transform(self, img1, img2, disp):
|
||||
# randomly sample scale
|
||||
ht, wd = img1.shape[:2]
|
||||
min_scale = np.maximum(
|
||||
(self.crop_size[0] + 8) / float(ht),
|
||||
(self.crop_size[1] + 8) / float(wd))
|
||||
|
||||
scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)
|
||||
if scale>min_scale:
|
||||
scale = np.random.uniform(min_scale, scale)
|
||||
scale_x = scale
|
||||
scale_y = scale
|
||||
if np.random.rand() < self.stretch_prob:
|
||||
scale_x *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)
|
||||
scale_y *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)
|
||||
|
||||
scale_x = np.clip(scale_x, min_scale, 2*min_scale)
|
||||
scale_y = np.clip(scale_y, min_scale, 2*min_scale)
|
||||
|
||||
if np.random.rand() < self.spatial_aug_prob or min_scale >= 1.0:
|
||||
# rescale the images
|
||||
img1 = cv2.resize(img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
img2 = cv2.resize(img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
disp = cv2.resize(disp, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
disp = disp * scale_x
|
||||
|
||||
if self.do_flip:
|
||||
if np.random.rand() < self.v_flip_prob and self.do_flip == 'v': # v-flip
|
||||
img1 = img1[::-1, :]
|
||||
img2 = img2[::-1, :]
|
||||
disp = disp[::-1, :]
|
||||
|
||||
if self.yjitter:
|
||||
y0 = np.random.randint(2, img1.shape[0] - self.crop_size[0] - 2)
|
||||
x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1] - 0)
|
||||
|
||||
y1 = y0 + np.random.randint(-2, 2 + 1)
|
||||
y1 = np.clip(y1, 0, img1.shape[0] - self.crop_size[0])
|
||||
img1 = img1[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
img2 = img2[y1:y1 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
disp = disp[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
|
||||
else:
|
||||
y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0])
|
||||
x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1])
|
||||
|
||||
img1 = img1[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
img2 = img2[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
disp = disp[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
|
||||
return img1, img2, disp
|
||||
|
||||
def __call__(self, img1, img2, disp):
|
||||
img1, img2 = self.color_transform(img1, img2)
|
||||
img1, img2 = self.eraser_transform(img1, img2)
|
||||
img1, img2, disp = self.spatial_transform(img1, img2, disp)
|
||||
|
||||
img1 = np.ascontiguousarray(img1)
|
||||
img2 = np.ascontiguousarray(img2)
|
||||
disp = np.ascontiguousarray(disp)
|
||||
|
||||
return img1, img2, disp
|
||||
|
||||
|
||||
class SparseDispAugmentor:
|
||||
def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=False, yjitter=False,
|
||||
saturation_range=[0.7, 1.3], gamma=[1, 1, 1, 1]):
|
||||
# spatial augmentation params
|
||||
self.crop_size = crop_size
|
||||
self.min_scale = min_scale
|
||||
self.max_scale = max_scale
|
||||
self.spatial_aug_prob = 0.8
|
||||
self.stretch_prob = 0.8
|
||||
self.max_stretch = 0.2
|
||||
|
||||
# flip augmentation params
|
||||
self.do_flip = do_flip
|
||||
self.v_flip_prob = 0.1
|
||||
|
||||
# photometric augmentation params
|
||||
self.photo_aug = Compose(
|
||||
[ColorJitter(brightness=0.3, contrast=0.3, saturation=saturation_range, hue=0.3/3.14),
|
||||
AdjustGamma(*gamma)])
|
||||
self.asymmetric_color_aug_prob = 0.2
|
||||
self.eraser_aug_prob = 0.5
|
||||
|
||||
def color_transform(self, img1, img2):
|
||||
image_stack = np.concatenate([img1, img2], axis=0)
|
||||
image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8)
|
||||
img1, img2 = np.split(image_stack, 2, axis=0)
|
||||
return img1, img2
|
||||
|
||||
def eraser_transform(self, img1, img2):
|
||||
ht, wd = img1.shape[:2]
|
||||
if np.random.rand() < self.eraser_aug_prob:
|
||||
mean_color = np.mean(img2.reshape(-1, 3), axis=0)
|
||||
for _ in range(np.random.randint(1, 3)):
|
||||
x0 = np.random.randint(0, wd)
|
||||
y0 = np.random.randint(0, ht)
|
||||
dx = np.random.randint(50, 100)
|
||||
dy = np.random.randint(50, 100)
|
||||
img2[y0:y0 + dy, x0:x0 + dx, :] = mean_color
|
||||
|
||||
return img1, img2
|
||||
|
||||
def resize_sparse_flow_map(self, disp, valid, fx=1.0, fy=1.0):
|
||||
ht, wd = disp.shape[:2]
|
||||
coords = np.meshgrid(np.arange(wd), np.arange(ht))
|
||||
coords = np.stack(coords, axis=-1)
|
||||
|
||||
coords = coords.reshape(-1, 2).astype(np.float32)
|
||||
disp = disp.reshape(-1).astype(np.float32)
|
||||
valid = valid.reshape(-1).astype(np.float32)
|
||||
|
||||
coords0 = coords[valid >= 1]
|
||||
disp0 = disp[valid >= 1]
|
||||
|
||||
ht1 = int(round(ht * fy))
|
||||
wd1 = int(round(wd * fx))
|
||||
|
||||
coords1 = coords0 * [fx, fy]
|
||||
disp1 = disp0 * fx
|
||||
|
||||
xx = np.round(coords1[:, 0]).astype(np.int32)
|
||||
yy = np.round(coords1[:, 1]).astype(np.int32)
|
||||
|
||||
v = (xx > 0) & (xx < wd1) & (yy > 0) & (yy < ht1)
|
||||
xx = xx[v]
|
||||
yy = yy[v]
|
||||
disp1 = disp1[v]
|
||||
|
||||
disp_img = np.zeros([ht1, wd1], dtype=np.float32)
|
||||
valid_img = np.zeros([ht1, wd1], dtype=np.int32)
|
||||
|
||||
disp_img[yy, xx] = disp1
|
||||
valid_img[yy, xx] = 1
|
||||
|
||||
return disp_img, valid_img
|
||||
|
||||
def spatial_transform(self, img1, img2, disp, valid):
|
||||
# randomly sample scale
|
||||
|
||||
ht, wd = img1.shape[:2]
|
||||
min_scale = np.maximum(
|
||||
(self.crop_size[0] + 1) / float(ht),
|
||||
(self.crop_size[1] + 1) / float(wd))
|
||||
|
||||
scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)
|
||||
if scale>min_scale:
|
||||
scale = np.random.uniform(min_scale, 2*min_scale)
|
||||
scale_x = scale
|
||||
scale_y = scale
|
||||
|
||||
scale_x = np.clip(scale_x, min_scale, 2*min_scale)
|
||||
scale_y = np.clip(scale_y, min_scale, 2*min_scale)
|
||||
|
||||
if np.random.rand() < self.spatial_aug_prob or min_scale >= 1.0:
|
||||
# rescale the images
|
||||
img1 = cv2.resize(img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
img2 = cv2.resize(img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
disp, valid = self.resize_sparse_flow_map(disp, valid, fx=scale_x, fy=scale_y)
|
||||
|
||||
if self.do_flip:
|
||||
if np.random.rand() < self.v_flip_prob and self.do_flip == 'v': # v-flip
|
||||
img1 = img1[::-1, :]
|
||||
img2 = img2[::-1, :]
|
||||
disp = disp[::-1, :]
|
||||
valid = valid[::-1, :]
|
||||
|
||||
y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0])
|
||||
x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1])
|
||||
|
||||
img1 = img1[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
img2 = img2[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
disp = disp[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
valid = valid[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
return img1, img2, disp, valid
|
||||
|
||||
def __call__(self, img1, img2, disp, valid):
|
||||
img1, img2 = self.color_transform(img1, img2)
|
||||
img1, img2 = self.eraser_transform(img1, img2)
|
||||
img1, img2, disp, valid = self.spatial_transform(img1, img2, disp, valid)
|
||||
|
||||
img1 = np.ascontiguousarray(img1)
|
||||
img2 = np.ascontiguousarray(img2)
|
||||
disp = np.ascontiguousarray(disp)
|
||||
valid = np.ascontiguousarray(valid)
|
||||
|
||||
return img1, img2, disp, valid
|
||||
@@ -0,0 +1,105 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
# https://github.com/open-mmlab/mmcv/blob/7540cf73ac7e5d1e14d0ffbd9b6759e83929ecfc/mmcv/runner/dist_utils.py
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
import torch
|
||||
import torch.multiprocessing as mp
|
||||
from torch import distributed as dist
|
||||
|
||||
|
||||
def init_dist(launcher, backend='nccl', **kwargs):
|
||||
if mp.get_start_method(allow_none=True) is None:
|
||||
mp.set_start_method('spawn')
|
||||
if launcher == 'pytorch':
|
||||
_init_dist_pytorch(backend, **kwargs)
|
||||
elif launcher == 'mpi':
|
||||
_init_dist_mpi(backend, **kwargs)
|
||||
elif launcher == 'slurm':
|
||||
_init_dist_slurm(backend, **kwargs)
|
||||
else:
|
||||
raise ValueError(f'Invalid launcher type: {launcher}')
|
||||
|
||||
|
||||
def _init_dist_pytorch(backend, **kwargs):
|
||||
# TODO: use local_rank instead of rank % num_gpus
|
||||
rank = int(os.environ['RANK'])
|
||||
num_gpus = torch.cuda.device_count()
|
||||
torch.cuda.set_device(rank % num_gpus)
|
||||
dist.init_process_group(backend=backend, **kwargs)
|
||||
|
||||
|
||||
def _init_dist_mpi(backend, **kwargs):
|
||||
# TODO: use local_rank instead of rank % num_gpus
|
||||
rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
|
||||
num_gpus = torch.cuda.device_count()
|
||||
torch.cuda.set_device(rank % num_gpus)
|
||||
dist.init_process_group(backend=backend, **kwargs)
|
||||
|
||||
|
||||
def _init_dist_slurm(backend, port=None):
|
||||
"""Initialize slurm distributed training environment.
|
||||
If argument ``port`` is not specified, then the master port will be system
|
||||
environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system
|
||||
environment variable, then a default port ``29500`` will be used.
|
||||
Args:
|
||||
backend (str): Backend of torch.distributed.
|
||||
port (int, optional): Master port. Defaults to None.
|
||||
"""
|
||||
proc_id = int(os.environ['SLURM_PROCID'])
|
||||
ntasks = int(os.environ['SLURM_NTASKS'])
|
||||
node_list = os.environ['SLURM_NODELIST']
|
||||
num_gpus = torch.cuda.device_count()
|
||||
torch.cuda.set_device(proc_id % num_gpus)
|
||||
addr = subprocess.getoutput(
|
||||
f'scontrol show hostname {node_list} | head -n1')
|
||||
# specify master port
|
||||
if port is not None:
|
||||
os.environ['MASTER_PORT'] = str(port)
|
||||
elif 'MASTER_PORT' in os.environ:
|
||||
pass # use MASTER_PORT in the environment variable
|
||||
else:
|
||||
# 29500 is torch.distributed default port
|
||||
os.environ['MASTER_PORT'] = '29500'
|
||||
# use MASTER_ADDR in the environment variable if it already exists
|
||||
if 'MASTER_ADDR' not in os.environ:
|
||||
os.environ['MASTER_ADDR'] = addr
|
||||
os.environ['WORLD_SIZE'] = str(ntasks)
|
||||
os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
|
||||
os.environ['RANK'] = str(proc_id)
|
||||
dist.init_process_group(backend=backend)
|
||||
|
||||
|
||||
def get_dist_info():
|
||||
# if (TORCH_VERSION != 'parrots'
|
||||
# and digit_version(TORCH_VERSION) < digit_version('1.0')):
|
||||
# initialized = dist._initialized
|
||||
# else:
|
||||
if dist.is_available():
|
||||
initialized = dist.is_initialized()
|
||||
else:
|
||||
initialized = False
|
||||
if initialized:
|
||||
rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
else:
|
||||
rank = 0
|
||||
world_size = 1
|
||||
return rank, world_size
|
||||
|
||||
|
||||
# from DETR repo
|
||||
def setup_for_distributed(is_master):
|
||||
"""
|
||||
This function disables printing when not in master process
|
||||
"""
|
||||
import builtins as __builtin__
|
||||
builtin_print = __builtin__.print
|
||||
|
||||
def print(*args, **kwargs):
|
||||
force = kwargs.pop('force', False)
|
||||
if is_master or force:
|
||||
builtin_print(*args, **kwargs)
|
||||
|
||||
__builtin__.print = print
|
||||
@@ -0,0 +1,286 @@
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from os.path import *
|
||||
import re
|
||||
import json
|
||||
import imageio
|
||||
import os
|
||||
import math
|
||||
|
||||
os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
|
||||
import cv2
|
||||
cv2.setNumThreads(0)
|
||||
cv2.ocl.setUseOpenCL(False)
|
||||
|
||||
TAG_CHAR = np.array([202021.25], np.float32)
|
||||
|
||||
|
||||
def readFlow(fn):
|
||||
""" Read .flo file in Middlebury format"""
|
||||
# Code adapted from:
|
||||
# http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy
|
||||
|
||||
# WARNING: this will work on little-endian architectures (eg Intel x86) only!
|
||||
# print 'fn = %s'%(fn)
|
||||
with open(fn, 'rb') as f:
|
||||
magic = np.fromfile(f, np.float32, count=1)
|
||||
if 202021.25 != magic:
|
||||
print('Magic number incorrect. Invalid .flo file')
|
||||
return None
|
||||
else:
|
||||
w = np.fromfile(f, np.int32, count=1)
|
||||
h = np.fromfile(f, np.int32, count=1)
|
||||
# print 'Reading %d x %d flo file\n' % (w, h)
|
||||
data = np.fromfile(f, np.float32, count=2*int(w)*int(h))
|
||||
# Reshape data into 3D array (columns, rows, bands)
|
||||
# The reshape here is for visualization, the original code is (w,h,2)
|
||||
return np.resize(data, (int(h), int(w), 2))
|
||||
|
||||
|
||||
def readPFM(file):
|
||||
file = open(file, 'rb')
|
||||
|
||||
color = None
|
||||
width = None
|
||||
height = None
|
||||
scale = None
|
||||
endian = None
|
||||
|
||||
header = file.readline().rstrip()
|
||||
if header == b'PF':
|
||||
color = True
|
||||
elif header == b'Pf':
|
||||
color = False
|
||||
else:
|
||||
raise Exception('Not a PFM file.')
|
||||
|
||||
dim_match = re.match(rb'^(\d+)\s(\d+)\s$', file.readline())
|
||||
if dim_match:
|
||||
width, height = map(int, dim_match.groups())
|
||||
else:
|
||||
raise Exception('Malformed PFM header.')
|
||||
|
||||
scale = float(file.readline().rstrip())
|
||||
if scale < 0: # little-endian
|
||||
endian = '<'
|
||||
scale = -scale
|
||||
else:
|
||||
endian = '>' # big-endian
|
||||
|
||||
data = np.fromfile(file, endian + 'f')
|
||||
shape = (height, width, 3) if color else (height, width)
|
||||
|
||||
data = np.reshape(data, shape)
|
||||
data = np.flipud(data)
|
||||
return data
|
||||
|
||||
|
||||
def writePFM(file, array):
|
||||
import os
|
||||
assert type(file) is str and type(array) is np.ndarray and \
|
||||
os.path.splitext(file)[1] == ".pfm"
|
||||
with open(file, 'wb') as f:
|
||||
H, W = array.shape
|
||||
headers = ["Pf\n", f"{W} {H}\n", "-1\n"]
|
||||
for header in headers:
|
||||
f.write(str.encode(header))
|
||||
array = np.flip(array, axis=0).astype(np.float32)
|
||||
f.write(array.tobytes())
|
||||
|
||||
|
||||
def writeFlow(filename,uv,v=None):
|
||||
""" Write optical flow to file.
|
||||
|
||||
If v is None, uv is assumed to contain both u and v channels,
|
||||
stacked in depth.
|
||||
Original code by Deqing Sun, adapted from Daniel Scharstein.
|
||||
"""
|
||||
nBands = 2
|
||||
|
||||
if v is None:
|
||||
assert(uv.ndim == 3)
|
||||
assert(uv.shape[2] == 2)
|
||||
u = uv[:,:,0]
|
||||
v = uv[:,:,1]
|
||||
else:
|
||||
u = uv
|
||||
|
||||
assert(u.shape == v.shape)
|
||||
height,width = u.shape
|
||||
f = open(filename,'wb')
|
||||
# write the header
|
||||
f.write(TAG_CHAR)
|
||||
np.array(width).astype(np.int32).tofile(f)
|
||||
np.array(height).astype(np.int32).tofile(f)
|
||||
# arrange into matrix form
|
||||
tmp = np.zeros((height, width*nBands))
|
||||
tmp[:,np.arange(width)*2] = u
|
||||
tmp[:,np.arange(width)*2 + 1] = v
|
||||
tmp.astype(np.float32).tofile(f)
|
||||
f.close()
|
||||
|
||||
|
||||
def readFlowKITTI(filename):
|
||||
flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH|cv2.IMREAD_COLOR)
|
||||
flow = flow[:,:,::-1].astype(np.float32)
|
||||
flow, valid = flow[:, :, :2], flow[:, :, 2]
|
||||
flow = (flow - 2**15) / 64.0
|
||||
return flow, valid
|
||||
|
||||
|
||||
def readDispKITTI(filename):
|
||||
disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 256.0
|
||||
valid = disp > 0.0
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDispInStereo2K(filename):
|
||||
disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 100.0
|
||||
valid = disp > 0.0
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDispVKITTI2(filename):
|
||||
depth = cv2.imread(filename, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH).astype(np.float32) / 100.0
|
||||
valid = depth > 0.0
|
||||
baseline = 0.532725
|
||||
focus_length = 725.0087
|
||||
disp = baseline*focus_length/(depth+1e-8)
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDispCreStereo(filename):
|
||||
disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 32
|
||||
valid = disp > -1e-8
|
||||
return disp, valid
|
||||
|
||||
|
||||
# Method taken from /n/fs/raft-depth/RAFT-Stereo/datasets/SintelStereo/sdk/python/sintel_io.py
|
||||
def readDispSintelStereo(file_name):
|
||||
a = np.array(Image.open(file_name))
|
||||
d_r, d_g, d_b = np.split(a, axis=2, indices_or_sections=3)
|
||||
disp = (d_r.astype('float64') * 4 + d_g.astype('float64') / (2**6) + d_b.astype('float64') / (2**14))[..., 0]
|
||||
mask = np.array(Image.open(file_name.replace('disparities', 'occlusions')))
|
||||
valid = ((mask == 0) & (disp > -1e-8))
|
||||
return disp, valid
|
||||
|
||||
|
||||
# Method taken from https://research.nvidia.com/sites/default/files/pubs/2018-06_Falling-Things/readme_0.txt
|
||||
def readDispFallingThings(file_name):
|
||||
a = np.array(Image.open(file_name))
|
||||
with open('/'.join(file_name.split('/')[:-1] + ['_camera_settings.json']), 'r') as f:
|
||||
intrinsics = json.load(f)
|
||||
if 'left' in file_name:
|
||||
idx = 0
|
||||
else:
|
||||
idx = 1
|
||||
fx = intrinsics['camera_settings'][idx]['intrinsic_settings']['fx']
|
||||
disp = (fx * 6.0 * 100) / a.astype(np.float32)
|
||||
valid = disp > -1e-8
|
||||
return disp, valid
|
||||
|
||||
|
||||
# Method taken from https://github.com/castacks/tartanair_tools/blob/master/data_type.md
|
||||
def readDispTartanAir(file_name):
|
||||
depth = np.load(file_name)
|
||||
disp = 80.0 / depth
|
||||
valid = disp > -1e-8
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDispBooster(file_name):
|
||||
disp = np.load(file_name)
|
||||
valid = disp > 0
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDisp3DKenBurns(file_name):
|
||||
depth = cv2.imread(file_name, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
|
||||
meta_file_name = file_name.replace('-depth', '')[:-7]+'-meta.json'
|
||||
fltFov = json.loads(open(meta_file_name, 'r').read())['fltFov']
|
||||
fltFocal = 0.5 * 512 * math.tan(math.radians(90.0) - (0.5 * math.radians(fltFov)))
|
||||
fltBaseline = 40.0
|
||||
disp = (fltFocal * fltBaseline) / depth
|
||||
valid = disp > 0
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDispMiddlebury0(file_name):
|
||||
if basename(file_name) == 'disp0GT.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
assert len(disp.shape) == 2
|
||||
nocc_pix = file_name.replace('disp0GT.pfm', 'mask0nocc.png')
|
||||
assert exists(nocc_pix)
|
||||
nocc_pix = imageio.imread(nocc_pix) == 255
|
||||
assert np.any(nocc_pix)
|
||||
return disp, nocc_pix
|
||||
elif basename(file_name) == 'disp1GT.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
assert len(disp.shape) == 2
|
||||
nocc_pix = file_name.replace('disp1GT.pfm', 'mask1nocc.png')
|
||||
assert exists(nocc_pix)
|
||||
nocc_pix = imageio.imread(nocc_pix) == 255
|
||||
assert np.any(nocc_pix)
|
||||
return disp, nocc_pix
|
||||
elif basename(file_name) == 'disp0.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
valid = disp < 1e3
|
||||
return disp, valid
|
||||
elif basename(file_name) == 'disp1.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
valid = disp < 1e3
|
||||
return disp, valid
|
||||
elif splitext(file_name)[-1] == '.png':
|
||||
disp = np.array(Image.open(file_name)).astype(np.float32)
|
||||
valid = disp > 0.0
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDispMiddlebury(file_name):
|
||||
if basename(file_name) == 'disp0GT.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
return disp, disp<1e3
|
||||
elif basename(file_name) == 'disp1GT.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
return disp, disp<1e3
|
||||
elif basename(file_name) == 'disp0.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
valid = disp < 1e3
|
||||
return disp, valid
|
||||
elif basename(file_name) == 'disp1.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
valid = disp < 1e3
|
||||
return disp, valid
|
||||
elif splitext(file_name)[-1] == '.png':
|
||||
disp = np.array(Image.open(file_name)).astype(np.float32)
|
||||
valid = disp > 0.0
|
||||
return disp, valid
|
||||
|
||||
|
||||
def writeFlowKITTI(filename, uv):
|
||||
uv = 64.0 * uv + 2**15
|
||||
valid = np.ones([uv.shape[0], uv.shape[1], 1])
|
||||
uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16)
|
||||
cv2.imwrite(filename, uv[..., ::-1])
|
||||
|
||||
|
||||
def read_gen(file_name, pil=False):
|
||||
ext = splitext(file_name)[-1]
|
||||
if ext == '.png' or ext == '.jpeg' or ext == '.ppm' or ext == '.jpg':
|
||||
return Image.open(file_name)
|
||||
elif ext == '.bin' or ext == '.raw':
|
||||
return np.load(file_name)
|
||||
elif ext == '.flo':
|
||||
return readFlow(file_name).astype(np.float32)
|
||||
elif ext == '.pfm':
|
||||
flow = readPFM(file_name).astype(np.float32)
|
||||
if len(flow.shape) == 2:
|
||||
return flow
|
||||
else:
|
||||
return flow[:, :, :-1]
|
||||
elif ext == '.exr':
|
||||
disp = cv2.imread(file_name, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
|
||||
if len(disp.shape) > 2:
|
||||
disp = disp[..., 0]
|
||||
return disp
|
||||
return []
|
||||
@@ -0,0 +1,242 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from scipy import interpolate
|
||||
import glob
|
||||
import os.path as osp
|
||||
|
||||
|
||||
def get_danv2_io_size(h, w, nds, max_i_size=2688, multiple_of=14):
|
||||
"""compute the input and output sizes of danv2 network"""
|
||||
danv2_oh, danv2_ow = h//2**nds, w//2**nds
|
||||
danv2_io_factor = 3.5 # more precise, 14/8=3.5
|
||||
ih, iw = danv2_io_factor*danv2_oh, danv2_io_factor*danv2_ow
|
||||
ih = int(np.ceil(ih / multiple_of) * multiple_of)
|
||||
iw = int(np.ceil(iw / multiple_of) * multiple_of)
|
||||
|
||||
max_i_size = int(np.floor(max_i_size / multiple_of) * multiple_of)
|
||||
|
||||
if ih <= max_i_size and iw <= max_i_size:
|
||||
danv2_ih, danv2_iw = ih, iw
|
||||
else:
|
||||
factor_h = max_i_size/ih
|
||||
factor_w = max_i_size/iw
|
||||
|
||||
if factor_w > factor_h:
|
||||
danv2_ih = max_i_size
|
||||
danv2_iw = int(np.ceil(factor_h * iw / multiple_of) * multiple_of)
|
||||
else:
|
||||
danv2_iw = max_i_size
|
||||
danv2_ih = int(np.ceil(factor_w * ih / multiple_of) * multiple_of)
|
||||
|
||||
return danv2_ih, danv2_iw, danv2_oh, danv2_ow
|
||||
|
||||
|
||||
class InputPadder:
|
||||
""" Pads images such that dimensions are divisible by 8 """
|
||||
def __init__(self, dims, mode='sintel', divis_by=8):
|
||||
self.ht, self.wd = dims[-2:]
|
||||
pad_ht = (((self.ht // divis_by) + 1) * divis_by - self.ht) % divis_by
|
||||
pad_wd = (((self.wd // divis_by) + 1) * divis_by - self.wd) % divis_by
|
||||
if mode == 'sintel':
|
||||
self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2]
|
||||
else:
|
||||
self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht]
|
||||
|
||||
def pad(self, *inputs):
|
||||
assert all((x.ndim == 4) for x in inputs)
|
||||
return [F.pad(x, self._pad, mode='replicate') for x in inputs]
|
||||
|
||||
def unpad(self, x):
|
||||
assert x.ndim == 4
|
||||
ht, wd = x.shape[-2:]
|
||||
c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]
|
||||
return x[..., c[0]:c[1], c[2]:c[3]]
|
||||
|
||||
|
||||
def forward_interpolate(flow):
|
||||
flow = flow.detach().cpu().numpy()
|
||||
dx, dy = flow[0], flow[1]
|
||||
|
||||
ht, wd = dx.shape
|
||||
x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht))
|
||||
|
||||
x1 = x0 + dx
|
||||
y1 = y0 + dy
|
||||
|
||||
x1 = x1.reshape(-1)
|
||||
y1 = y1.reshape(-1)
|
||||
dx = dx.reshape(-1)
|
||||
dy = dy.reshape(-1)
|
||||
|
||||
valid = (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht)
|
||||
x1 = x1[valid]
|
||||
y1 = y1[valid]
|
||||
dx = dx[valid]
|
||||
dy = dy[valid]
|
||||
|
||||
flow_x = interpolate.griddata(
|
||||
(x1, y1), dx, (x0, y0), method='nearest', fill_value=0)
|
||||
|
||||
flow_y = interpolate.griddata(
|
||||
(x1, y1), dy, (x0, y0), method='nearest', fill_value=0)
|
||||
|
||||
flow = np.stack([flow_x, flow_y], axis=0)
|
||||
return torch.from_numpy(flow).float()
|
||||
|
||||
|
||||
def bilinear_sampler(img, coords, mode='bilinear', mask=False):
|
||||
""" Wrapper for grid_sample, uses pixel coordinates """
|
||||
H, W = img.shape[-2:]
|
||||
xgrid, ygrid = coords.split([1, 1], dim=-1)
|
||||
xgrid = 2*xgrid/(W-1) - 1
|
||||
if H > 1:
|
||||
ygrid = 2*ygrid/(H-1) - 1
|
||||
|
||||
grid = torch.cat([xgrid, ygrid], dim=-1)
|
||||
img = F.grid_sample(img, grid, align_corners=True)
|
||||
# img = bilinear_grid_sample(img, grid, align_corners=True)
|
||||
|
||||
if mask:
|
||||
mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
|
||||
return img, mask.float()
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def coords_grid(batch, ht, wd):
|
||||
coords = torch.meshgrid(torch.arange(ht), torch.arange(wd))
|
||||
coords = torch.stack(coords[::-1], dim=0).float()
|
||||
return coords[None].repeat(batch, 1, 1, 1)
|
||||
|
||||
|
||||
def upflow(flow, factor=8, mode='bilinear', sacle=True):
|
||||
new_size = (factor * flow.shape[2], factor * flow.shape[3])
|
||||
if sacle:
|
||||
return factor * F.interpolate(flow, size=new_size, mode=mode, align_corners=True)
|
||||
else:
|
||||
return F.interpolate(flow, size=new_size, mode=mode, align_corners=True)
|
||||
|
||||
|
||||
def gauss_blur(input, N=5, std=1):
|
||||
B, D, H, W = input.shape
|
||||
x, y = torch.meshgrid(torch.arange(N).float() - N//2, torch.arange(N).float() - N//2)
|
||||
unnormalized_gaussian = torch.exp(-(x.pow(2) + y.pow(2)) / (2 * std ** 2))
|
||||
weights = unnormalized_gaussian / unnormalized_gaussian.sum().clamp(min=1e-4)
|
||||
weights = weights.view(1, 1, N, N).to(input)
|
||||
output = F.conv2d(input.reshape(B*D, 1, H, W), weights, padding=N//2)
|
||||
return output.view(B, D, H, W)
|
||||
|
||||
|
||||
# Ref: https://zenn.dev/pinto0309/scraps/7d4032067d0160
|
||||
def bilinear_grid_sample(im, grid, align_corners=False):
|
||||
"""Given an input and a flow-field grid, computes the output using input
|
||||
values and pixel locations from grid. Supported only bilinear interpolation
|
||||
method to sample the input pixels.
|
||||
|
||||
Args:
|
||||
im (torch.Tensor): Input feature map, shape (N, C, H, W)
|
||||
grid (torch.Tensor): Point coordinates, shape (N, Hg, Wg, 2)
|
||||
align_corners {bool}: If set to True, the extrema (-1 and 1) are
|
||||
considered as referring to the center points of the input’s
|
||||
corner pixels. If set to False, they are instead considered as
|
||||
referring to the corner points of the input’s corner pixels,
|
||||
making the sampling more resolution agnostic.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: A tensor with sampled points, shape (N, C, Hg, Wg)
|
||||
"""
|
||||
n, c, h, w = im.shape
|
||||
gn, gh, gw, _ = grid.shape
|
||||
assert n == gn
|
||||
|
||||
x = grid[:, :, :, 0]
|
||||
y = grid[:, :, :, 1]
|
||||
|
||||
if align_corners:
|
||||
x = ((x + 1) / 2) * (w - 1)
|
||||
y = ((y + 1) / 2) * (h - 1)
|
||||
else:
|
||||
x = ((x + 1) * w - 1) / 2
|
||||
y = ((y + 1) * h - 1) / 2
|
||||
|
||||
x = x.view(n, -1)
|
||||
y = y.view(n, -1)
|
||||
|
||||
x0 = torch.floor(x).long()
|
||||
y0 = torch.floor(y).long()
|
||||
x1 = x0 + 1
|
||||
y1 = y0 + 1
|
||||
|
||||
wa = ((x1 - x) * (y1 - y)).unsqueeze(1)
|
||||
wb = ((x1 - x) * (y - y0)).unsqueeze(1)
|
||||
wc = ((x - x0) * (y1 - y)).unsqueeze(1)
|
||||
wd = ((x - x0) * (y - y0)).unsqueeze(1)
|
||||
|
||||
# Apply default for grid_sample function zero padding
|
||||
im_padded = torch.nn.functional.pad(im, pad=[1, 1, 1, 1], mode='constant', value=0)
|
||||
padded_h = h + 2
|
||||
padded_w = w + 2
|
||||
# save points positions after padding
|
||||
x0, x1, y0, y1 = x0 + 1, x1 + 1, y0 + 1, y1 + 1
|
||||
|
||||
# Clip coordinates to padded image size
|
||||
x0 = torch.where(x0 < 0, torch.tensor(0, device=im.device), x0)
|
||||
x0 = torch.where(x0 > padded_w - 1, torch.tensor(padded_w - 1, device=im.device), x0)
|
||||
x1 = torch.where(x1 < 0, torch.tensor(0, device=im.device), x1)
|
||||
x1 = torch.where(x1 > padded_w - 1, torch.tensor(padded_w - 1, device=im.device), x1)
|
||||
y0 = torch.where(y0 < 0, torch.tensor(0, device=im.device), y0)
|
||||
y0 = torch.where(y0 > padded_h - 1, torch.tensor(padded_h - 1, device=im.device), y0)
|
||||
y1 = torch.where(y1 < 0, torch.tensor(0, device=im.device), y1)
|
||||
y1 = torch.where(y1 > padded_h - 1, torch.tensor(padded_h - 1, device=im.device), y1)
|
||||
|
||||
im_padded = im_padded.view(n, c, -1)
|
||||
|
||||
x0_y0 = (x0 + y0 * padded_w).unsqueeze(1).expand(-1, c, -1)
|
||||
x0_y1 = (x0 + y1 * padded_w).unsqueeze(1).expand(-1, c, -1)
|
||||
x1_y0 = (x1 + y0 * padded_w).unsqueeze(1).expand(-1, c, -1)
|
||||
x1_y1 = (x1 + y1 * padded_w).unsqueeze(1).expand(-1, c, -1)
|
||||
|
||||
Ia = torch.gather(im_padded, 2, x0_y0)
|
||||
Ib = torch.gather(im_padded, 2, x0_y1)
|
||||
Ic = torch.gather(im_padded, 2, x1_y0)
|
||||
Id = torch.gather(im_padded, 2, x1_y1)
|
||||
|
||||
return (Ia * wa + Ib * wb + Ic * wc + Id * wd).reshape(n, c, gh, gw)
|
||||
|
||||
|
||||
def read_kitti_calib_file(path):
|
||||
"""Read KITTI calibration file
|
||||
(from https://github.com/hunse/kitti)
|
||||
"""
|
||||
float_chars = set("0123456789.e+- ")
|
||||
data = {}
|
||||
with open(path, 'r') as f:
|
||||
for line in f.readlines():
|
||||
key, value = line.split(':', 1)
|
||||
value = value.strip()
|
||||
data[key] = value
|
||||
if float_chars.issuperset(value):
|
||||
# try to cast to float array
|
||||
try:
|
||||
data[key] = np.array(list(map(float, value.split(' '))))
|
||||
except ValueError:
|
||||
# casting error: data[key] already eq. value, so pass
|
||||
pass
|
||||
|
||||
return data
|
||||
|
||||
|
||||
# from https://github.com/ozendelait/rvc_devkit/blob/master/stereo/stereo_devkit.py
|
||||
def ReadMiddlebury2014CalibFile(path):
|
||||
result = dict()
|
||||
with open(path, 'rb') as calib_file:
|
||||
for line in calib_file.readlines():
|
||||
line = line.decode('UTF-8').rstrip('\n')
|
||||
if len(line) == 0:
|
||||
continue
|
||||
eq_pos = line.find('=')
|
||||
if eq_pos < 0:
|
||||
raise Exception('Cannot parse Middlebury 2014 calib file: ' + path)
|
||||
result[line[:eq_pos]] = line[eq_pos + 1:]
|
||||
return result
|
||||
Reference in New Issue
Block a user