def forward(self, src_input, tgt_input, accum_field=None,
_index=slice(None), **kwargs):
prev_level = None
for i, aligner in zip(reversed(range(len(self.list))[_index]),
reversed(self.list[_index])):
if isinstance(src_input, list) and isinstance(tgt_input, list):
src, tgt = src_input[i], tgt_input[i]
else:
src, tgt = downsample(i)(src_input), downsample(i)(tgt_input)
if prev_level is not None:
accum_field = (upsample(prev_level - i)
(accum_field.permute(0, 3, 1, 2))
.permute(0, 2, 3, 1))
src = gridsample_residual(src, accum_field,
padding_mode='border')
factor = 2 / src.shape[-1] # scale to [-1,1]
res_field = aligner(src, tgt, **kwargs) * factor
if accum_field is not None:
resampled = gridsample_residual(
accum_field.permute(0, 3, 1, 2), res_field,
padding_mode='border').permute(0, 2, 3, 1)
accum_field = res_field + resampled
else:
accum_field = res_field
prev_level = i
accum_field = (upsample(prev_level)
(accum_field.permute(0, 3, 1, 2))
.permute(0, 2, 3, 1))
return accum_field
def forward(self,
src_input,
tgt_input,
accum_field=None,
_index=slice(None),
**kwargs):
assert _index.step is None
_index = slice(max(1, _index.start or 0),
_index.stop) # Skip the lowest aligner level (speedup)
prev_level = None
for i, aligner in zip(reversed(range(len(self.list))[_index]),
reversed(self.list[_index])):
if isinstance(src_input, list) and isinstance(tgt_input, list):
src, tgt = src_input[i], tgt_input[i]
else:
src, tgt = downsample(i, type='max')(src_input), downsample(
i, type='max')(tgt_input)
if prev_level is not None:
accum_field = accum_field.up(mips=(prev_level - i))
src = accum_field.sample(src)
res_field = aligner(src, tgt, **kwargs)
res_field = res_field * 128 / src_input[0].shape[
-2] # Weird factor inherited from old architecture
if accum_field is not None:
accum_field = res_field.compose_with(accum_field)
else:
accum_field = res_field
prev_level = i
accum_field = accum_field.up(prev_level)
return accum_field
def forward(self, src, tgt, **kwargs):
src_encodings = []
tgt_encodings = []
for module in self.list:
src, tgt = module(src, tgt, **kwargs)
src_encodings.append(src)
tgt_encodings.append(tgt)
src, tgt = downsample()(src), downsample()(tgt)
return src_encodings, tgt_encodings
def forward(self, src_input, tgt_input, accum_field=None):
for i in reversed(range(self.height)):
if isinstance(src_input, list) and isinstance(tgt_input, list):
src, tgt = src_input[i], tgt_input[i]
else:
src, tgt = downsample(i)(src_input), downsample(i)(tgt_input)
if accum_field is not None:
accum_field = (upsample()(accum_field.permute(0, 3, 1,
2)).permute(
0, 2, 3, 1))
src = gridsample_residual(src,
accum_field,
padding_mode='border')
factor = 2 / src.shape[-1] # scale to [-1,1]
res_field = self.list[i](src, tgt) * factor
if accum_field is not None:
resampled = gridsample_residual(
accum_field.permute(0, 3, 1, 2),
res_field,
padding_mode='border').permute(0, 2, 3, 1)
accum_field = res_field + resampled
else:
accum_field = res_field
return accum_field