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