dasha_f
6e1a22ba8b
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>
90 lines
4.4 KiB
Python
90 lines
4.4 KiB
Python
import sys
|
|
sys.path.append('core')
|
|
|
|
import argparse
|
|
import glob
|
|
import numpy as np
|
|
import torch
|
|
from tqdm import tqdm
|
|
from pathlib import Path
|
|
from core.defom_stereo import DEFOMStereo
|
|
from utils.utils import InputPadder
|
|
from PIL import Image
|
|
from matplotlib import pyplot as plt
|
|
|
|
|
|
DEVICE = 'cuda'
|
|
|
|
def load_image(imfile):
|
|
img = np.array(Image.open(imfile)).astype(np.uint8)
|
|
img = torch.from_numpy(img).permute(2, 0, 1).float()
|
|
return img[None].to(DEVICE)
|
|
|
|
def demo(args):
|
|
model = DEFOMStereo(args)
|
|
checkpoint = torch.load(args.restore_ckpt, map_location='cuda')
|
|
if 'model' in checkpoint:
|
|
model.load_state_dict(checkpoint['model'])
|
|
else:
|
|
model.load_state_dict(checkpoint)
|
|
|
|
model.to(DEVICE)
|
|
model.eval()
|
|
|
|
output_directory = Path(args.output_directory)
|
|
output_directory.mkdir(exist_ok=True)
|
|
|
|
with torch.no_grad():
|
|
left_images = sorted(glob.glob(args.left_imgs, recursive=True))
|
|
right_images = sorted(glob.glob(args.right_imgs, recursive=True))
|
|
print(f"Found {len(left_images)} images. Saving files to {output_directory}/")
|
|
|
|
for (imfile1, imfile2) in tqdm(list(zip(left_images, right_images))):
|
|
image1 = load_image(imfile1)
|
|
image2 = load_image(imfile2)
|
|
|
|
padder = InputPadder(image1.shape, divis_by=32)
|
|
image1, image2 = padder.pad(image1, image2)
|
|
|
|
with torch.no_grad():
|
|
disp_pr = model(image1, image2, iters=args.valid_iters, scale_iters=args.scale_iters, test_mode=True)
|
|
disp_pr = padder.unpad(disp_pr).cpu().squeeze().numpy()
|
|
|
|
file_stem = imfile1.split('/')[-1].split('_')[0]+'_'+args.restore_ckpt.split('/')[-1][:-4]
|
|
if args.save_numpy:
|
|
np.save(output_directory / f"{file_stem}.npy", disp_pr)
|
|
plt.imsave(output_directory / f"{file_stem}.png", disp_pr, cmap='jet')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--restore_ckpt', help="restore checkpoint", required=True)
|
|
parser.add_argument('--save_numpy', action='store_true', help='save output as numpy arrays')
|
|
parser.add_argument('-l', '--left_imgs', help="path to all first (left) frames", default="demo/*_left.png")
|
|
parser.add_argument('-r', '--right_imgs', help="path to all second (right) frames", default="demo/*_right.png")
|
|
parser.add_argument('--output_directory', help="directory to save output", default="demo")
|
|
parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision')
|
|
parser.add_argument('--valid_iters', type=int, default=32, help='number of flow-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.")
|
|
|
|
# Architecture choices
|
|
parser.add_argument('--dinov2_encoder', type=str, default='vitl', 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()
|
|
|
|
demo(args)
|