Upload models/model.py

models/model.py

@@ -0,0 +1,196 @@
+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