|
import torch |
|
import torch.nn as nn |
|
|
|
from .utils import extract_patches, combine_patches, efficient_cdist, get_NNs_Dists |
|
|
|
def make_criteria(conf): |
|
if conf['type'] == 'PatchCoherentLoss': |
|
return PatchCoherentLoss(conf['patch_size'], stride=conf['stride'], loop=conf['loop'], coherent_alpha=conf['coherent_alpha']) |
|
elif conf['type'] == 'SWDLoss': |
|
raise NotImplementedError('SWDLoss is not implemented') |
|
else: |
|
raise ValueError('Invalid criteria: {}'.format(conf['criteria'])) |
|
|
|
class PatchCoherentLoss(torch.nn.Module): |
|
def __init__(self, patch_size=7, stride=1, loop=False, coherent_alpha=None, cache=False): |
|
super(PatchCoherentLoss, self).__init__() |
|
self.patch_size = patch_size |
|
self.stride = stride |
|
self.loop = loop |
|
self.coherent_alpha = coherent_alpha |
|
assert self.stride == 1, "Only support stride of 1" |
|
|
|
self.cache = cache |
|
if cache: |
|
self.cached_data = None |
|
|
|
def forward(self, X, Ys, dist_wrapper=None, ext=None, return_blended_results=False): |
|
"""For each patch in input X find its NN in target Y and sum the their distances""" |
|
assert X.shape[0] == 1, "Only support batch size of 1" |
|
dist_fn = lambda X, Y: dist_wrapper(efficient_cdist, X, Y) if dist_wrapper is not None else efficient_cdist(X, Y) |
|
|
|
x_patches = extract_patches(X, self.patch_size, self.stride, loop=self.loop) |
|
|
|
if not self.cache or self.cached_data is None: |
|
y_patches = [] |
|
for y in Ys: |
|
y_patches += [extract_patches(y, self.patch_size, self.stride, loop=False)] |
|
y_patches = torch.cat(y_patches, dim=1) |
|
self.cached_data = y_patches |
|
else: |
|
y_patches = self.cached_data |
|
|
|
nnf, dist = get_NNs_Dists(dist_fn, x_patches.squeeze(0), y_patches.squeeze(0), self.coherent_alpha) |
|
|
|
if return_blended_results: |
|
return combine_patches(X.shape, y_patches[:, nnf, :], self.patch_size, self.stride, loop=self.loop), dist.mean() |
|
else: |
|
return dist.mean() |
|
|
|
def clean_cache(self): |
|
self.cached_data = None |
