Добавлены пропсы конвейера и стереодвижки, задействованные в прогоне

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>
2026-08-01 13:12:07 +00:00
parent 0d32f32db0
commit 6e1a22ba8b
184 changed files with 17666 additions and 3 deletions
@@ -0,0 +1,142 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from core.update import BasicMultiUpdateBlock, ScaleBasicMultiUpdateBlock
+from core.extractor import BasicEncoder, MultiBasicEncoder, ResidualBlock, DefomEncoder
+from core.corr import CorrBlock1D, PytorchAlternateCorrBlock1D, CorrBlockFast1D, AlternateCorrBlock
+from core.utils.utils import coords_grid, upflow, get_danv2_io_size
+
+
+try:
+    autocast = torch.cuda.amp.autocast
+except:
+    # dummy autocast for PyTorch < 1.6
+    class autocast:
+        def __init__(self, enabled):
+            pass
+        def __enter__(self):
+            pass
+        def __exit__(self, *args):
+            pass
+
+
+class DEFOMStereo(nn.Module):
+    def __init__(self, args):
+        super(DEFOMStereo, self).__init__()
+        self.args = args
+
+        self.register_buffer('mean', torch.tensor([[0.485, 0.456, 0.406]])[..., None, None] * 255)
+        self.register_buffer('std', torch.tensor([[0.229, 0.224, 0.225]])[..., None, None] * 255)
+
+        self.defomencoder = DefomEncoder(args.dinov2_encoder, idepth_scale=args.idepth_scale)
+
+        context_dims = args.hidden_dims
+
+        self.fnet = BasicEncoder(self.defomencoder.out_dim, output_dim=256, norm_fn='instance', downsample=args.n_downsample)
+
+        self.context_zqr_convs = nn.ModuleList([nn.Conv2d(context_dims[i], args.hidden_dims[i]*3, 3, padding=3//2) for i in range(self.args.n_gru_layers)])
+
+        self.update_block = BasicMultiUpdateBlock(self.args, hidden_dims=args.hidden_dims)
+        self.scale_update_block = ScaleBasicMultiUpdateBlock(self.args, hidden_dims=args.hidden_dims)
+        
+        self.cnet = MultiBasicEncoder(self.defomencoder.out_dim, output_dim=[args.hidden_dims, context_dims],
+                                      norm_fn=args.context_norm, downsample=args.n_downsample)
+
+    def freeze_bn(self):
+        for m in self.modules():
+            if isinstance(m, nn.BatchNorm2d):
+                m.eval()
+
+    def initialize_coords(self, img):
+        """ Disparity is represented as difference between two vertical coordinate grids disp
+            = coords0[:, :1] - coords1[:, :1] """
+        N, _, H, W = img.shape
+
+        coords = coords_grid(N, H, W)[:, :1].to(img.device)
+
+        return coords
+
+    def upsample_flow(self, flow, mask):
+        """ Upsample disparity field [H/scale, W/scale, 1] -> [H, W, 1] using convex combination """
+        N, D, H, W = flow.shape
+        factor = 2 ** self.args.n_downsample
+        mask = mask.view(N, 1, 9, factor, factor, H, W)
+        mask = torch.softmax(mask, dim=2)
+
+        up_flow = F.unfold(factor * flow, [3, 3], padding=1)
+        up_flow = up_flow.view(N, D, 9, 1, 1, H, W)
+
+        up_flow = torch.sum(mask * up_flow, dim=2)
+        up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
+        return up_flow.reshape(N, D, factor * H, factor * W)
+
+    def forward(self, image1, image2, iters=12, scale_iters=3, test_mode=False):
+        """ Estimate optical flow between pair of frames """
+
+        image1 = ((image1 - self.mean)/self.std).contiguous()
+        image2 = ((image2 - self.mean)/self.std).contiguous()
+
+        bs, _, h, w = image1.shape
+        danv2_io_sizes = get_danv2_io_size(h, w, self.args.n_downsample)
+
+        # run the context network
+        with autocast(enabled=self.args.mixed_precision):
+            d_features, dfeat1, dfeat2, disp = self.defomencoder([image1, image2], danv2_io_sizes)
+
+            cnet_list = self.cnet(image1, d_features)
+            fmap1, fmap2 = self.fnet([image1, image2], [dfeat1, dfeat2])
+            net_list = [torch.tanh(x[0]) for x in cnet_list]
+            inp_list = [torch.relu(x[1]) for x in cnet_list]
+            # Rather than running the GRU's conv layers on the context features multiple times, we do it once at the beginning
+            inp_list = [list(conv(i).split(split_size=conv.out_channels//3, dim=1)) for i, conv in zip(inp_list, self.context_zqr_convs)]
+
+        coords = self.initialize_coords(net_list[0])
+
+        fmap1, fmap2 = fmap1.float(), fmap2.float()
+        disp = disp.float()
+        corr_fn = CorrBlock1D(fmap1, fmap2, coords, radius=self.args.corr_radius, num_levels=self.args.corr_levels,
+                              scale_list=self.args.scale_list, scale_corr_radius=self.args.scale_corr_radius)
+
+        disp_predictions = []
+        for itr in range(iters):
+            disp = disp.detach()
+
+            if itr < scale_iters:
+                corr = corr_fn(disp, scaling=True)  # index correlation volume
+                with autocast(enabled=self.args.mixed_precision):
+                    net_list, up_mask, scale_disp = self.scale_update_block(net_list, inp_list, corr, disp,
+                                                                            iter32=self.args.n_gru_layers == 3,
+                                                                            iter16=self.args.n_gru_layers >= 2)
+
+                # F(t+1) = \Scale(t) x F(t)
+                disp = scale_disp * disp
+            else:
+                corr = corr_fn(disp, scaling=False)  # index correlation volume
+                with autocast(enabled=self.args.mixed_precision):
+                    net_list, up_mask, delta_disp = self.update_block(net_list, inp_list, corr, disp,
+                                                                      iter32=self.args.n_gru_layers == 3,
+                                                                      iter16=self.args.n_gru_layers >= 2)
+
+                    # To avoid unstability, we limit the disparity update within the searching range.
+                    delta_disp = torch.clip(delta_disp, min=-2**(self.args.corr_levels-1)*self.args.corr_radius,
+                                            max=2**(self.args.corr_levels-1)*self.args.corr_radius)
+
+                # F(t+1) = F(t) + \Delta(t)
+                disp = disp + delta_disp
+
+            # We do not need to upsample or output intermediate results in test_mode
+            if test_mode and itr < iters - 1:
+                continue
+
+            # upsample predictions
+            if up_mask is None:
+                disp_up = upflow(disp, factor=2 ** self.n_downsample)
+            else:
+                disp_up = self.upsample_flow(disp, up_mask)
+
+            disp_predictions.append(disp_up)
+
+        if test_mode:
+            return disp_up
+
+        return disp_predictions