import numpy as np import cv2 import os import sys import glob from ppdet.utils.logger import setup_logger import copy logger = setup_logger('ppdet_cam') import paddle from ppdet.engine import Trainer def get_test_images(infer_dir, infer_img): """ Get image path list in TEST mode """ assert infer_img is not None or infer_dir is not None, \ "--infer_img or --infer_dir should be set" assert infer_img is None or os.path.isfile(infer_img), \ "{} is not a file".format(infer_img) assert infer_dir is None or os.path.isdir(infer_dir), \ "{} is not a directory".format(infer_dir) # infer_img has a higher priority if infer_img and os.path.isfile(infer_img): return [infer_img] images = set() infer_dir = os.path.abspath(infer_dir) assert os.path.isdir(infer_dir), \ "infer_dir {} is not a directory".format(infer_dir) exts = ['jpg', 'jpeg', 'png', 'bmp'] exts += [ext.upper() for ext in exts] for ext in exts: images.update(glob.glob('{}/*.{}'.format(infer_dir, ext))) images = list(images) assert len(images) > 0, "no image found in {}".format(infer_dir) logger.info("Found {} inference images in total.".format(len(images))) return images def compute_ious(boxes1, boxes2): """[Compute pairwise IOU matrix for given two sets of boxes] Args: boxes1 ([numpy ndarray with shape N,4]): [representing bounding boxes with format (xmin,ymin,xmax,ymax)] boxes2 ([numpy ndarray with shape M,4]): [representing bounding boxes with format (xmin,ymin,xmax,ymax)] Returns: pairwise IOU maxtrix with shape (N,M)，where the value at ith row jth column hold the iou between ith box and jth box from box1 and box2 respectively. """ lu = np.maximum( boxes1[:, None, :2], boxes2[:, :2] ) # lu with shape N,M,2 ; boxes1[:,None,:2] with shape (N,1,2) boxes2 with shape(M,2) rd = np.minimum(boxes1[:, None, 2:], boxes2[:, 2:]) # rd same to lu intersection_wh = np.maximum(0.0, rd - lu) intersection_area = intersection_wh[:, :, 0] * intersection_wh[:, :, 1] # with shape (N,M) boxes1_wh = np.maximum(0.0, boxes1[:, 2:] - boxes1[:, :2]) boxes1_area = boxes1_wh[:, 0] * boxes1_wh[:, 1] # with shape (N,) boxes2_wh = np.maximum(0.0, boxes2[:, 2:] - boxes2[:, :2]) boxes2_area = boxes2_wh[:, 0] * boxes2_wh[:, 1] # with shape (M,) union_area = np.maximum( boxes1_area[:, None] + boxes2_area - intersection_area, 1e-8) # with shape (N,M) ious = np.clip(intersection_area / union_area, 0.0, 1.0) return ious def grad_cam(feat, grad): """ Args: feat: CxHxW grad: CxHxW Returns: cam: HxW """ exp = (feat * grad.mean((1, 2), keepdims=True)).mean(axis=0) exp = np.maximum(-exp, 0) return exp def resize_cam(explanation, resize_shape) -> np.ndarray: """ Args: explanation: (width, height) resize_shape: (width, height) Returns: """ assert len(explanation.shape) == 2, f"{explanation.shape}. " \ f"Currently support 2D explanation results for visualization. " \ "Reduce higher dimensions to 2D for visualization." explanation = (explanation - explanation.min()) / ( explanation.max() - explanation.min()) explanation = cv2.resize(explanation, resize_shape) explanation = np.uint8(255 * explanation) explanation = cv2.applyColorMap(explanation, cv2.COLORMAP_JET) explanation = cv2.cvtColor(explanation, cv2.COLOR_BGR2RGB) return explanation class BBoxCAM: def __init__(self, FLAGS, cfg): self.FLAGS = FLAGS self.cfg = cfg # build model self.trainer = self.build_trainer(cfg) # num_class self.num_class = cfg.num_classes # set hook for extraction of featuremaps and grads self.set_hook(cfg) self.nms_idx_need_divid_numclass_arch = ['FasterRCNN', 'MaskRCNN', 'CascadeRCNN'] """ In these networks, the bbox array shape before nms contain num_class, the nms_keep_idx of the bbox need to divide the num_class; """ # cam image output_dir try: os.makedirs(FLAGS.cam_out) except: print('Path already exists.') pass def build_trainer(self, cfg): # build trainer trainer = Trainer(cfg, mode='test') # load weights trainer.load_weights(cfg.weights) # set for get extra_data before nms trainer.model.use_extra_data=True # set for record the bbox index before nms if cfg.architecture in ['FasterRCNN', 'MaskRCNN']: trainer.model.bbox_post_process.nms.return_index = True elif cfg.architecture in ['YOLOv3', 'PPYOLOE', 'PPYOLOEWithAuxHead']: if trainer.model.post_process is not None: # anchor based YOLOs: YOLOv3,PP-YOLO trainer.model.post_process.nms.return_index = True else: # anchor free YOLOs: PP-YOLOE, PP-YOLOE+ trainer.model.yolo_head.nms.return_index = True elif cfg.architecture=='BlazeFace' or cfg.architecture=='SSD': trainer.model.post_process.nms.return_index = True elif cfg.architecture=='RetinaNet': trainer.model.head.nms.return_index = True else: print( cfg.architecture+' is not supported for cam temporarily!' ) sys.exit() # Todo: Unify the head/post_process name in each model return trainer def set_hook(self, cfg): # set hook for extraction of featuremaps and grads self.target_feats = {} self.target_layer_name = cfg.target_feature_layer_name # such as trainer.model.backbone, trainer.model.bbox_head.roi_extractor def hook(layer, input, output): self.target_feats[layer._layer_name_for_hook] = output try: exec('self.trainer.'+self.target_layer_name+'._layer_name_for_hook = self.target_layer_name') # self.trainer.target_layer_name._layer_name_for_hook = self.target_layer_name exec('self.trainer.'+self.target_layer_name+'.register_forward_post_hook(hook)') # self.trainer.target_layer_name.register_forward_post_hook(hook) except: print("Error! " "The target_layer_name--"+self.target_layer_name+" is not in model! " "Please check the spelling and " "the network's architecture!") sys.exit() def get_bboxes(self): # get inference images images = get_test_images(self.FLAGS.infer_dir, self.FLAGS.infer_img) # inference result = self.trainer.predict( images, draw_threshold=self.FLAGS.draw_threshold, output_dir=self.FLAGS.output_dir, save_results=self.FLAGS.save_results, visualize=False)[0] return result def get_bboxes_cams(self): # Get the bboxes prediction(after nms result) of the input inference_result = self.get_bboxes() # read input image # Todo: Support folder multi-images process from PIL import Image img = np.array(Image.open(self.cfg.infer_img)) # data for calaulate bbox grad_cam extra_data = inference_result['extra_data'] """ Example of Faster_RCNN based architecture: extra_data: {'scores': tensor with shape [num_of_bboxes_before_nms, num_classes], for example: [1000, 80] 'nms_keep_idx': tensor with shape [num_of_bboxes_after_nms, 1], for example: [300, 1] } Example of YOLOv3 based architecture: extra_data: {'scores': tensor with shape [1, num_classes, num_of_yolo_bboxes_before_nms], #for example: [1, 80, 8400] 'nms_keep_idx': tensor with shape [num_of_yolo_bboxes_after_nms, 1], # for example: [300, 1] } """ # array index of the predicted bbox before nms if self.cfg.architecture in self.nms_idx_need_divid_numclass_arch: # some network's bbox array shape before nms may be like [num_of_bboxes_before_nms, num_classes, 4], # we need to divide num_classes to get the before_nms_index； # currently, only include the rcnn architectures （fasterrcnn, maskrcnn, cascadercnn); before_nms_indexes = extra_data['nms_keep_idx'].cpu().numpy( ) // self.num_class # num_class else : before_nms_indexes = extra_data['nms_keep_idx'].cpu().numpy() # Calculate and visualize the heatmap of per predict bbox for index, target_bbox in enumerate(inference_result['bbox']): # target_bbox: [cls, score, x1, y1, x2, y2] # filter bboxes with low predicted scores if target_bbox[1] < self.FLAGS.draw_threshold: continue target_bbox_before_nms = int(before_nms_indexes[index]) if len(extra_data['scores'].shape)==2: score_out = extra_data['scores'][target_bbox_before_nms] else: score_out = extra_data['scores'][0, :, target_bbox_before_nms] """ There are two kinds array shape of bbox score output : 1) [num_of_bboxes_before_nms, num_classes], for example: [1000, 80] 2) [num_of_image, num_classes, num_of_yolo_bboxes_before_nms], for example: [1, 80, 1000] """ # construct one_hot label and do backward to get the gradients predicted_label = paddle.argmax(score_out) label_onehot = paddle.nn.functional.one_hot( predicted_label, num_classes=len(score_out)) label_onehot = label_onehot.squeeze() target = paddle.sum(score_out * label_onehot) target.backward(retain_graph=True) if 'backbone' in self.target_layer_name or \ 'neck' in self.target_layer_name: # backbone/neck level feature if isinstance(self.target_feats[self.target_layer_name], list): # when the featuremap contains of multiple scales, # take the featuremap of the last scale # Todo: fuse the cam result from multisclae featuremaps if self.target_feats[self.target_layer_name][ -1].shape[-1]==1: """ if the last level featuremap is 1x1 size, we take the second last one """ cam_grad = self.target_feats[self.target_layer_name][ -2].grad.squeeze().cpu().numpy() cam_feat = self.target_feats[self.target_layer_name][ -2].squeeze().cpu().numpy() else: cam_grad = self.target_feats[self.target_layer_name][ -1].grad.squeeze().cpu().numpy() cam_feat = self.target_feats[self.target_layer_name][ -1].squeeze().cpu().numpy() else: cam_grad = self.target_feats[ self.target_layer_name].grad.squeeze().cpu().numpy() cam_feat = self.target_feats[ self.target_layer_name].squeeze().cpu().numpy() else: # roi level feature cam_grad = self.target_feats[ self.target_layer_name].grad.squeeze().cpu().numpy()[target_bbox_before_nms] cam_feat = self.target_feats[ self.target_layer_name].squeeze().cpu().numpy()[target_bbox_before_nms] # grad_cam: exp = grad_cam(cam_feat, cam_grad) if 'backbone' in self.target_layer_name or \ 'neck' in self.target_layer_name: """ when use backbone/neck featuremap, we first do the cam on whole image, and then set the area outside the predic bbox to 0 """ # reshape the cam image to the input image size resized_exp = resize_cam(exp, (img.shape[1], img.shape[0])) mask = np.zeros((img.shape[0], img.shape[1], 3)) mask[int(target_bbox[3]):int(target_bbox[5]), int(target_bbox[2]): int(target_bbox[4]), :] = 1 resized_exp = resized_exp * mask # add the bbox cam back to the input image overlay_vis = np.uint8(resized_exp * 0.4 + img * 0.6) elif 'roi' in self.target_layer_name: # get the bbox part of the image bbox_img = copy.deepcopy(img[int(target_bbox[3]):int(target_bbox[5]), int(target_bbox[2]):int(target_bbox[4]), :]) # reshape the cam image to the bbox size resized_exp = resize_cam(exp, (bbox_img.shape[1], bbox_img.shape[0])) # add the bbox cam back to the bbox image bbox_overlay_vis = np.uint8(resized_exp * 0.4 + bbox_img * 0.6) # put the bbox_cam image to the original image overlay_vis = copy.deepcopy(img) overlay_vis[int(target_bbox[3]):int(target_bbox[5]), int(target_bbox[2]):int(target_bbox[4]), :] = bbox_overlay_vis else: print( 'Only supported cam for backbone/neck feature and roi feature, the others are not supported temporarily!' ) sys.exit() # put the bbox rectangle on image cv2.rectangle( overlay_vis, (int(target_bbox[2]), int(target_bbox[3])), (int(target_bbox[4]), int(target_bbox[5])), (0, 0, 255), 2) # save visualization result cam_image = Image.fromarray(overlay_vis) cam_image.save(self.FLAGS.cam_out + '/' + str(index) + '.jpg') # clear gradients after each bbox grad_cam target.clear_gradient() for n, v in self.trainer.model.named_sublayers(): v.clear_gradients()