Spaces:
Sleeping
Sleeping
File size: 2,039 Bytes
2267636 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | """Shared DPT reassemble + RefineNet cascade.
`DPTRefineNetStack` owns the `scratch` reassemble layers and the four
`FeatureFusionBlock` refinenets used by every DPT decoder in this repo (the depth
heads in ``dpt_decoder.py`` and the saliency decoder in
``dpt_segmentation_decoder.py``). Decoders subclass it so the parameter names stay
flat (``scratch.*`` / ``refinenet{1..4}.*``) and existing checkpoints keep loading;
each subclass provides its own input projection and output head.
"""
import torch.nn as nn
from .blocks import FeatureFusionBlock, _make_scratch
class DPTRefineNetStack(nn.Module):
def __init__(self, features=256, use_bn=False, out_channels=(256, 512, 1024, 1024)):
super().__init__()
self.features = features
self.scratch = _make_scratch(list(out_channels), features)
self.refinenet4 = FeatureFusionBlock(features, nn.ReLU(inplace=False), bn=use_bn)
self.refinenet3 = FeatureFusionBlock(features, nn.ReLU(inplace=False), bn=use_bn)
self.refinenet2 = FeatureFusionBlock(features, nn.ReLU(inplace=False), bn=use_bn)
self.refinenet1 = FeatureFusionBlock(features, nn.ReLU(inplace=False), bn=use_bn)
def fuse(self, layers, keep_layer1_size=False):
"""Run the coarse-to-fine RefineNet cascade; returns the layer-1 feature map.
``keep_layer1_size=True`` stops the final block from doing its default 2x
upsample (used by the high-res depth decoder, which upsamples in its head).
"""
l1, l2, l3, l4 = layers
l1 = self.scratch.layer1_rn(l1)
l2 = self.scratch.layer2_rn(l2)
l3 = self.scratch.layer3_rn(l3)
l4 = self.scratch.layer4_rn(l4)
path = self.refinenet4(l4, size=l3.shape[2:])
path = self.refinenet3(path, l3, size=l2.shape[2:])
path = self.refinenet2(path, l2, size=l1.shape[2:])
if keep_layer1_size:
path = self.refinenet1(path, l1, size=l1.shape[2:])
else:
path = self.refinenet1(path, l1)
return path
|