from .pytorch_grad_cam import GradCAM from .pytorch_grad_cam.utils.image import scale_cam_image from . import clip from .utils import reshape_transform, zeroshot_classifier, ClipOutputTarget, scoremap2bbox from torchvision import transforms import torch import numpy as np import cv2 class CAMGenerator: def __init__(self, categories, clip_path): # self.device = "cuda" if torch.cuda.is_available() else "cpu" self.clip_path = clip_path self.clip_model, _ = clip.load(self.clip_path, device="cpu") self.target_layers = [self.clip_model.visual.transformer.resblocks[-1].ln_1] self.cam = GradCAM(model=self.clip_model, target_layers=self.target_layers, reshape_transform=reshape_transform, use_cuda=True) self.categories = categories # if categories is None: # self.categories = ['vehicle', 'baseballfield', 'groundtrackfield', 'windmill', 'bridge', 'overpass', 'ship', 'airplane', 'tenniscourt', 'airport', # 'expressway-service-area', 'basketballcourt', 'stadium', 'storagetank', 'chimney', 'dam', 'expressway-toll-station', 'golffield', 'trainstation', 'harbor'] self.background_categories = ['ground','land','grass','tree','building','wall','sky','lake','water','river','sea','railway','railroad','keyboard','helmet', 'cloud','house','mountain','ocean','road','rock','street','valley','bridge','sign',] self.normalize = transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)) def _prepare(self): self.bg_text_features = zeroshot_classifier(self.background_categories, ['a clean origami {}.'], self.clip_model, self.device) self.fg_text_features = zeroshot_classifier(self.categories, ['a clean origami {}.'], self.clip_model, self.device) def to(self, device, dtype): self.device = device self.clip_model.to(device) self.cam.set_device(device) self._prepare() def re_normalize(self, image): # image = (image / 2) + 0.5 image = self.normalize(image) return image def get_label_list(self, captions): label_list = [] label_id_list = [] for caption in captions: if caption in self.categories and caption not in label_list: label_list.append(caption) label_id_list.append(self.categories.index(caption)) return label_list, label_id_list def get_label_list_with_bboxes(self, captions, bboxes): label_list, label_id_list, bboxes_list = [], [], [] for caption, bbox in zip(captions, bboxes): if caption not in self.categories: continue if caption not in label_list: label_list.append(caption) label_id_list.append(self.categories.index(caption)) bboxes_list.append([bbox]) else: bboxes_list[label_list.index(caption)].append(bbox) return label_list, label_id_list, bboxes_list def __call__(self, image, captions, bboxes, gt_bboxes_only=False): image = self.re_normalize(image) # label_list, label_id_list = self.get_label_list(captions[0][1:]) label_list, label_id_list, bboxes_list = self.get_label_list_with_bboxes(captions[0][1:], bboxes[0]) h, w = image.shape[-2], image.shape[-1] image_features, attn_weight_list = self.clip_model.encode_image(image, h, w) bg_features_temp = self.bg_text_features fg_features_temp = self.fg_text_features[label_id_list] text_features_temp = torch.cat([fg_features_temp, bg_features_temp], dim=0) input_tensor = [image_features, text_features_temp, h, w] keys, refined_cam_list = [], [] for idx, (label, bbox) in enumerate(zip(label_list, bboxes_list)): keys.append(self.categories.index(label)) targets = [ClipOutputTarget(label_list.index(label))] grayscale_cam, logits_per_image, attn_weight_last = self.cam(input_tensor=input_tensor, targets=targets, target_size=None) grayscale_cam = grayscale_cam[0, :] # [32, 32] # grayscale_cam_highres = cv2.resize(grayscale_cam, (ori_width, ori_height)) if idx == 0: attn_weight_list.append(attn_weight_last) attn_weight = [aw[:, 1:, 1:] for aw in attn_weight_list] # (b, hxw, hxw) attn_weight = torch.stack(attn_weight, dim=0)[-8:] # [8, 1, 1024, 1024] attn_weight = torch.mean(attn_weight, dim=0) # [1, 1024, 1024] # attn_weight = attn_weight[0].detach() # [1024, 1024] # original detach attn_weight = attn_weight[0] #.detach() # [1024, 1024] attn_weight = attn_weight.float() gt_box, gt_cnt = (np.array(bbox) * [grayscale_cam.shape[1], grayscale_cam.shape[0], grayscale_cam.shape[1], grayscale_cam.shape[0]]).astype(int), len(bbox) if gt_bboxes_only: box, cnt = gt_box, gt_cnt else: box, cnt = scoremap2bbox(scoremap=grayscale_cam.cpu().data.numpy(), threshold=0.4, multi_contour_eval=True) box, cnt = np.concatenate([box, gt_box], axis=0), cnt + gt_cnt aff_mask = torch.zeros_like(grayscale_cam) for i_ in range(cnt): x0_, y0_, x1_, y1_ = box[i_] aff_mask[y0_:y1_, x0_:x1_] = 1 aff_mask = aff_mask.view(1, grayscale_cam.shape[0] * grayscale_cam.shape[1]) aff_mat = attn_weight trans_mat = aff_mat / torch.sum(aff_mat, dim=0, keepdim=True) trans_mat = trans_mat / torch.sum(trans_mat, dim=1, keepdim=True) for _ in range(2): trans_mat = trans_mat / torch.sum(trans_mat, dim=0, keepdim=True) trans_mat = trans_mat / torch.sum(trans_mat, dim=1, keepdim=True) trans_mat = (trans_mat + trans_mat.transpose(1, 0)) / 2 for _ in range(1): trans_mat = torch.matmul(trans_mat, trans_mat) trans_mat = trans_mat * aff_mask cam_to_refine = grayscale_cam.view(-1, 1) cam_refined = torch.matmul(trans_mat, cam_to_refine).reshape(h //16, w // 16) cam_refined = cam_refined - cam_refined.min() cam_refined = cam_refined / (cam_refined.max() + 1e-7) refined_cam_list.append(cam_refined) keys = torch.tensor(keys) refined_cams = torch.stack(refined_cam_list, dim=0) return refined_cams, keys if __name__ == '__main__': cam = CAMGenerator()