Добавлены пропсы конвейера и стереодвижки, задействованные в прогоне
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,237 @@
|
||||
from __future__ import print_function, division
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import time
|
||||
import os
|
||||
import cv2
|
||||
import sys
|
||||
|
||||
|
||||
from core.defom_stereo import DEFOMStereo, autocast
|
||||
|
||||
import core.stereo_datasets as datasets
|
||||
from core.utils.utils import InputPadder
|
||||
from core.utils.frame_utils import writePFM
|
||||
|
||||
|
||||
def makedirs(path):
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path)
|
||||
|
||||
|
||||
def StrToBytes(text):
|
||||
if sys.version_info[0] == 2:
|
||||
return text
|
||||
else:
|
||||
return bytes(text, 'UTF-8')
|
||||
|
||||
|
||||
def count_parameters(model):
|
||||
return sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test_eth3d(model, save_path, iters=32, scale_iters=3, mixed_prec=False):
|
||||
""" Peform validation using the ETH3D (train) split """
|
||||
model.eval()
|
||||
aug_params = {}
|
||||
|
||||
test_dataset = datasets.ETH3D(aug_params, split='testing', is_test=True)
|
||||
training_dataset = datasets.ETH3D(aug_params, split='training', is_test=True)
|
||||
dataset = test_dataset + training_dataset
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
for test_id in tqdm(range(len(dataset))):
|
||||
img1, img2, imageL_file = dataset[test_id]
|
||||
image1 = img1[None].cuda()
|
||||
image2 = img2[None].cuda()
|
||||
|
||||
padder = InputPadder(image1.shape, divis_by=32)
|
||||
image1, image2 = padder.pad(image1, image2)
|
||||
|
||||
with autocast(enabled=mixed_prec):
|
||||
start = time.time()
|
||||
disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True)
|
||||
end = time.time()
|
||||
runtime = end - start
|
||||
disp = padder.unpad(disp_pr).cpu().squeeze().numpy()
|
||||
disp[disp < 0] = 0
|
||||
disp[disp > 64] = 64
|
||||
|
||||
names = imageL_file.split("/")
|
||||
save_sub_path = os.path.join(save_path, "low_res_"+names[-3])
|
||||
makedirs(save_sub_path)
|
||||
|
||||
disp_path = os.path.join(save_sub_path, names[-2] + '.pfm')
|
||||
writePFM(disp_path, disp)
|
||||
|
||||
txt_path = os.path.join(save_sub_path, names[-2] + '.txt')
|
||||
with open(txt_path, 'wb') as time_file:
|
||||
time_file.write(StrToBytes('runtime ' + str(runtime)))
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test_kitti(model, save_path, iters=32, scale_iters=3, split='15', mixed_prec=False):
|
||||
""" Peform testing on the KITTI-2015 (test) split """
|
||||
model.eval()
|
||||
aug_params = {}
|
||||
save_path = os.path.join(save_path, "disp_0")
|
||||
makedirs(save_path)
|
||||
|
||||
test_dataset = datasets.KITTI(aug_params, split=split, image_set='testing', is_test=True)
|
||||
|
||||
runtime_sum = 0.0
|
||||
runtime_count = 0
|
||||
|
||||
for test_id in tqdm(range(len(test_dataset))):
|
||||
img1, img2, imageL_file = test_dataset[test_id]
|
||||
image1 = img1[None].cuda()
|
||||
image2 = img2[None].cuda()
|
||||
|
||||
padder = InputPadder(image1.shape, divis_by=32)
|
||||
image1, image2 = padder.pad(image1, image2)
|
||||
|
||||
with autocast(enabled=mixed_prec):
|
||||
start = time.time()
|
||||
disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True)
|
||||
end = time.time()
|
||||
runtime = end - start
|
||||
runtime_sum += runtime
|
||||
runtime_count += 1
|
||||
|
||||
disp = padder.unpad(disp_pr).cpu().squeeze().numpy()
|
||||
disp[disp < 0] = 0
|
||||
disp[disp > 240] = 240
|
||||
disp = np.uint16(disp*256)
|
||||
|
||||
name = imageL_file.split('/')[-1]
|
||||
path = os.path.join(save_path, name)
|
||||
cv2.imwrite(path, disp, [cv2.IMWRITE_PNG_COMPRESSION, 9])
|
||||
|
||||
print('The average runtime on Kitti test images is (you will need this for the submission): '
|
||||
+ str(runtime_sum / runtime_count) + " seconds")
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test_middlebury(model, save_path, iters=32, scale_iters=8, split='F', mixed_prec=False, method_name="DEFOM-Stereo"):
|
||||
""" Peform validation using the Middlebury-V3 dataset """
|
||||
model.eval()
|
||||
aug_params = {}
|
||||
test_dataset = datasets.Middlebury(aug_params, split=split, image_set='test', is_test=True)
|
||||
training_dataset = datasets.Middlebury(aug_params, split=split, image_set='training', is_test=True)
|
||||
dataset = test_dataset + training_dataset
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
for test_id in tqdm(range(len(dataset))):
|
||||
img1, img2, imageL_file = dataset[test_id]
|
||||
image1 = img1[None].cuda()
|
||||
image2 = img2[None].cuda()
|
||||
|
||||
padder = InputPadder(image1.shape, divis_by=32)
|
||||
image1, image2 = padder.pad(image1, image2)
|
||||
|
||||
with autocast(enabled=mixed_prec):
|
||||
start = time.time()
|
||||
disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True)
|
||||
end = time.time()
|
||||
runtime = end - start
|
||||
disp = padder.unpad(disp_pr).cpu().squeeze().numpy()
|
||||
disp[disp < 0] = 0
|
||||
disp[disp > 800] = 800
|
||||
|
||||
names = imageL_file.split("/")
|
||||
save_sub_path = os.path.join(save_path, names[-3], names[-2])
|
||||
makedirs(save_sub_path)
|
||||
|
||||
disp_path = os.path.join(save_sub_path, 'disp0' + method_name + '.pfm')
|
||||
writePFM(disp_path, disp)
|
||||
|
||||
txt_path = os.path.join(save_sub_path, 'time' + method_name + '.txt')
|
||||
with open(txt_path, 'wb') as time_file:
|
||||
time_file.write(StrToBytes(str(runtime)))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument('--restore_ckpt', help="restore checkpoint", default=None)
|
||||
parser.add_argument('--datasets', nargs='+', type=str, help="dataset for evaluation", default=["kitti12", "kitti15"],
|
||||
choices=["eth3d", "kitti12", "kitti15"] + [f"middlebury_{s}" for s in 'FHQ'])
|
||||
parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision')
|
||||
parser.add_argument('--valid_iters', type=int, default=32, help='number of disparity field updates during forward pass')
|
||||
parser.add_argument('--scale_iters', type=int, default=8, help="number of scaling updates to the disparity field in each forward pass.")
|
||||
parser.add_argument('--method_name', default="DEFOM-Stereo", help="the method to test")
|
||||
|
||||
# Architecure choices
|
||||
parser.add_argument('--dinov2_encoder', type=str, default='vits', choices=['vits', 'vitb', 'vitl', 'vitg'])
|
||||
parser.add_argument('--idepth_scale', type=float, default=0.5,
|
||||
help="the scale of inverse depth to initialize disparity")
|
||||
parser.add_argument('--hidden_dims', nargs='+', type=int, default=[128] * 3,
|
||||
help="hidden state and context dimensions")
|
||||
parser.add_argument('--corr_implementation', choices=["reg", "alt", "reg_cuda", "alt_cuda"], default="reg",
|
||||
help="correlation volume implementation")
|
||||
parser.add_argument('--shared_backbone', action='store_true',
|
||||
help="use a single backbone for the context and feature encoders")
|
||||
parser.add_argument('--corr_levels', type=int, default=2, help="number of levels in the correlation pyramid")
|
||||
parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid")
|
||||
parser.add_argument('--scale_list', type=float, nargs='+', default=[0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0],
|
||||
help='the list of scaling factors of disparity')
|
||||
parser.add_argument('--scale_corr_radius', type=int, default=2,
|
||||
help="width of the correlation pyramid for scaled disparity")
|
||||
|
||||
parser.add_argument('--n_downsample', type=int, default=2, choices=[2, 3],
|
||||
help="resolution of the disparity field (1/2^K)")
|
||||
parser.add_argument('--context_norm', type=str, default="batch", choices=['group', 'batch', 'instance', 'none'],
|
||||
help="normalization of context encoder")
|
||||
parser.add_argument('--n_gru_layers', type=int, default=3, help="number of hidden GRU levels")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
model = DEFOMStereo(args)
|
||||
|
||||
logging.basicConfig(level=logging.INFO,
|
||||
format='%(asctime)s %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s')
|
||||
|
||||
if args.restore_ckpt is not None:
|
||||
assert args.restore_ckpt.endswith(".pth")
|
||||
logging.info("Loading checkpoint...")
|
||||
checkpoint = torch.load(args.restore_ckpt, map_location='cuda')
|
||||
model.load_state_dict(checkpoint, strict=True)
|
||||
logging.info(f"Done loading checkpoint")
|
||||
|
||||
model.cuda()
|
||||
model.eval()
|
||||
|
||||
print(f"The model has {format(count_parameters(model)/1e6, '.2f')}M learnable parameters.")
|
||||
|
||||
# The CUDA implementations of the correlation volume prevent half-precision
|
||||
# rounding errors in the correlation lookup. This allows us to use mixed precision
|
||||
# in the entire forward pass, not just in the GRUs & feature extractors.
|
||||
use_mixed_precision = args.corr_implementation.endswith("_cuda")
|
||||
|
||||
if 'eth3d' in args.datasets:
|
||||
save_path = os.path.abspath(args.restore_ckpt).split('.')[0] + '_' + "eth3d"
|
||||
makedirs(save_path)
|
||||
test_eth3d(model, save_path, iters=args.valid_iters, scale_iters=args.scale_iters, mixed_prec=use_mixed_precision)
|
||||
|
||||
if 'kitti12' in args.datasets:
|
||||
save_path = os.path.abspath(args.restore_ckpt).split('.')[0] + '_' + "kitti12"
|
||||
makedirs(save_path)
|
||||
test_kitti(model, save_path, iters=args.valid_iters, scale_iters=args.scale_iters, mixed_prec=use_mixed_precision, split='12')
|
||||
|
||||
if 'kitti15' in args.datasets:
|
||||
save_path = os.path.abspath(args.restore_ckpt).split('.')[0] + '_' + "kitti15"
|
||||
makedirs(save_path)
|
||||
test_kitti(model, save_path, iters=args.valid_iters, scale_iters=args.scale_iters, mixed_prec=use_mixed_precision, split='15')
|
||||
|
||||
for s in 'FHQ':
|
||||
if f"middlebury_{s}" in args.datasets:
|
||||
save_path = os.path.abspath(args.restore_ckpt).split('.')[0] + '_' + f"middlebury_{s}"
|
||||
makedirs(save_path)
|
||||
test_middlebury(model, save_path, iters=args.valid_iters, scale_iters=args.scale_iters, split=s,
|
||||
method_name=args.method_name, mixed_prec=use_mixed_precision)
|
||||
|
||||
Reference in New Issue
Block a user