|
- import warnings
-
- warnings.filterwarnings('ignore')
- warnings.simplefilter('ignore')
- import torch, yaml, cv2, os, shutil
- 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 DetectionModel as Model
- from ultralytics.utils.torch_utils import intersect_dicts
- from ultralytics.utils.ops import xywh2xyxy
- from pytorch_grad_cam import GradCAMPlusPlus, GradCAM, XGradCAM
- from pytorch_grad_cam.utils.image import show_cam_on_image
- from pytorch_grad_cam.activations_and_gradients import ActivationsAndGradients
-
-
- 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)
-
-
- class yolov8_heatmap:
- def __init__(self, weight, cfg, device, method, layer, backward_type, conf_threshold, ratio):
- device = torch.device(device)
- ckpt = torch.load(weight)
- model_names = ckpt['model'].names
- csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
- model = Model(cfg, ch=3, nc=len(model_names)).to(device)
- csd = intersect_dicts(csd, model.state_dict(), exclude=['anchor']) # intersect
- model.load_state_dict(csd, strict=False) # load
- model.eval()
- print(f'Transferred {len(csd)}/{len(model.state_dict())} items')
-
- target_layers = [eval(layer)]
- method = eval(method)
-
- colors = np.random.uniform(0, 255, size=(len(model_names), 3)).astype(np.int32)
- self.__dict__.update(locals())
-
- def post_process(self, result):
- logits_ = result[:, 4:]
- boxes_ = result[:, :4]
- sorted, indices = torch.sort(logits_.max(1)[0], descending=True)
- return torch.transpose(logits_[0], dim0=0, dim1=1)[indices[0]], torch.transpose(boxes_[0], dim0=0, dim1=1)[
- indices[0]], xywh2xyxy(torch.transpose(boxes_[0], dim0=0, dim1=1)[indices[0]]).cpu().detach().numpy()
-
- def draw_detections(self, box, color, name, img):
- xmin, ymin, xmax, ymax = list(map(int, list(box)))
- cv2.rectangle(img, (xmin, ymin), (xmax, ymax), tuple(int(x) for x in color), 2)
- cv2.putText(img, str(name), (xmin, ymin - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.8, tuple(int(x) for x in color), 2,
- lineType=cv2.LINE_AA)
- return img
-
- def __call__(self, img_path, save_path):
- # remove dir if exist
- if os.path.exists(save_path):
- shutil.rmtree(save_path)
- # make dir if not exist
- os.makedirs(save_path, exist_ok=True)
-
- # img process
- img = cv2.imread(img_path)
- img = letterbox(img)[0]
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
- img = np.float32(img) / 255.0
- tensor = torch.from_numpy(np.transpose(img, axes=[2, 0, 1])).unsqueeze(0).to(self.device)
-
- # init ActivationsAndGradients
- grads = ActivationsAndGradients(self.model, self.target_layers, reshape_transform=None)
-
- # get ActivationsAndResult
- result = grads(tensor)
- activations = grads.activations[0].cpu().detach().numpy()
-
- # postprocess to yolo output
- post_result, pre_post_boxes, post_boxes = self.post_process(result[0])
- for i in trange(int(post_result.size(0) * self.ratio)):
- if float(post_result[i].max()) < self.conf_threshold:
- break
-
- self.model.zero_grad()
- # get max probability for this prediction
- if self.backward_type == 'class' or self.backward_type == 'all':
- score = post_result[i].max()
- score.backward(retain_graph=True)
-
- if self.backward_type == 'box' or self.backward_type == 'all':
- for j in range(4):
- score = pre_post_boxes[i, j]
- score.backward(retain_graph=True)
-
- # process heatmap
- if self.backward_type == 'class':
- gradients = grads.gradients[0]
- elif self.backward_type == 'box':
- gradients = grads.gradients[0] + grads.gradients[1] + grads.gradients[2] + grads.gradients[3]
- else:
- gradients = grads.gradients[0] + grads.gradients[1] + grads.gradients[2] + grads.gradients[3] + \
- grads.gradients[4]
- b, k, u, v = gradients.size()
- weights = self.method.get_cam_weights(self.method, None, None, None, activations,
- gradients.detach().numpy())
- weights = weights.reshape((b, k, 1, 1))
- saliency_map = np.sum(weights * activations, axis=1)
- saliency_map = np.squeeze(np.maximum(saliency_map, 0))
- saliency_map = cv2.resize(saliency_map, (tensor.size(3), tensor.size(2)))
- saliency_map_min, saliency_map_max = saliency_map.min(), saliency_map.max()
- if (saliency_map_max - saliency_map_min) == 0:
- continue
- saliency_map = (saliency_map - saliency_map_min) / (saliency_map_max - saliency_map_min)
-
- # add heatmap and box to image
- cam_image = show_cam_on_image(img.copy(), saliency_map, use_rgb=True)
- "不想在图片中绘画出边界框和置信度,注释下面的一行代码即可"
- cam_image = self.draw_detections(post_boxes[i], self.colors[int(post_result[i, :].argmax())],
- f'{self.model_names[int(post_result[i, :].argmax())]} {float(post_result[i].max()):.2f}',
- cam_image)
- cam_image = Image.fromarray(cam_image)
- cam_image.save(f'{save_path}/{i}.png')
-
-
- def get_params():
- params = {
- 'weight': 'yolov8n.pt', # 训练出来的权重文件
- 'cfg': 'ultralytics/cfg/models/v8/yolov8n.yaml', # 训练权重对应的yaml配置文件
- 'device': 'cuda:0',
- 'method': 'GradCAM', # GradCAMPlusPlus, GradCAM, XGradCAM , 使用的热力图库文件不同的效果不一样可以多尝试
- 'layer': 'model.model[9]', # 想要检测的对应层
- 'backward_type': 'all', # class, box, all
- 'conf_threshold': 0.01, # 0.6 # 置信度阈值,有的时候你的进度条到一半就停止了就是因为没有高于此值的了
- 'ratio': 0.02 # 0.02-0.1
- }
- return params
-
-
- if __name__ == '__main__':
- model = yolov8_heatmap(**get_params())
- model(r'ultralytics/assets/bus.jpg', 'result') # 第一个是检测的文件, 第二个是保存的路径
|