import warnings warnings.filterwarnings('ignore') warnings.simplefilter('ignore') import torch, yaml, cv2, os, shutil, sys, glob import numpy as np np.random.seed(0) import matplotlib.pyplot as plt from tqdm import trange from PIL import Image from ultralytics.nn.tasks import attempt_load_weights from timm.utils import AverageMeter import matplotlib.pyplot as plt plt.rcParams["font.family"] = "Times New Roman" import seaborn as sns def get_activation(feat, backbone_idx=-1): def hook(model, inputs, outputs): if backbone_idx != -1: for _ in range(5 - len(outputs)): outputs.insert(0, None) feat.append(outputs[backbone_idx]) else: feat.append(outputs) return hook def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32): # Resize and pad image while meeting stride-multiple constraints shape = im.shape[:2] # current shape [height, width] if isinstance(new_shape, int): new_shape = (new_shape, new_shape) # Scale ratio (new / old) r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) if not scaleup: # only scale down, do not scale up (for better val mAP) r = min(r, 1.0) # Compute padding ratio = r, r # width, height ratios new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding if auto: # minimum rectangle dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding elif scaleFill: # stretch dw, dh = 0.0, 0.0 new_unpad = (new_shape[1], new_shape[0]) ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios dw /= 2 # divide padding into 2 sides dh /= 2 if shape[::-1] != new_unpad: # resize im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR) top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border return im, ratio, (dw, dh) def get_rectangle(data, thresh): h, w = data.shape all_sum = np.sum(data) for i in range(1, h // 2): selected_area = data[h // 2 - i:h // 2 + 1 + i, w // 2 - i:w // 2 + 1 + i] area_sum = np.sum(selected_area) if area_sum / all_sum > thresh: return i * 2 + 1, (i * 2 + 1) / h * (i * 2 + 1) / w return None def heatmap(data, camp='RdYlGn', figsize=(10, 10.75), ax=None, save_path=None): plt.figure(figsize=figsize, dpi=40) ax = sns.heatmap(data, xticklabels=False, yticklabels=False, cmap=camp, center=0, annot=False, ax=ax, cbar=True, annot_kws={"size": 24}, fmt='.2f') plt.tight_layout() plt.savefig(save_path) class yolov8_erf: feature, hooks = [], [] def __init__(self, weight, device, layer, dataset, num_images, save_path) -> None: device = torch.device(device) ckpt = torch.load(weight) model = attempt_load_weights(weight, device) model.info() for p in model.parameters(): p.requires_grad_(True) model.eval() optimizer = torch.optim.SGD(model.parameters(), lr=0, weight_decay=0) meter = AverageMeter() optimizer.zero_grad() if '-' in layer: layer_first, layer_second = layer.split('-') self.hooks.append(model.model[int(layer_first)].register_forward_hook(get_activation(self.feature, backbone_idx=int(layer_second)))) else: self.hooks.append(model.model[int(layer)].register_forward_hook(get_activation(self.feature))) self.__dict__.update(locals()) def get_input_grad(self, samples): _ = self.model(samples) outputs = self.feature[-1] self.feature.clear() out_size = outputs.size() central_point = torch.nn.functional.relu(outputs[:, :, out_size[2] // 2, out_size[3] // 2]).sum() grad = torch.autograd.grad(central_point, samples) grad = grad[0] grad = torch.nn.functional.relu(grad) aggregated = grad.sum((0, 1)) grad_map = aggregated.cpu().numpy() return grad_map def process(self): for image_path in os.listdir(self.dataset): if self.meter.count == self.num_images: break img = cv2.imread(f'{self.dataset}/{image_path}') img = letterbox(img, auto=False)[0] img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = np.float32(img) / 255.0 samples = torch.from_numpy(np.transpose(img, axes=[2, 0, 1])).unsqueeze(0).to(self.device) samples.requires_grad = True self.optimizer.zero_grad() contribution_scores = self.get_input_grad(samples) if np.isnan(np.sum(contribution_scores)): print('got NAN, next image') continue else: print(f'{self.meter.count}/{self.num_images} calculate....') self.meter.update(contribution_scores) # Set figure parameters large = 24; med = 24; small = 24 params = {'axes.titlesize': large, 'legend.fontsize': med, 'figure.figsize': (16, 10), 'axes.labelsize': med, 'xtick.labelsize': med, 'ytick.labelsize': med, 'figure.titlesize': large} plt.rcParams.update(params) plt.style.use('seaborn-whitegrid') sns.set_style("white") plt.rc('font', **{'family': 'Times New Roman'}) plt.rcParams['axes.unicode_minus'] = False data = self.meter.avg print(f'max value:{np.max(data):.3f} min value:{np.min(data):.3f}') data = np.log10(data + 1) # the scores differ in magnitude. take the logarithm for better readability data = data / np.max(data) # rescale to [0,1] for the comparability among models print('======================= the high-contribution area ratio =====================') for thresh in [0.2, 0.3, 0.5, 0.99]: side_length, area_ratio = get_rectangle(data, thresh) print('thresh, rectangle side length, area ratio: ', thresh, side_length, area_ratio) heatmap(data, save_path=self.save_path) def get_params(): params = { 'weight': 'yolov8n.pt', # 只需要指定权重即可 'device': 'cuda:0', 'layer': '10', # string 'dataset': '', 'num_images': 50, 'save_path': 'result.png' } return params if __name__ == '__main__': cfg = get_params() yolov8_erf(**cfg).process()