Upload 9 files

tools/grad_cam_CNN.py

@@ -0,0 +1,72 @@
+import os
+import sys
+sys.path.append('.')
+
+import matplotlib.pyplot as plt
+from utils import GradCAM, show_cam_on_image, center_crop_img
+
+import argparse
+from utils.config import Config
+from train import *
+
+def get_args():
+    parser = argparse.ArgumentParser('description=Change detection of remote sensing images')
+    parser.add_argument("-c", "--config", type=str, default="configs\cdxformer.py")
+    parser.add_argument("--output_dir", default=None)
+    parser.add_argument("--layer", default=None)
+    return parser.parse_args()
+
+def main():
+    args = get_args()
+
+    if args.layer == None:
+        raise NameError("Please ensure the parameter '--layer' is not None!\n e.g. --layer=model.net.decoderhead.LHBlock2.mlp_l")
+    
+    cfg = Config.fromfile(args.config)
+
+    model = myTrain.load_from_checkpoint(cfg.test_ckpt_path, cfg = cfg)
+    model = model.to('cuda')
+
+    test_loader = build_dataloader(cfg.dataset_config, mode='test')
+
+    if args.output_dir:
+        base_dir = args.output_dir
+    else:
+        base_dir = os.path.dirname(cfg.test_ckpt_path)
+    gradcam_output_dir = os.path.join(base_dir, "grad_cam", args.layer) 
+    if os.path.exists(gradcam_output_dir):
+        raise NameError("Please ensure gradcam_output_dir does not exist!")
+    
+    os.makedirs(gradcam_output_dir)
+
+    for input in tqdm(test_loader):
+        target_layers = [eval(args.layer)] # name of the network layer
+        mask, img_id =  input[2].cuda(), input[3]
+
+        cam = GradCAM(cfg, model=model.net, target_layers=target_layers, use_cuda=True)
+        target_category = 1  # tabby, tabby cat
+
+        grayscale_cam_all = cam(input_tensor=(input[0], input[1]), target_category=target_category)
+        
+        for i in range(grayscale_cam_all.shape[0]):
+            grayscale_cam = grayscale_cam_all[i, :]
+            visualization = show_cam_on_image(0,
+                                            grayscale_cam,
+                                            use_rgb=True)
+            fig = plt.figure()
+            ax = fig.add_subplot(111)
+            ax.imshow(visualization)
+            # ax = fig.add_subplot(122)
+            # ax.imshow(mask[i].cpu().numpy())
+            ax.set_xticks([])
+            ax.set_yticks([])
+            ax.spines['top'].set_visible(False)
+            ax.spines['right'].set_visible(False)
+            ax.spines['bottom'].set_visible(False)
+            ax.spines['left'].set_visible(False)
+            plt.savefig(os.path.join(gradcam_output_dir, '{}.png'.format(img_id[i])))
+            plt.close()
+
+
+if __name__ == '__main__':
+    main()

tools/grad_cam_transformer.py

