Добавлены пропсы конвейера и стереодвижки, задействованные в прогоне
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
+202
@@ -0,0 +1,202 @@
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from os.path import basename, exists, splitext
|
||||
import os,sys
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(f'{code_dir}/../../')
|
||||
import re
|
||||
import json
|
||||
import imageio
|
||||
import cv2
|
||||
from turbojpeg import TurboJPEG, TJPF_GRAY, TJSAMP_GRAY, TJFLAG_PROGRESSIVE, TJFLAG_FASTUPSAMPLE, TJFLAG_FASTDCT
|
||||
jpeg = TurboJPEG()
|
||||
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
|
||||
|
||||
# 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 * 4 + d_g / (2**6) + d_b / (2**14))[..., 0]
|
||||
mask = np.array(Image.open(file_name.replace('disparities', 'occlusions')))
|
||||
valid = ((mask == 0) & (disp > 0))
|
||||
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)
|
||||
fx = intrinsics['camera_settings'][0]['intrinsic_settings']['fx']
|
||||
disp = (fx * 6.0 * 100) / a.astype(np.float32)
|
||||
valid = disp > 0
|
||||
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 > 0
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDispMiddlebury(file_name):
|
||||
assert 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
|
||||
|
||||
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 in ['.jpeg','.jpg']:
|
||||
with open(file_name, 'rb') as ff:
|
||||
bgr_array = jpeg.decode(ff.read())
|
||||
img = bgr_array[...,::-1]
|
||||
return img
|
||||
elif ext == '.png' or ext == '.ppm':
|
||||
img = cv2.imread(file_name)[...,:3]
|
||||
if len(img.shape)==3:
|
||||
img = img[...,::-1]
|
||||
elif len(img.shape)==2:
|
||||
img = np.tile(img[...,None], (1,1,3))
|
||||
return img
|
||||
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]
|
||||
return []
|
||||
Executable
+122
@@ -0,0 +1,122 @@
|
||||
import torch,os,sys
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(f'{code_dir}/../../')
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
|
||||
class InputPadder:
|
||||
""" Pads images such that dimensions are divisible by 8 """
|
||||
def __init__(self, dims, mode='sintel', divis_by=8, force_square=False):
|
||||
self.ht, self.wd = dims[-2:]
|
||||
if force_square:
|
||||
max_side = max(self.ht, self.wd)
|
||||
pad_ht = ((max_side // divis_by) + 1) * divis_by - self.ht
|
||||
pad_wd = ((max_side // divis_by) + 1) * divis_by - self.wd
|
||||
else:
|
||||
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]]
|
||||
|
||||
|
||||
@torch.compile
|
||||
def bilinear_sampler1d(img, x_coords, mode='bilinear', align_corners=True):
|
||||
"""
|
||||
1D bilinear sampling along width dimension only (for stereo applications)
|
||||
Much faster than grid_sample for stereo where y is constant
|
||||
|
||||
Args:
|
||||
img: (B, C, 1, W) input tensor
|
||||
x_coords: (B, 1, W_out, 1) x coordinates in pixel space [0, W-1]
|
||||
mode: interpolation mode ('bilinear' or 'nearest')
|
||||
align_corners: if True, corner pixels are aligned (like grid_sample)
|
||||
|
||||
Returns:
|
||||
sampled: (B, C, 1, W_coords) sampled tensor
|
||||
mask: (B, 1, H, W) validity mask (if mask=True)
|
||||
"""
|
||||
B, C, H_img, W = img.shape
|
||||
x = x_coords.reshape(B,-1) # (B, W_out)
|
||||
|
||||
if align_corners:
|
||||
# align_corners=True: coordinate range [0, W-1] maps to pixel centers
|
||||
# This matches grid_sample with align_corners=True behavior
|
||||
x_normalized = x
|
||||
else:
|
||||
# align_corners=False: coordinate range [0, W-1] maps to pixel edges
|
||||
# Need to adjust coordinates to match grid_sample with align_corners=False
|
||||
# grid_sample maps [-1, 1] to [0, W-1] when align_corners=False
|
||||
# So our [0, W-1] input should be treated as [0.5, W-0.5] in pixel space
|
||||
x_normalized = x + 0.5
|
||||
|
||||
if mode == 'nearest':
|
||||
# Nearest neighbor sampling with zero padding outside [0, W-1]
|
||||
if align_corners:
|
||||
x_nearest = torch.round(x_normalized).long()
|
||||
else:
|
||||
x_nearest = torch.floor(x_normalized).long()
|
||||
valid = (x_nearest >= 0) & (x_nearest < W) # (B, W_out)
|
||||
x_index = torch.clamp(x_nearest, 0, W-1)
|
||||
sampled = torch.gather(img, 3, x_index.view(B,1,1,-1).expand(B,C,1,-1))
|
||||
sampled = sampled * valid.view(B,1,1,-1).to(img.dtype)
|
||||
|
||||
else: # bilinear
|
||||
# Get integer and fractional parts
|
||||
x_floor = torch.floor(x_normalized)
|
||||
x_ceil = x_floor + 1
|
||||
x_frac = x_normalized - x_floor # (B, W_out)
|
||||
|
||||
# Zero padding behavior: mark validity and zero-out invalid contributions
|
||||
valid_floor = (x_floor >= 0) & (x_floor < W)
|
||||
valid_ceil = (x_ceil >= 0) & (x_ceil < W)
|
||||
x_floor_clamped = torch.clamp(x_floor, 0, W-1)
|
||||
x_ceil_clamped = torch.clamp(x_ceil, 0, W-1)
|
||||
|
||||
# Create index tensors
|
||||
batch_idx = torch.arange(B, device=img.device).view(B, 1)
|
||||
img_floor = torch.gather(img, 3, x_floor_clamped.view(B,1,1,-1).expand(B,C,1,-1).long())
|
||||
img_ceil = torch.gather(img, 3, x_ceil_clamped.view(B,1,1,-1).expand(B,C,1,-1).long())
|
||||
|
||||
# Apply validity masks (zero out-of-bounds samples)
|
||||
img_floor = img_floor * valid_floor.view(B,1,1,-1).to(img.dtype)
|
||||
img_ceil = img_ceil * valid_ceil.view(B,1,1,-1).to(img.dtype)
|
||||
|
||||
# Linear interpolation
|
||||
x_frac = x_frac.view(B,1,1,-1)
|
||||
sampled = img_floor * (1 - x_frac) + img_ceil * x_frac
|
||||
|
||||
return sampled
|
||||
|
||||
|
||||
def bilinear_sampler(img, coords, mode='bilinear', mask=False, low_memory=False, use1d=False):
|
||||
""" Wrapper for grid_sample, uses pixel coordinates """
|
||||
H, W = img.shape[-2:]
|
||||
coords[...,0] = 2*coords[...,0]/(W-1) - 1
|
||||
if low_memory:
|
||||
B = img.shape[0]
|
||||
out = []
|
||||
bs = 102400
|
||||
for b in np.arange(0,B,bs):
|
||||
tmp = F.grid_sample(img[b:b+bs], coords[b:b+bs], align_corners=True)
|
||||
out.append(tmp)
|
||||
img = torch.cat(out, dim=0)
|
||||
else:
|
||||
img = F.grid_sample(img, coords, align_corners=True)
|
||||
if mask:
|
||||
mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
|
||||
return img, mask.float()
|
||||
return img
|
||||
|
||||
Reference in New Issue
Block a user