def remd_loss(X,Y, h=None, cos_d=True, splits= [3+64+64+128+128+256+256+256+512+512],return_mat=False):
d = X.size(1)
if d == 3:
X = utils.rgb_to_yuv_pc(X.transpose(0,1).contiguous().view(d,-1)).transpose(0,1)
Y = utils.rgb_to_yuv_pc(Y.transpose(0,1).contiguous().view(d,-1)).transpose(0,1)
else:
X = X.transpose(0,1).contiguous().view(d,-1).transpose(0,1)
Y = Y.transpose(0,1).contiguous().view(d,-1).transpose(0,1)
# Y = Y.transpose(0, 1)
#Relaxed EMD
CX_M = get_DMat(X,Y,1.,cos_d=True, splits=splits)
if return_mat:
return CX_M
if d==3:
CX_M = CX_M+get_DMat(X,Y,1.,cos_d=False, splits=splits)
m1,m1_inds = CX_M.min(1)
m2,m2_inds = CX_M.min(0)
if m1.mean() > m2.mean():
used_style_feats = Y[m1_inds,:]
else:
used_style_feats = Y
remd = torch.max(m1.mean(),m2.mean())
return remd, used_style_feats
def remd_loss_g(X,
Y,
GX,
GY,
h=1.0,
splits=[3 + 64 + 64 + 128 + 128 + 256 + 256 + 256 + 512 + 512
]):
d = X.size(1)
if d == 3:
X = utils.rgb_to_yuv_pc(X.transpose(0, 1).contiguous().view(
d, -1)).transpose(0, 1)
Y = utils.rgb_to_yuv_pc(Y.transpose(0, 1).contiguous().view(
d, -1)).transpose(0, 1)
GX = utils.rgb_to_yuv_pc(GX.transpose(0, 1).contiguous().view(
d, -1)).transpose(0, 1)
GY = utils.rgb_to_yuv_pc(GY.transpose(0, 1).contiguous().view(
d, -1)).transpose(0, 1)
else:
X = X.transpose(0, 1).contiguous().view(d, -1).transpose(0, 1)
Y = Y.transpose(0, 1).contiguous().view(d, -1).transpose(0, 1)
GX = GX.transpose(0, 1).contiguous().view(d, -1).transpose(0, 1)
GY = GY.transpose(0, 1).contiguous().view(d, -1).transpose(0, 1)
c1 = 10000.
c2 = 1.
CX_M = get_DMat(X, Y, 1., cos_d=True, splits=splits)
if d == 3:
CX_M = CX_M + get_DMat(X, Y, 1., cos_d=False, splits=splits)
CX_M_2 = get_DMat(GX, GY, 1., cos_d=True, splits=splits) + get_DMat(
GX, GY, 1., cos_d=False, splits=splits) #CX_M[i:,i:].clone()
for i in range(GX.size(0) - 1):
CX_M_2[(i + 1):, i] = CX_M_2[(i + 1):, i] * 1000.
CX_M_2[i, (i + 1):] = CX_M_2[i, (i + 1):] * 1000.
m1, m1_inds = CX_M.min(1)
m2, m2_inds = CX_M.min(0)
m2, min_inds = torch.topk(m2, m1.size(0), largest=False)
if m1.mean() > m2.mean():
used_style_feats = Y[m1_inds, :]
else:
used_style_feats = Y[min_inds, :]
m12, _ = CX_M_2.min(1)
m22, _ = CX_M_2.min(0)
used_style_feats = Y[m1_inds, :]
remd = torch.max(m1.mean() * h,
m2.mean()) + c2 * torch.max(m12.mean() * h, m22.mean())
return remd, used_style_feats