Spaces:
Running
Running
| # from MIT licensed https://github.com/nemodleo/pytorch-histogram-matching | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class Histogram_Matching(nn.Module): | |
| def __init__(self, differentiable=False): | |
| super(Histogram_Matching, self).__init__() | |
| self.differentiable = differentiable | |
| def forward(self, dst, ref): | |
| # B C | |
| B, C, H, W = dst.size() | |
| # assertion | |
| assert dst.device == ref.device | |
| # [B*C 256] | |
| hist_dst = self.cal_hist(dst) | |
| hist_ref = self.cal_hist(ref) | |
| # [B*C 256] | |
| tables = self.cal_trans_batch(hist_dst, hist_ref) | |
| # [B C H W] | |
| rst = dst.clone() | |
| for b in range(B): | |
| for c in range(C): | |
| rst[b,c] = tables[b*c, (dst[b,c] * 255).long()] | |
| # [B C H W] | |
| rst /= 255. | |
| return rst | |
| def cal_hist(self, img): | |
| B, C, H, W = img.size() | |
| # [B*C 256] | |
| if self.differentiable: | |
| hists = self.soft_histc_batch(img * 255, bins=256, min=0, max=256, sigma=3*25) | |
| else: | |
| hists = torch.stack([torch.histc(img[b,c] * 255, bins=256, min=0, max=255) for b in range(B) for c in range(C)]) | |
| hists = hists.float() | |
| hists = F.normalize(hists, p=1) | |
| # BC 256 | |
| bc, n = hists.size() | |
| # [B*C 256 256] | |
| triu = torch.ones(bc, n, n, device=hists.device).triu() | |
| # [B*C 256] | |
| hists = torch.bmm(hists[:,None,:], triu)[:,0,:] | |
| return hists | |
| def soft_histc_batch(self, x, bins=256, min=0, max=256, sigma=3*25): | |
| # B C H W | |
| B, C, H, W = x.size() | |
| # [B*C H*W] | |
| x = x.view(B*C, -1) | |
| # 1 | |
| delta = float(max - min) / float(bins) | |
| # [256] | |
| centers = float(min) + delta * (torch.arange(bins, device=x.device, dtype=torch.bfloat16) + 0.5) | |
| # [B*C 1 H*W] | |
| x = torch.unsqueeze(x, 1) | |
| # [1 256 1] | |
| centers = centers[None,:,None] | |
| # [B*C 256 H*W] | |
| x = x - centers | |
| # [B*C 256 H*W] | |
| x = x.type(torch.bfloat16) | |
| # [B*C 256 H*W] | |
| x = torch.sigmoid(sigma * (x + delta/2)) - torch.sigmoid(sigma * (x - delta/2)) | |
| # [B*C 256] | |
| x = x.sum(dim=2) | |
| # [B*C 256] | |
| x = x.type(torch.float32) | |
| # prevent oom | |
| # torch.cuda.empty_cache() | |
| return x | |
| def cal_trans_batch(self, hist_dst, hist_ref): | |
| # [B*C 256 256] | |
| hist_dst = hist_dst[:,None,:].repeat(1,256,1) | |
| # [B*C 256 256] | |
| hist_ref = hist_ref[:,:,None].repeat(1,1,256) | |
| # [B*C 256 256] | |
| table = hist_dst - hist_ref | |
| # [B*C 256 256] | |
| table = torch.where(table>=0, 1., 0.) | |
| # [B*C 256] | |
| table = torch.sum(table, dim=1) - 1 | |
| # [B*C 256] | |
| table = torch.clamp(table, min=0, max=255) | |
| return table | |