@@ -0,0 +1,95 @@
+import os
+import sys
+sys.path.append('.')
+
+import matplotlib.pyplot as plt
+from utilss import GradCAM, show_cam_on_image, center_crop_img
+import math
+import argparse
+from utils.config import Config
+from train import *
+
+def get_args():
+    # input x: B, L, C
+    # if not, please adjust the order
+    parser = argparse.ArgumentParser('description=Change detection of remote sensing images')
+    parser.add_argument("-c", "--config", type=str, default="configs/cdmask.py")
+    parser.add_argument("--output_dir", default=None)
+    parser.add_argument("--layer", default=None)
+    parser.add_argument("--imgsize", default=256)
+    return parser.parse_args()
+
+class ResizeTransform:
+    def __init__(self, im_h: int, im_w: int):
+        self.height = im_h
+        self.width = im_w
+
+    def __call__(self, x):
+        # input x: B, L, C
+        result = x.reshape(x.size(0),
+                           self.height,
+                           self.width,
+                           x.size(2))
+
+        # Bring the channels to the first dimension,
+        # like in CNNs.
+        # [batch_size, H, W, C] -> [batch, C, H, W]
+        result = result.permute(0, 3, 1, 2)
+
+        return result
+
+def main():
+    args = get_args()
+
+    if args.layer == None:
+        raise NameError("Please ensure the parameter '--layer' is not None!\n e.g. --layer=model.net.decoderhead.LHBlock2.mlp_l")
+        
+    cfg = Config.fromfile(args.config)
+
+    model = myTrain.load_from_checkpoint(cfg.test_ckpt_path, cfg = cfg)
+    model = model.to('cuda')
+
+    test_loader = build_dataloader(cfg.dataset_config, mode='test')
+
+    if args.output_dir:
+        base_dir = args.output_dir
+    else:
+        base_dir = os.path.dirname(cfg.test_ckpt_path)
+    gradcam_output_dir = os.path.join(base_dir, "grad_cam", args.layer) 
+    if os.path.exists(gradcam_output_dir):
+        raise NameError("Please ensure gradcam_output_dir does not exist!")
+    
+    os.makedirs(gradcam_output_dir)
+
+    for input in tqdm(test_loader):
+        target_layers = [eval(args.layer)] # name of the network layer
+        mask, img_id =  input[2].cuda(), input[3]
+
+        cam = GradCAM(cfg, model=model.net, target_layers=target_layers, use_cuda=True, 
+                      reshape_transform=ResizeTransform(im_h=args.imgsize, im_w=args.imgsize))
+        target_category = 1  # tabby, tabby cat
+
+        grayscale_cam_all = cam(input_tensor=(input[0], input[1]), target_category=target_category)
+
+        for i in range(grayscale_cam_all.shape[0]):
+            grayscale_cam = grayscale_cam_all[i, :]
+            visualization = show_cam_on_image(0,
+                                            grayscale_cam,
+                                            use_rgb=True)
+            fig = plt.figure()
+            ax = fig.add_subplot(111)
+            ax.imshow(visualization)
+            # ax = fig.add_subplot(122)
+            # ax.imshow(mask[i].cpu().numpy())
+            ax.set_xticks([])
+            ax.set_yticks([])
+            ax.spines['top'].set_visible(False)
+            ax.spines['right'].set_visible(False)
+            ax.spines['bottom'].set_visible(False)
+            ax.spines['left'].set_visible(False)
+            plt.savefig(os.path.join(gradcam_output_dir, '{}.png'.format(img_id[i])))
+            plt.close()
+
+
+if __name__ == '__main__':
+    main()

tools/mask_convert.py

