import torch
import torch.nn as nn
import torch.nn.functional as F
from models.network import HourGlass2, SpixelNet, ColorProbNet
from models.transformer2d import EncoderLayer, DecoderLayer, TransformerEncoder, TransformerDecoder
from models.position_encoding import build_position_encoding
from models import basic, clusterkit, anchor_gen
from collections import OrderedDict
from utils import util, cielab
class SpixelSeg(nn.Module):
def __init__(self, inChannel=1, outChannel=9, batchNorm=True):
super(SpixelSeg, self).__init__()
self.net = SpixelNet(inChannel=inChannel, outChannel=outChannel, batchNorm=batchNorm)
def get_trainable_params(self, lr=1.0):
#print('=> [optimizer] finetune backbone with smaller lr')
params = []
for name, param in self.named_parameters():
if 'xxx' in name:
params.append({'params': param, 'lr': lr})
else:
params.append({'params': param})
return params
def forward(self, input_grays):
pred_probs = self.net(input_grays)
return pred_probs
class AnchorColorProb(nn.Module):
def __init__(self, inChannel=1, outChannel=313, sp_size=16, d_model=64, use_dense_pos=True, spix_pos=False, learning_pos=False, \
random_hint=False, hint2regress=False, enhanced=False, use_mask=False, rank=0, colorLabeler=None):
super(AnchorColorProb, self).__init__()
self.sp_size = sp_size
self.spix_pos = spix_pos
self.use_token_mask = use_mask
self.hint2regress = hint2regress
self.segnet = SpixelSeg(inChannel=1, outChannel=9, batchNorm=True)
self.repnet = ColorProbNet(inChannel=inChannel, outChannel=64)
self.enhanced = enhanced
if self.enhanced:
self.enhanceNet = HourGlass2(inChannel=64+1, outChannel=2, resNum=3, normLayer=nn.BatchNorm2d)
## transformer architecture
self.n_vocab = 313
d_model, dim_feedforward, nhead = d_model, 4*d_model, 8
dropout, activation = 0.1, "relu"
n_enc_layers, n_dec_layers = 6, 6
enc_layer = EncoderLayer(d_model, nhead, dim_feedforward, dropout, activation, use_dense_pos)
self.wildpath = TransformerEncoder(enc_layer, n_enc_layers, use_dense_pos)
self.hintpath = TransformerEncoder(enc_layer, n_enc_layers, use_dense_pos)
if self.spix_pos:
n_pos_x, n_pos_y = 256, 256
else:
n_pos_x, n_pos_y = 256//sp_size, 16//sp_size
self.pos_enc = build_position_encoding(d_model//2, n_pos_x, n_pos_y, is_learned=False)
self.mid_word_prj = nn.Linear(d_model, self.n_vocab, bias=False)
if self.hint2regress:
self.trg_word_emb = nn.Linear(d_model+2+1, d_model, bias=False)
self.trg_word_prj = nn.Linear(d_model, 2, bias=False)
else:
self.trg_word_emb = nn.Linear(d_model+self.n_vocab+1, d_model, bias=False)
self.trg_word_prj = nn.Linear(d_model, self.n_vocab, bias=False)
self.colorLabeler = colorLabeler
anchor_mode = 'random' if random_hint else 'clustering'
self.anchorGen = anchor_gen.AnchorAnalysis(mode=anchor_mode, colorLabeler=self.colorLabeler)
self._reset_parameters()
def _reset_parameters(self):
for p in self.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
def load_and_froze_weight(self, checkpt_path):
data_dict = torch.load(checkpt_path, map_location=torch.device('cpu'))
'''
for param_tensor in data_dict['state_dict']:
print(param_tensor,'\t',data_dict['state_dict'][param_tensor].size())
'''
self.segnet.load_state_dict(data_dict['state_dict'])
for name, param in self.segnet.named_parameters():
param.requires_grad = False
self.segnet.eval()
def set_train(self):
## running mode only affect certain modules, e.g. Dropout, BN, etc.
self.repnet.train()
self.wildpath.train()
self.hintpath.train()
if self.enhanced:
self.enhanceNet.train()
def get_entry_mask(self, mask_tensor):
if mask_tensor is None:
return None
## flatten (N,1,H,W) to (N,HW)
return mask_tensor.flatten(1)
def forward(self, input_grays, input_colors, n_anchors=8, sampled_T=0):
'''
Notice: function was customized for inferece only
'''
affinity_map = self.segnet(input_grays)
pred_feats = self.repnet(input_grays)
if self.spix_pos:
full_pos_feats = self.pos_enc(pred_feats)
proxy_feats = torch.cat([pred_feats, input_colors, full_pos_feats], dim=1)
pooled_proxy_feats, conf_sum = basic.poolfeat(proxy_feats, affinity_map, self.sp_size, self.sp_size, True)
feat_tokens = pooled_proxy_feats[:,:64,:,:]
spix_colors = pooled_proxy_feats[:,64:66,:,:]
pos_feats = pooled_proxy_feats[:,66:,:,:]
else:
proxy_feats = torch.cat([pred_feats, input_colors], dim=1)
pooled_proxy_feats, conf_sum = basic.poolfeat(proxy_feats, affinity_map, self.sp_size, self.sp_size, True)
feat_tokens = pooled_proxy_feats[:,:64,:,:]
spix_colors = pooled_proxy_feats[:,64:,:,:]
pos_feats = self.pos_enc(feat_tokens)
token_labels = torch.max(self.colorLabeler.encode_ab2ind(spix_colors), dim=1, keepdim=True)[1]
spixel_sizes = basic.get_spixel_size(affinity_map, self.sp_size, self.sp_size)
all_one_map = torch.ones(spixel_sizes.shape, device=input_grays.device)
empty_entries = torch.where(spixel_sizes < 25/(self.sp_size**2), all_one_map, 1-all_one_map)
src_pad_mask = self.get_entry_mask(empty_entries) if self.use_token_mask else None
trg_pad_mask = src_pad_mask
## parallel prob
N,C,H,W = feat_tokens.shape
## (N,C,H,W) -> (HW,N,C)
src_pos_seq = pos_feats.flatten(2).permute(2, 0, 1)
src_seq = feat_tokens.flatten(2).permute(2, 0, 1)
## color prob branch
enc_out, _ = self.wildpath(src_seq, src_pos_seq, src_pad_mask)
pal_logit = self.mid_word_prj(enc_out)
pal_logit = pal_logit.permute(1, 2, 0).view(N,self.n_vocab,H,W)
## seed prob branch
## mask(N,1,H,W): sample anchors at clustering layers
color_feat = enc_out.permute(1, 2, 0).view(N,C,H,W)
hint_mask, cluster_mask = self.anchorGen(color_feat, n_anchors, spixel_sizes, use_sklearn_kmeans=False)
pred_prob = torch.softmax(pal_logit, dim=1)
color_feat2 = src_seq.permute(1, 2, 0).view(N,C,H,W)
#pred_prob, adj_matrix = self.anchorGen._detect_correlation(color_feat, pred_prob, hint_mask, thres=0.1)
if sampled_T < 0:
## GT anchor colors
sampled_spix_colors = spix_colors
elif sampled_T > 0:
top1_spix_colors = self.anchorGen._sample_anchor_colors(pred_prob, hint_mask, T=0)
top2_spix_colors = self.anchorGen._sample_anchor_colors(pred_prob, hint_mask, T=1)
top3_spix_colors = self.anchorGen._sample_anchor_colors(pred_prob, hint_mask, T=2)
## duplicate meta tensors
sampled_spix_colors = torch.cat((top1_spix_colors,top2_spix_colors,top3_spix_colors), dim=0)
N = 3*N
input_grays = input_grays.expand(N,-1,-1,-1)
hint_mask = hint_mask.expand(N,-1,-1,-1)
affinity_map = affinity_map.expand(N,-1,-1,-1)
src_seq = src_seq.expand(-1, N,-1)
src_pos_seq = src_pos_seq.expand(-1, N,-1)
else:
sampled_spix_colors = self.anchorGen._sample_anchor_colors(pred_prob, hint_mask, T=sampled_T)
## debug: controllable
if False:
hint_mask, sampled_spix_colors = basic.io_user_control(hint_mask, spix_colors, output=False)
sampled_token_labels = torch.max(self.colorLabeler.encode_ab2ind(sampled_spix_colors), dim=1, keepdim=True)[1]
## hint based prediction
## (N,C,H,W) -> (HW,N,C)
mask_seq = hint_mask.flatten(2).permute(2, 0, 1)
if self.hint2regress:
spix_colors_ = sampled_spix_colors
gt_seq = spix_colors_.flatten(2).permute(2, 0, 1)
hint_seq = self.trg_word_emb(torch.cat([src_seq, mask_seq * gt_seq, mask_seq], dim=2))
dec_out, _ = self.hintpath(hint_seq, src_pos_seq, src_pad_mask)
else:
token_labels_ = sampled_token_labels
label_map = F.one_hot(token_labels_, num_classes=313).squeeze(1).float()
label_seq = label_map.permute(0, 3, 1, 2).flatten(2).permute(2, 0, 1)
hint_seq = self.trg_word_emb(torch.cat([src_seq, mask_seq * label_seq, mask_seq], dim=2))
dec_out, _ = self.hintpath(hint_seq, src_pos_seq, src_pad_mask)
ref_logit = self.trg_word_prj(dec_out)
Ct = 2 if self.hint2regress else self.n_vocab
ref_logit = ref_logit.permute(1, 2, 0).view(N,Ct,H,W)
## pixelwise enhancement
pred_colors = None
if self.enhanced:
proc_feats = dec_out.permute(1, 2, 0).view(N,64,H,W)
full_feats = basic.upfeat(proc_feats, affinity_map, self.sp_size, self.sp_size)
pred_colors = self.enhanceNet(torch.cat((input_grays,full_feats), dim=1))
pred_colors = torch.tanh(pred_colors)
return pal_logit, ref_logit, pred_colors, affinity_map, spix_colors, hint_mask