训练数据前置处理与提升训练效率
This commit is contained in:
lee
164
tools/getHeatMap.py
Normal file
164
tools/getHeatMap.py
Normal file
@ -0,0 +1,164 @@
|
||||
# -*- coding: UTF-8 -*-
|
||||
import os
|
||||
|
||||
import torch
|
||||
from torchvision import models
|
||||
import torch.nn as nn
|
||||
import torchvision.transforms as tfs
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
from PIL import Image
|
||||
import cv2
|
||||
# from tools.config import cfg
|
||||
# from comparative.tools.initmodel import initSimilarityModel
|
||||
import yaml
|
||||
from dataset import get_transform
|
||||
|
||||
|
||||
class cal_cam(nn.Module):
|
||||
def __init__(self, model, conf):
|
||||
super(cal_cam, self).__init__()
|
||||
self.conf = conf
|
||||
self.device = self.conf['base']['device']
|
||||
|
||||
self.model = model
|
||||
self.model.to(self.device)
|
||||
|
||||
# 要求梯度的层
|
||||
self.feature_layer = conf['heatmap']['feature_layer']
|
||||
# 记录梯度
|
||||
self.gradient = []
|
||||
# 记录输出的特征图
|
||||
self.output = []
|
||||
_, self.transform = get_transform(self.conf)
|
||||
|
||||
def get_conf(self, yaml_pth):
|
||||
with open(yaml_pth, 'r') as f:
|
||||
conf = yaml.load(f, Loader=yaml.FullLoader)
|
||||
return conf
|
||||
|
||||
def save_grad(self, grad):
|
||||
self.gradient.append(grad)
|
||||
|
||||
def get_grad(self):
|
||||
return self.gradient[-1].cpu().data
|
||||
|
||||
def get_feature(self):
|
||||
return self.output[-1][0]
|
||||
|
||||
def process_img(self, input):
|
||||
input = self.transform(input)
|
||||
input = input.unsqueeze(0)
|
||||
return input
|
||||
|
||||
# 计算最后一个卷积层的梯度，输出梯度和最后一个卷积层的特征图
|
||||
def getGrad(self, input_):
|
||||
self.gradient = [] # 清除之前的梯度
|
||||
self.output = [] # 清除之前的特征图
|
||||
# print(f"cuda.memory_allocated 1 {torch.cuda.memory_allocated()/ (1024 ** 3)}G")
|
||||
input_ = input_.to(self.device).requires_grad_(True)
|
||||
num = 1
|
||||
for name, module in self.model._modules.items():
|
||||
# print(f'module_name: {name}')
|
||||
# print(f'module: {module}')
|
||||
if (num == 1):
|
||||
input = module(input_)
|
||||
num = num + 1
|
||||
continue
|
||||
# 是待提取特征图的层
|
||||
if (name == self.feature_layer):
|
||||
input = module(input)
|
||||
input.register_hook(self.save_grad)
|
||||
self.output.append([input])
|
||||
# 马上要到全连接层了
|
||||
elif (name == "avgpool"):
|
||||
input = module(input)
|
||||
input = input.reshape(input.shape[0], -1)
|
||||
# 普通的层
|
||||
else:
|
||||
input = module(input)
|
||||
|
||||
# print(f"cuda.memory_allocated 2 {torch.cuda.memory_allocated() / (1024 ** 3)}G")
|
||||
# 到这里input就是最后全连接层的输出了
|
||||
index = torch.max(input, dim=-1)[1]
|
||||
one_hot = torch.zeros((1, input.shape[-1]), dtype=torch.float32)
|
||||
one_hot[0][index] = 1
|
||||
confidenct = one_hot * input.cpu()
|
||||
confidenct = torch.sum(confidenct, dim=-1).requires_grad_(True)
|
||||
|
||||
# print(f"cuda.memory_allocated 3 {torch.cuda.memory_allocated() / (1024 ** 3)}G")
|
||||
# 清除之前的所有梯度
|
||||
self.model.zero_grad()
|
||||
# 反向传播获取梯度
|
||||
grad_output = torch.ones_like(confidenct)
|
||||
confidenct.backward(grad_output)
|
||||
# 获取特征图的梯度
|
||||
grad_val = self.get_grad()
|
||||
feature = self.get_feature()
|
||||
|
||||
# print(f"cuda.memory_allocated 4 {torch.cuda.memory_allocated() / (1024 ** 3)}G")
|
||||
return grad_val, feature, input_.grad
|
||||
|
||||
# 计算CAM
|
||||
def getCam(self, grad_val, feature):
|
||||
# 对特征图的每个通道进行全局池化
|
||||
alpha = torch.mean(grad_val, dim=(2, 3)).cpu()
|
||||
feature = feature.cpu()
|
||||
# 将池化后的结果和相应通道特征图相乘
|
||||
cam = torch.zeros((feature.shape[2], feature.shape[3]), dtype=torch.float32)
|
||||
for idx in range(alpha.shape[1]):
|
||||
cam = cam + alpha[0][idx] * feature[0][idx]
|
||||
# 进行ReLU操作
|
||||
cam = np.maximum(cam.detach().numpy(), 0)
|
||||
|
||||
# plt.imshow(cam)
|
||||
# plt.colorbar()
|
||||
# plt.savefig("cam.jpg")
|
||||
|
||||
# 将cam区域放大到输入图片大小
|
||||
cam_ = cv2.resize(cam, (224, 224))
|
||||
cam_ = cam_ - np.min(cam_)
|
||||
cam_ = cam_ / np.max(cam_)
|
||||
# plt.imshow(cam_)
|
||||
# plt.savefig("cam_.jpg")
|
||||
cam = torch.from_numpy(cam)
|
||||
|
||||
return cam, cam_
|
||||
|
||||
def show_img(self, cam_, img, heatmap_save_pth, imgname):
|
||||
heatmap = cv2.applyColorMap(np.uint8(255 * cam_), cv2.COLORMAP_JET)
|
||||
cam_img = 0.3 * heatmap + 0.7 * np.float32(img)
|
||||
# cv2.imwrite("img.jpg", cam_img)
|
||||
cv2.imwrite(os.sep.join([heatmap_save_pth, imgname]), cam_img)
|
||||
|
||||
def get_hot_map(self, img_pth):
|
||||
img = Image.open(img_pth)
|
||||
img = img.resize((224, 224))
|
||||
input = self.process_img(img)
|
||||
grad_val, feature, input_grad = self.getGrad(input)
|
||||
cam, cam_ = self.getCam(grad_val, feature)
|
||||
heatmap = cv2.applyColorMap(np.uint8(255 * cam_), cv2.COLORMAP_JET)
|
||||
cam_img = 0.3 * heatmap + 0.7 * np.float32(img)
|
||||
cam_img = Image.fromarray(np.uint8(cam_img))
|
||||
return cam_img
|
||||
|
||||
# def __call__(self, img_root, heatmap_save_pth):
|
||||
# for imgname in os.listdir(img_root):
|
||||
# img = Image.open(os.sep.join([img_root, imgname]))
|
||||
# img = img.resize((224, 224))
|
||||
# # plt.imshow(img)
|
||||
# # plt.savefig("airplane.jpg")
|
||||
# input = self.process_img(img)
|
||||
# grad_val, feature, input_grad = self.getGrad(input)
|
||||
# cam, cam_ = self.getCam(grad_val, feature)
|
||||
# self.show_img(cam_, img, heatmap_save_pth, imgname)
|
||||
# return cam
|
||||
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cam = cal_cam()
|
||||
img_root = "test_img/"
|
||||
heatmap_save_pth = "heatmap_result"
|
||||
cam(img_root, heatmap_save_pth)
|
||||
@ -188,7 +188,7 @@ class PairGenerator:
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
original_path = '/home/lc/data_center/scatter/val_extar'
|
||||
original_path = '/home/lc/data_center/scatter/v4/val'
|
||||
parent_dir = str(Path(original_path).parent)
|
||||
generator = PairGenerator()
|
||||
|
||||
|
||||
@ -1,33 +1,50 @@
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
from tools.getHeatMap import cal_cam
|
||||
import os
|
||||
|
||||
|
||||
def merge_imgs(img1_path, img2_path, save_path, similar=None, label=None):
|
||||
position = (50, 50) # 文字的左上角坐标
|
||||
color = (255, 0, 0) # 红色文字，格式为 RGB
|
||||
if not os.path.exists(os.sep.join([save_path, str(label)])):
|
||||
os.makedirs(os.sep.join([save_path, str(label)]))
|
||||
save_path = os.sep.join([save_path, str(label)])
|
||||
img_name = os.path.basename(img1_path).split('.')[0]+'_'+os.path.basename(img2_path).split('.')[0]+'.png'
|
||||
def merge_imgs(img1_path, img2_path, conf, similar=None, label=None, cam=None):
|
||||
save = True
|
||||
position = (50, 50) # 文字的左上角坐标
|
||||
color = (255, 0, 0) # 红色文字，格式为 RGB
|
||||
# if not os.path.exists(os.sep.join([save_path, str(label)])):
|
||||
# os.makedirs(os.sep.join([save_path, str(label)]))
|
||||
# save_path = os.sep.join([save_path, str(label)])
|
||||
# img_name = os.path.basename(img1_path).split('.')[0] + '_' + os.path.basename(img2_path).split('.')[0] + '.png'
|
||||
if not conf['heatmap']['show_heatmap']:
|
||||
img1 = Image.open(img1_path)
|
||||
img2 = Image.open(img2_path)
|
||||
img1 = img1.resize((224,224))
|
||||
img2 = img2.resize((224,224))
|
||||
print('img1_path', img1)
|
||||
print('img2_path', img2)
|
||||
assert img1.height == img2.height
|
||||
img1 = img1.resize((224, 224))
|
||||
img2 = img2.resize((224, 224))
|
||||
save_path = conf['data']['image_joint_pth']
|
||||
else:
|
||||
assert cam is not None, 'cam is None'
|
||||
img1 = cam.get_hot_map(img1_path)
|
||||
img2 = cam.get_hot_map(img2_path)
|
||||
save_path = conf['heatmap']['image_joint_pth']
|
||||
# print('img1_path', img1)
|
||||
# print('img2_path', img2)
|
||||
if not os.path.exists(os.sep.join([save_path, str(label)])):
|
||||
os.makedirs(os.sep.join([save_path, str(label)]))
|
||||
save_path = os.sep.join([save_path, str(label)])
|
||||
img_name = os.path.basename(img1_path).split('.')[0] + '_' + os.path.basename(img2_path).split('.')[0] + '.png'
|
||||
assert img1.height == img2.height
|
||||
|
||||
new_img = Image.new('RGB', (img1.width + img2.width + 10, img1.height))
|
||||
new_img = Image.new('RGB', (img1.width + img2.width + 10, img1.height))
|
||||
|
||||
# print('new_img', new_img)
|
||||
new_img.paste(img1, (0, 0))
|
||||
new_img.paste(img2, (img1.width + 10, 0))
|
||||
# print('new_img', new_img)
|
||||
new_img.paste(img1, (0, 0))
|
||||
new_img.paste(img2, (img1.width + 10, 0))
|
||||
|
||||
if similar is not None:
|
||||
similar = str(similar)+'_'+str(label)
|
||||
draw = ImageDraw.Draw(new_img)
|
||||
draw.text(position, str(similar), color, font_size=36)
|
||||
os.makedirs(save_path, exist_ok=True)
|
||||
img_save = os.path.join(save_path, img_name)
|
||||
if similar is not None:
|
||||
if label == '1' and similar > 0.5:
|
||||
save = False
|
||||
elif label == '0' and similar < 0.5:
|
||||
save = False
|
||||
similar = str(similar) + '_' + str(label)
|
||||
draw = ImageDraw.Draw(new_img)
|
||||
draw.text(position, str(similar), color, font_size=36)
|
||||
os.makedirs(save_path, exist_ok=True)
|
||||
img_save = os.path.join(save_path, img_name)
|
||||
if save:
|
||||
new_img.save(img_save)
|
||||
|
||||
|
||||
@ -96,7 +96,7 @@ if __name__ == '__main__':
|
||||
rknn.config(
|
||||
mean_values=[[127.5, 127.5, 127.5]],
|
||||
std_values=[[127.5, 127.5, 127.5]],
|
||||
target_platform='rk3588',
|
||||
target_platform='rk3566',
|
||||
model_pruning=False,
|
||||
compress_weight=False,
|
||||
single_core_mode=True,
|
||||
|
||||
Reference in New Issue
Block a user