@@ -0,0 +1,103 @@
+import numpy as np
+import argparse
+import glob
+import os
+import sys
+import torch
+import cv2
+import random
+import time
+import multiprocessing.pool as mpp
+import multiprocessing as mp
+SEED = 66
+
+def seed_everything(seed):
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = True
+def label2rgb(mask, mask_pred):
+    real_1 = (mask == 1)
+    real_0 = (mask == 0)
+    pred_1 = (mask_pred == 1)
+    pred_0 = (mask_pred == 0)
+
+    TP = np.logical_and(real_1, pred_1)
+    TN = np.logical_and(real_0, pred_0)
+    FN = np.logical_and(real_1, pred_0)
+    FP = np.logical_and(real_0, pred_1)
+
+    mask_TP = TP[np.newaxis, :, :]
+    mask_TN = TN[np.newaxis, :, :]
+    mask_FN = FN[np.newaxis, :, :]
+    mask_FP = FP[np.newaxis, :, :]
+
+    h, w = mask.shape[0], mask.shape[1]
+    mask_rgb = np.zeros(shape=(h, w, 3), dtype=np.uint8)
+    mask_rgb[np.all(mask_TP, axis=0)] = [255, 255, 255] # TP
+    mask_rgb[np.all(mask_TN, axis=0)] = [0, 0, 0] # TN
+    mask_rgb[np.all(mask_FN, axis=0)] = [0, 255, 0] # FN
+    mask_rgb[np.all(mask_FP, axis=0)] = [255, 0, 0] # FP
+
+    return mask_rgb
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", default="Vaihingen")
+    parser.add_argument("--mask-dir", default="data/Test/masks")
+    parser.add_argument("--output-mask-dir", default="data/Test/masks_rgb")
+    return parser.parse_args()
+
+def mask_save(inp):
+    (mask, mask_pred, masks_output_dir, file_name) = inp
+    out_mask_path = os.path.join(masks_output_dir, "{}.png".format(file_name))
+
+    label = label2rgb(mask.copy(), mask_pred.copy())
+
+    rgb_label = cv2.cvtColor(label, cv2.COLOR_BGR2RGB)
+    cv2.imwrite(out_mask_path, rgb_label)
+
+# def get_rgb(inp):
+#     (mask_path, masks_output_dir,dataset) = inp
+#     mask_filename = os.path.splitext(os.path.basename(mask_path))[0]
+#     mask_bgr = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
+#     mask = cv2.cvtColor(mask_bgr, cv2.COLOR_BGR2RGB)
+#     if dataset == "LoveDA":
+#         rgb_label = loveda_label2rgb(mask.copy())
+#     elif dataset == "Vaihingen":
+#         rgb_label = vaihingen_label2rgb(mask.copy())
+#     elif dataset == "Potsdam":
+#         rgb_label = potsdam_label2rgb(mask.copy())
+#     elif dataset == "uavid":
+#         rgb_label = uavid_label2rgb(mask.copy())
+#     else: return
+#     #rgb_label = cv2.cvtColor(rgb_label, cv2.COLOR_RGB2BGR)
+
+#     out_mask_path_rgb = os.path.join(masks_output_dir, "{}.png".format(mask_filename))
+#     rgb_label = cv2.cvtColor(rgb_label, cv2.COLOR_BGR2RGB)
+#     cv2.imwrite(out_mask_path_rgb, rgb_label)
+
+# if __name__ == '__main__':
+#     base_path = "/home/xwma/lrr/rssegmentation/"
+#     args = parse_args()
+#     dataset = args.dataset
+
+#     seed_everything(SEED)
+#     masks_dir = args.mask_dir
+#     masks_output_dir = args.output_mask_dir
+#     masks_dir = base_path + masks_dir
+#     masks_output_dir = base_path + masks_output_dir
+
+#     mask_paths = glob.glob(os.path.join(masks_dir, "*.png"))
+#     inp = [(mask_path, masks_output_dir, dataset) for mask_path in mask_paths]
+#     if not os.path.exists(masks_output_dir):
+#         os.makedirs(masks_output_dir)
+
+#     t0 = time.time()
+#     mpp.Pool(processes=mp.cpu_count()).map(get_rgb, inp)
+#     t1 = time.time()
+#     split_time = t1 - t0
+#     print('images spliting spends: {} s'.format(split_time))

tools/params_flops.py

@@ -0,0 +1,55 @@
+import sys
+import torch
+sys.path.append('.')
+from train import *
+from fvcore.nn import FlopCountAnalysis, flop_count_table, flop_count, parameter_count
+from rscd.models.backbones.lamba_util.csms6s import flops_selective_scan_fn, flops_selective_scan_ref, selective_scan_flop_jit
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='count params and flops')
+    parser.add_argument("-c", "--config", type=str, default="configs/cdlama.py")
+    parser.add_argument("--size", type=int, default=256)
+    args = parser.parse_args()
+    return args
+
+def flops_mamba(model, shape=(3, 224, 224)):
+    # shape = self.__input_shape__[1:]
+    supported_ops = {
+        "aten::silu": None,  # as relu is in _IGNORED_OPS
+        "aten::neg": None,  # as relu is in _IGNORED_OPS
+        "aten::exp": None,  # as relu is in _IGNORED_OPS
+        "aten::flip": None,  # as permute is in _IGNORED_OPS
+        # "prim::PythonOp.CrossScan": None,
+        # "prim::PythonOp.CrossMerge": None,
+        "prim::PythonOp.SelectiveScanCuda": selective_scan_flop_jit,
+        "prim::PythonOp.SelectiveScanMamba": selective_scan_flop_jit,
+        "prim::PythonOp.SelectiveScanOflex": selective_scan_flop_jit,
+        "prim::PythonOp.SelectiveScanCore": selective_scan_flop_jit,
+        "prim::PythonOp.SelectiveScanNRow": selective_scan_flop_jit,
+    }
+
+    model.cuda().eval()
+
+    input1 = torch.randn((1, *shape), device=next(model.parameters()).device)
+    input2 = torch.randn((1, *shape), device=next(model.parameters()).device)
+    params = parameter_count(model)[""]
+    Gflops, unsupported = flop_count(model=model, inputs=(input1,input2), supported_ops=supported_ops)
+
+    del model, input1, input2
+    # return sum(Gflops.values()) * 1e9
+    return f"params {params / 1e6} GFLOPs {sum(Gflops.values())}"
+
+if __name__ == "__main__":
+    args = parse_args()
+    cfg = Config.fromfile(args.config)
+    net = myTrain(cfg).net.cuda()
+
+    size = args.size
+    input = torch.rand((1, 3, size, size)).cuda()
+    
+    net.eval()
+    flops = FlopCountAnalysis(net, (input, input))
+    print(flop_count_table(flops, max_depth = 2))
+
+    print(flops_mamba(net, (3, size, size)))

