add gaitedge training code

2022-07-17 13:47:50 +08:00
parent 4205c5f283
commit b183455eb8
17 changed files with 814 additions and 11 deletions
@@ -0,0 +1,135 @@
+import torch
+from kornia import morphology as morph
+import torch.optim as optim
+
+from ..base_model import BaseModel
+from .gaitgl import GaitGL
+from ..modules import SilhouetteCropAndResize
+from torchvision.transforms import Resize
+from utils import get_valid_args, get_attr_from, is_list_or_tuple
+import os.path as osp
+
+
+class Segmentation(BaseModel):
+
+    def forward(self, inputs):
+        ipts, labs, typs, vies, seqL = inputs
+        del seqL
+        # ratios = ipts[0]
+        rgbs = ipts[1]
+        sils = ipts[2]
+        # del ipts
+        n, s, c, h, w = rgbs.size()
+        rgbs = rgbs.view(n*s, c, h, w)
+        sils = sils.view(n*s, 1, h, w)
+        logi = self.Backbone(rgbs)  # [n*s, c, h, w]
+        logits = torch.sigmoid(logi)
+        pred = (logits > 0.5).float()  # [n*s, c, h, w]
+
+        retval = {
+            'training_feat': {
+                'bce': {'logits': logits, 'labels': sils}
+            },
+            'visual_summary': {
+                'image/sils': sils, 'image/logits': logits, "image/pred": pred
+            },
+            'inference_feat': {
+                'pred': pred, 'mask': sils
+            }
+        }
+        return retval
+
+
+class GaitEdge(GaitGL):
+
+    def build_network(self, model_cfg):
+        super(GaitEdge, self).build_network(model_cfg["GaitGL"])
+        self.Backbone = self.get_backbone(model_cfg['Segmentation'])
+        self.align = model_cfg['align']
+        self.CROP = SilhouetteCropAndResize()
+        self.resize = Resize((64, 44))
+        self.is_edge = model_cfg['edge']
+        self.seg_lr = model_cfg['seg_lr']
+
+    def finetune_parameters(self):
+        fine_tune_params = list()
+        others_params = list()
+        for name, p in self.named_parameters():
+            if not p.requires_grad:
+                continue
+            if 'Backbone' in name:
+                fine_tune_params.append(p)
+            else:
+                others_params.append(p)
+        return [{'params': fine_tune_params, 'lr': self.seg_lr}, {'params': others_params}]
+
+    def get_optimizer(self, optimizer_cfg):
+        self.msg_mgr.log_info(optimizer_cfg)
+        optimizer = get_attr_from([optim], optimizer_cfg['solver'])
+        valid_arg = get_valid_args(optimizer, optimizer_cfg, ['solver'])
+        optimizer = optimizer(self.finetune_parameters(), **valid_arg)
+        return optimizer
+
+    def resume_ckpt(self, restore_hint):
+        if is_list_or_tuple(restore_hint):
+            for restore_hint_i in restore_hint:
+                self.resume_ckpt(restore_hint_i)
+            return
+        if isinstance(restore_hint, int):
+            save_name = self.engine_cfg['save_name']
+            save_name = osp.join(
+                self.save_path, 'checkpoints/{}-{:0>5}.pt'.format(save_name, restore_hint))
+            self.iteration = restore_hint
+        elif isinstance(restore_hint, str):
+            save_name = restore_hint
+            self.iteration = 0
+        else:
+            raise ValueError(
+                "Error type for -Restore_Hint-, supported: int or string.")
+        self._load_ckpt(save_name)
+
+    def forward(self, inputs):
+        ipts, labs, _, _, seqL = inputs
+
+        ratios = ipts[0]
+        rgbs = ipts[1]
+        sils = ipts[2]
+        # if len(sils.size()) == 4:
+        #     sils = sils.unsqueeze(2)
+        n, s, c, h, w = rgbs.size()
+        rgbs = rgbs.view(n*s, c, h, w)
+        sils = sils.view(n*s, 1, h, w)
+        logis = self.Backbone(rgbs)  # [n, s, c, h, w]
+        logits = torch.sigmoid(logis)
+        mask = torch.round(logits).float()
+        if self.is_edge:
+            kernel_1 = torch.ones((3, 3)).cuda()
+            kernel_2 = torch.ones((3, 3)).cuda()
+
+            dilated_mask = (morph.dilation(sils, kernel_1).detach()
+                            ) > 0.5  # Dilation
+            eroded_mask = (morph.erosion(sils, kernel_2).detach()
+                           ) > 0.5   # Dilation
+            edge_mask = dilated_mask ^ eroded_mask
+
+            new_logits = edge_mask*logits+eroded_mask*sils
+            if self.align:
+                cropped_logits = self.CROP(
+                    new_logits, sils, ratios)
+            else:
+                cropped_logits = self.resize(new_logits)
+        else:
+            if self.align:
+                cropped_logits = self.CROP(
+                    logits, mask, ratios)
+            else:
+                cropped_logits = self.resize(logits)
+        _, c, H, W = cropped_logits.size()
+        cropped_logits = cropped_logits.view(n, s, H, W)
+        retval = super(GaitEdge, self).forward(
+            [[cropped_logits], labs, None, None, seqL])
+        retval['training_feat']['bce'] = {'logits': logits, 'labels': sils}
+        retval['visual_summary']['image/roi'] = cropped_logits.view(
+            n*s, 1, H, W)
+
+        return retval