# Copyright (c) OpenMMLab. All rights reserved. import torch import torch.nn as nn import torch.nn.functional as F from mmseg.core import add_prefix from mmseg.ops import resize from mmseg.models import builder from mmseg.models.builder import SEGMENTORS from mmseg.models.segmentors.base import BaseSegmentor from mmseg.models.segmentors.encoder_decoder import EncoderDecoder class Sobel(nn.Module): def __init__(self,in_Channels): super(Sobel, self).__init__() self.sobel_filter_x = nn.Conv2d(in_channels=in_Channels, out_channels=in_Channels, kernel_size=3, stride=1, padding=1, bias=False) self.sobel_filter_x.weight.data[...] = torch.Tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]) self.sobel_filter_y = nn.Conv2d(in_channels=in_Channels, out_channels=in_Channels, kernel_size=3, stride=1, padding=1, bias=False) self.sobel_filter_y.weight.data[...] = torch.Tensor([[-1, -2, -1], [ 0, 0, 0], [ 1, 2, 1]]) def forward(self, x): with torch.no_grad(): sobel_x = self.sobel_filter_x(x) sobel_y = self.sobel_filter_y(x) sobel_intensity=torch.sqrt(sobel_x ** 2 + sobel_y ** 2) sobel_intensity_mean=torch.mean(sobel_intensity,dim=1,keepdim=True) sobel_intensity_mean_norm=(sobel_intensity_mean-sobel_intensity_mean.min())/(sobel_intensity_mean.max()-sobel_intensity_mean.min()) sobel_intensity_mean_norm[:,:,:2]=sobel_intensity_mean_norm[:,:,-2:]=sobel_intensity_mean_norm[:,:,:,:2]=sobel_intensity_mean_norm[:,:,:,-2:]=0 result = sobel_intensity_mean_norm return result @SEGMENTORS.register_module() class TemporalEncoderDecoder(EncoderDecoder): """Encoder Decoder segmentors. EncoderDecoder typically consists of backbone, neck, decode_head, auxiliary_head. Note that auxiliary_head is only used for deep supervision during training, which could be dumped during inference. The backbone should return plain embeddings. The neck can process these to make them suitable for the chosen heads. The heads perform the final processing that will return the output. """ def __init__(self, backbone, decode_head, neck=None, auxiliary_head=None, edge_head=None, train_cfg=None, test_cfg=None, pretrained=None, init_cfg=None, frozen_backbone=False): super(EncoderDecoder, self).__init__(init_cfg) if pretrained is not None: assert backbone.get('pretrained') is None, \ 'both backbone and segmentor set pretrained weight' backbone.pretrained = pretrained self.backbone = builder.build_backbone(backbone) backbone_params = sum(p.numel() for p in self.backbone.parameters() if p.requires_grad) print("backbone params:", backbone_params) if frozen_backbone: for param in self.backbone.parameters(): param.reauires_grad=False if neck is not None: self.neck = builder.build_neck(neck) self._init_decode_head(decode_head) self._init_auxiliary_head(auxiliary_head) self._init_edge_head(edge_head) self.train_cfg = train_cfg self.test_cfg = test_cfg assert self.with_decode_head def _init_edge_head(self, edge_head): """Initialize ``edge_head``""" if edge_head is not None: self.with_edge_head=True if isinstance(edge_head, list): self.edge_head = nn.ModuleList() for head_cfg in edge_head: self.edge_head.append(builder.build_head(head_cfg)) else: self.edge_head = builder.build_head(edge_head) else: self.with_edge_head=False def encode_decode(self, img, img_metas): """Encode images with backbone and decode into a semantic segmentation map of the same size as input.""" x = self.extract_feat(img) out = self._decode_head_forward_test(x, img_metas) size = img.shape[-2:] out = resize( input=out, size=size, mode='bilinear', align_corners=self.align_corners) return out def _edge_head_forward_train(self, x, img_metas, gt_semantic_seg): """Run forward function and calculate loss for edge head in training.""" losses = dict() if isinstance(self.edge_head, nn.ModuleList): for idx, eg_head in enumerate(self.edge_head): loss_eg = eg_head.forward_train(x, img_metas, gt_semantic_seg, self.train_cfg) losses.update(add_prefix(loss_eg, f'eg_{idx}')) else: loss_eg = self.edge_head.forward_train( x, img_metas, gt_semantic_seg, self.train_cfg) losses.update(add_prefix(loss_eg, 'eg')) return losses def forward_train(self, img, img_metas, gt_semantic_seg): """Forward function for training. Args: img (Tensor): Input images. img_metas (list[dict]): List of image info dict where each dict has: 'img_shape', 'scale_factor', 'flip', and may also contain 'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'. For details on the values of these keys see `mmseg/datasets/pipelines/formatting.py:Collect`. gt_semantic_seg (Tensor): Semantic segmentation masks used if the architecture supports semantic segmentation task. Returns: dict[str, Tensor]: a dictionary of loss components """ x = self.extract_feat(img) losses = dict() loss_decode = self._decode_head_forward_train(x, img_metas, gt_semantic_seg) losses.update(loss_decode) if self.with_auxiliary_head: loss_aux = self._auxiliary_head_forward_train( x, img_metas, gt_semantic_seg) losses.update(loss_aux) if self.with_edge_head: B,C,T,W,H=img.shape img_new=img.view(B,C*T,W,H) sobel=Sobel(in_Channels=C*T).to(img_new.device) edge_labels=sobel(img_new) loss_eg=self._edge_head_forward_train( x,img_metas,edge_labels ) losses.update(loss_eg) return losses def slide_inference(self, img, img_meta, rescale): """Inference by sliding-window with overlap. If h_crop > h_img or w_crop > w_img, the small patch will be used to decode without padding. """ h_stride, w_stride = self.test_cfg.stride h_crop, w_crop = self.test_cfg.crop_size img_size = img.size() batch_size = img_size[0] h_img = img_size[-2] w_img = img_size[-1] out_channels = self.out_channels h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1 w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1 preds = img.new_zeros((batch_size, out_channels, h_img, w_img)) count_mat = img.new_zeros((batch_size, 1, h_img, w_img)) for h_idx in range(h_grids): for w_idx in range(w_grids): y1 = h_idx * h_stride x1 = w_idx * w_stride y2 = min(y1 + h_crop, h_img) x2 = min(x1 + w_crop, w_img) y1 = max(y2 - h_crop, 0) x1 = max(x2 - w_crop, 0) if len(img_size) == 4: crop_img = img[:, :, y1:y2, x1:x2] elif len(img_size) == 5: crop_img = img[:, :, :, y1:y2, x1:x2] crop_seg_logit = self.encode_decode(crop_img, img_meta) preds += F.pad(crop_seg_logit, (int(x1), int(preds.shape[3] - x2), int(y1), int(preds.shape[2] - y2))) count_mat[:, :, y1:y2, x1:x2] += 1 assert (count_mat == 0).sum() == 0 if torch.onnx.is_in_onnx_export(): count_mat = torch.from_numpy( count_mat.cpu().detach().numpy()).to(device=img.device) preds = preds / count_mat if rescale: resize_shape = img_meta[0]['img_shape'][:2] preds = preds[:, :, :resize_shape[0], :resize_shape[1]] preds = resize( preds, size=img_meta[0]['ori_shape'][:2], mode='bilinear', align_corners=self.align_corners, warning=False) return preds def whole_inference(self, img, img_meta, rescale): """Inference with full image.""" seg_logit = self.encode_decode(img, img_meta) if rescale: # support dynamic shape for onnx if torch.onnx.is_in_onnx_export(): size = img.shape[-2:] else: # remove padding area resize_shape = img_meta[0]['img_shape'][:2] seg_logit = seg_logit[:, :, :resize_shape[0], :resize_shape[1]] size = img_meta[0]['ori_shape'][:2] seg_logit = resize( seg_logit, size=size, mode='bilinear', align_corners=self.align_corners, warning=False) return seg_logit def inference(self, img, img_meta, rescale): """Inference with slide/whole style. Args: img (Tensor): The input image of shape (N, 3, H, W). img_meta (dict): Image info dict where each dict has: 'img_shape', 'scale_factor', 'flip', and may also contain 'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'. For details on the values of these keys see `mmseg/datasets/pipelines/formatting.py:Collect`. rescale (bool): Whether rescale back to original shape. Returns: Tensor: The output segmentation map. """ assert self.test_cfg.mode in ['slide', 'whole'] ori_shape = img_meta[0]['ori_shape'] assert all(_['ori_shape'] == ori_shape for _ in img_meta) if self.test_cfg.mode == 'slide': seg_logit = self.slide_inference(img, img_meta, rescale) else: seg_logit = self.whole_inference(img, img_meta, rescale) if self.out_channels == 1: output = F.sigmoid(seg_logit) else: output = F.softmax(seg_logit, dim=1) flip = ( img_meta[0]["flip"] if "flip" in img_meta[0] else False ) ##### if flip key is not there d not apply it if flip: flip_direction = img_meta[0]["flip_direction"] assert flip_direction in ["horizontal", "vertical"] if flip_direction == "horizontal": output = output.flip(dims=(3,)) elif flip_direction == "vertical": output = output.flip(dims=(2,)) return output def simple_test(self, img, img_meta, rescale=True): """Simple test with single image.""" seg_logit = self.inference(img, img_meta, rescale) if self.out_channels == 1: seg_pred = (seg_logit > self.decode_head.threshold).to(seg_logit).squeeze(1) else: seg_pred=seg_logit if torch.onnx.is_in_onnx_export(): seg_pred = seg_pred.unsqueeze(0) return seg_pred seg_pred = seg_pred.cpu().numpy() # unravel batch dim seg_pred = list(seg_pred) return seg_pred