tools/utilss.py

@@ -0,0 +1,249 @@
+import cv2
+import numpy as np
+from torch.nn import functional as F
+import torch
+
+class ActivationsAndGradients:
+    """ Class for extracting activations and
+    registering gradients from targeted intermediate layers """
+
+    def __init__(self, model, target_layers, reshape_transform):
+        self.model = model
+        self.gradients = []
+        self.activations = []
+        self.reshape_transform = reshape_transform
+        self.handles = []
+        for target_layer in target_layers:
+            self.handles.append(
+                target_layer.register_forward_hook(
+                    self.save_activation))
+            # Backward compatibility with older pytorch versions:
+            if hasattr(target_layer, 'register_full_backward_hook'):
+                self.handles.append(
+                    target_layer.register_full_backward_hook(
+                        self.save_gradient))
+            else:
+                self.handles.append(
+                    target_layer.register_backward_hook(
+                        self.save_gradient))
+
+    def save_activation(self, module, input, output):
+        activation = output
+        if self.reshape_transform is not None:
+            activation = self.reshape_transform(activation)
+        self.activations.append(activation.cpu().detach())
+
+    def save_gradient(self, module, grad_input, grad_output):
+        # Gradients are computed in reverse order
+        grad = grad_output[0]
+        if self.reshape_transform is not None:
+            grad = self.reshape_transform(grad)
+        self.gradients = [grad.cpu().detach()] + self.gradients
+
+    def __call__(self, x, y):
+        self.gradients = []
+        self.activations = []
+        return self.model(x, y)
+
+    def release(self):
+        for handle in self.handles:
+            handle.remove()
+
+
+class GradCAM:
+    def __init__(self,
+                 cfg,
+                 model,
+                 target_layers,
+                 reshape_transform=None,
+                 use_cuda=False):
+        self.cfg = cfg
+        self.model = model.eval()
+        self.target_layers = target_layers
+        self.reshape_transform = reshape_transform
+        self.cuda = use_cuda
+        if self.cuda:
+            self.model = model.cuda()
+        self.activations_and_grads = ActivationsAndGradients(
+            self.model, target_layers, reshape_transform)
+
+    """ Get a vector of weights for every channel in the target layer.
+        Methods that return weights channels,
+        will typically need to only implement this function. """
+
+    @staticmethod
+    def get_cam_weights(grads):
+        return np.mean(grads, axis=(2, 3), keepdims=True)
+
+    @staticmethod
+    def get_loss(output, target_category):
+        loss = 0
+        for i in range(len(target_category)):
+            loss = loss + output[i]
+        return loss
+
+    def get_cam_image(self, activations, grads):
+        weights = self.get_cam_weights(grads)
+        weighted_activations = weights * activations
+        cam = weighted_activations.sum(axis=1)
+
+        return cam
+
+    @staticmethod
+    def get_target_width_height(input_tensor):
+        width, height = input_tensor.size(-1), input_tensor.size(-2)
+        return width, height
+
+    def compute_cam_per_layer(self, input_tensor):
+        activations_list = [a.cpu().data.numpy()
+                            for a in self.activations_and_grads.activations]
+        grads_list = [g.cpu().data.numpy()
+                      for g in self.activations_and_grads.gradients]
+        target_size = self.get_target_width_height(input_tensor)
+
+        cam_per_target_layer = []
+        # Loop over the saliency image from every layer
+
+        for layer_activations, layer_grads in zip(activations_list, grads_list):
+            cam = self.get_cam_image(layer_activations, layer_grads)
+            cam[cam < 0] = 0  # works like mute the min-max scale in the function of scale_cam_image
+            scaled = self.scale_cam_image(cam, target_size)
+            cam_per_target_layer.append(scaled[:, None, :])
+
+        return cam_per_target_layer
+
+    def aggregate_multi_layers(self, cam_per_target_layer):
+        cam_per_target_layer = np.concatenate(cam_per_target_layer, axis=1)
+        cam_per_target_layer = np.maximum(cam_per_target_layer, 0)
+        result = np.mean(cam_per_target_layer, axis=1)
+        return self.scale_cam_image(result)
+
+    @staticmethod
+    def scale_cam_image(cam, target_size=None):
+        result = []
+        for img in cam:
+            img = img - np.min(img)
+            img = img / (1e-7 + np.max(img))
+            if target_size is not None:
+                img = cv2.resize(img, target_size)
+            result.append(img)
+        result = np.float32(result)
+
+        return result
+
+    def __call__(self, input_tensor, target_category=None):
+        x, y = input_tensor
+        if self.cuda:
+            x = x.cuda()
+            y = y.cuda()
+
+        # 正向传播得到网络输出logits(未经过softmax)
+        if self.cfg.net == 'cdmask':
+            o, outputs = self.activations_and_grads(x, y)
+            mask_cls_results = outputs["pred_logits"]
+            mask_pred_results = outputs["pred_masks"]
+            mask_pred_results = F.interpolate(
+                mask_pred_results,
+                scale_factor=(4,4),
+                mode="bilinear",
+                align_corners=False,
+            )
+            mask_cls = F.softmax(mask_cls_results, dim=-1)[...,1:]
+            mask_pred = mask_pred_results.sigmoid()  
+            output = torch.einsum("bqc,bqhw->bchw", mask_cls, mask_pred)
+        else:
+            output = self.activations_and_grads(x, y)
+
+        if isinstance(target_category, int):
+            target_category = [target_category] * x.size(0)
+
+        if target_category is None:
+            target_category = np.argmax(output.cpu().data.numpy(), axis=-1)
+            print(f"category id: {target_category}")
+        else:
+            assert (len(target_category) == x.size(0))
+
+        self.model.zero_grad()
+        loss = self.get_loss(output, target_category).sum()
+        loss.backward(retain_graph=True)
+
+        # In most of the saliency attribution papers, the saliency is
+        # computed with a single target layer.
+        # Commonly it is the last convolutional layer.
+        # Here we support passing a list with multiple target layers.
+        # It will compute the saliency image for every image,
+        # and then aggregate them (with a default mean aggregation).
+        # This gives you more flexibility in case you just want to
+        # use all conv layers for example, all Batchnorm layers,
+        # or something else.
+        cam_per_layer = self.compute_cam_per_layer(x)
+        return self.aggregate_multi_layers(cam_per_layer)
+
+    def __del__(self):
+        self.activations_and_grads.release()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        self.activations_and_grads.release()
+        if isinstance(exc_value, IndexError):
+            # Handle IndexError here...
+            print(
+                f"An exception occurred in CAM with block: {exc_type}. Message: {exc_value}")
+            return True
+
+
+def show_cam_on_image(img: np.ndarray,
+                      mask: np.ndarray,
+                      use_rgb: bool = False,
+                      colormap: int = cv2.COLORMAP_JET) -> np.ndarray:
+    """ This function overlays the cam mask on the image as an heatmap.
+    By default the heatmap is in BGR format.
+
+    :param img: The base image in RGB or BGR format.
+    :param mask: The cam mask.
+    :param use_rgb: Whether to use an RGB or BGR heatmap, this should be set to True if 'img' is in RGB format.
+    :param colormap: The OpenCV colormap to be used.
+    :returns: The default image with the cam overlay.
+    """
+
+    heatmap = cv2.applyColorMap(np.uint8(255 * mask), colormap)
+    if use_rgb:
+        heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
+    heatmap = np.float32(heatmap) / 255
+
+    if np.max(img) > 1:
+        raise Exception(
+            "The input image should np.float32 in the range [0, 1]")
+
+    cam = heatmap + img
+    cam = cam / np.max(cam)
+    return np.uint8(255 * cam)
+
+
+def center_crop_img(img: np.ndarray, size: int):
+    h, w, c = img.shape
+
+    if w == h == size:
+        return img
+
+    if w < h:
+        ratio = size / w
+        new_w = size
+        new_h = int(h * ratio)
+    else:
+        ratio = size / h
+        new_h = size
+        new_w = int(w * ratio)
+
+    img = cv2.resize(img, dsize=(new_w, new_h))
+
+    if new_w == size:
+        h = (new_h - size) // 2
+        img = img[h: h+size]
+    else:
+        w = (new_w - size) // 2
+        img = img[:, w: w+size]
+
+    return img