diff --git a/README.md b/README.md index d070aad..8280d8a 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,7 @@ image classifier implement in pytoch. - **简单的安装过程** 1. 安装好pytorch, torchvision(pytorch==1.12.0+ torchvision==0.13.0+) 可以在[pytorch](https://pytorch.org/get-started/previous-versions/)官网找到对应的命令进行安装. - 2. pip install -r requirements.txt + 2. pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple - **人性化的设定** 1. 大部分可视化数据(混淆矩阵,tsne,每个类别的指标)都会以csv或者log的格式保存到本地,方便后期美工图像. @@ -245,7 +245,7 @@ image classifier implement in pytoch. 是否采用测试阶段的数据增强. - **cam_visual** default: False - 是否进行热力图可视化. + 是否进行热力图可视化.(目前只支持在cpu上进行可视化热力图) - **cam_type** type: string, default: GradCAMPlusPlus, choices: ['GradCAM', 'HiResCAM', 'ScoreCAM', 'GradCAMPlusPlus', 'AblationCAM', 'XGradCAM', 'EigenCAM', 'FullGrad'] 热力图可视化的类型. @@ -533,7 +533,7 @@ image classifier implement in pytoch. 1. 在predict.py中的--cam_type参数,可以选择各种各样的热力图计算方法,如果遇到这个效果不好可以进行更换,当然效果不好除了跟热力图的计算方法有关系之外还有是你选择的层. 2. 在第一点中说到的层,热力图可视化需要选择层进行计算,这个是在predict.py中的cam_visual方法定义中进行定义. - 例如这样 cam_model = cam_visual(model, test_transform, DEVICE, model.cam_layer(), opt) 这个model.cam_layer()就是选择的层,默认就是卷积特征的最后一个block,这个可以在model文件夹中的每个模型文件中的模型进行修改. + 例如这样 cam_model = cam_visual(model, test_transform, DEVICE, model.cam_layer(), opt) 这个model.cam_layer()就是选择的层,默认就是卷积特征的最后一个layer,这个可以在model文件夹中的每个模型文件中的模型进行修改. 13. 关于类别平衡的问题. @@ -679,4 +679,11 @@ image classifier implement in pytoch. 1. predict.py支持检测灰度图,其读取后会检测是否为RGB通道,不是的话会进行转换. 2. 更新readme.md. -3. 修复一些bug. \ No newline at end of file +3. 修复一些bug. + +### pytorch-classifier v1.5 更新日志 + +1. 更新readme.md +2. 修改predict.py中的读取模型方式. +3. 修复predict.py中批量预测保存图片时候的内存溢出问题. +4. 修复predict.py中热力图可视化的报错问题,但是目前只支持cpu状态下可视化热力图. \ No newline at end of file diff --git a/model/__pycache__/__init__.cpython-38.pyc b/model/__pycache__/__init__.cpython-38.pyc index 177d5b3..24c9f6a 100644 Binary files a/model/__pycache__/__init__.cpython-38.pyc and b/model/__pycache__/__init__.cpython-38.pyc differ diff --git a/model/__pycache__/convnext.cpython-38.pyc b/model/__pycache__/convnext.cpython-38.pyc index 37ce41d..c0029f9 100644 Binary files a/model/__pycache__/convnext.cpython-38.pyc and b/model/__pycache__/convnext.cpython-38.pyc differ diff --git a/model/__pycache__/cspnet.cpython-38.pyc b/model/__pycache__/cspnet.cpython-38.pyc index aa6c3d8..5ffa9a1 100644 Binary files a/model/__pycache__/cspnet.cpython-38.pyc and b/model/__pycache__/cspnet.cpython-38.pyc differ diff --git a/model/__pycache__/densenet.cpython-38.pyc b/model/__pycache__/densenet.cpython-38.pyc index 448ea6b..b65e6a0 100644 Binary files a/model/__pycache__/densenet.cpython-38.pyc and b/model/__pycache__/densenet.cpython-38.pyc differ diff --git a/model/__pycache__/dpn.cpython-38.pyc b/model/__pycache__/dpn.cpython-38.pyc index 2930103..1a16282 100644 Binary files a/model/__pycache__/dpn.cpython-38.pyc and b/model/__pycache__/dpn.cpython-38.pyc differ diff --git a/model/__pycache__/efficientnetv2.cpython-38.pyc b/model/__pycache__/efficientnetv2.cpython-38.pyc index 3e9a487..0c089b8 100644 Binary files a/model/__pycache__/efficientnetv2.cpython-38.pyc and b/model/__pycache__/efficientnetv2.cpython-38.pyc differ diff --git a/model/__pycache__/ghostnet.cpython-38.pyc b/model/__pycache__/ghostnet.cpython-38.pyc index ad16ffd..e5d0c21 100644 Binary files a/model/__pycache__/ghostnet.cpython-38.pyc and b/model/__pycache__/ghostnet.cpython-38.pyc differ diff --git a/model/__pycache__/mnasnet.cpython-38.pyc b/model/__pycache__/mnasnet.cpython-38.pyc index 04e2329..b0bb409 100644 Binary files a/model/__pycache__/mnasnet.cpython-38.pyc and b/model/__pycache__/mnasnet.cpython-38.pyc differ diff --git a/model/__pycache__/mobilenetv2.cpython-38.pyc b/model/__pycache__/mobilenetv2.cpython-38.pyc index ee27d93..84d8c33 100644 Binary files a/model/__pycache__/mobilenetv2.cpython-38.pyc and b/model/__pycache__/mobilenetv2.cpython-38.pyc differ diff --git a/model/__pycache__/mobilenetv3.cpython-38.pyc b/model/__pycache__/mobilenetv3.cpython-38.pyc index d37248a..7a9164b 100644 Binary files a/model/__pycache__/mobilenetv3.cpython-38.pyc and b/model/__pycache__/mobilenetv3.cpython-38.pyc differ diff --git a/model/__pycache__/repghost.cpython-38.pyc b/model/__pycache__/repghost.cpython-38.pyc index 741c13c..12a0f53 100644 Binary files a/model/__pycache__/repghost.cpython-38.pyc and b/model/__pycache__/repghost.cpython-38.pyc differ diff --git a/model/__pycache__/repvgg.cpython-38.pyc b/model/__pycache__/repvgg.cpython-38.pyc index 92a9b21..228e904 100644 Binary files a/model/__pycache__/repvgg.cpython-38.pyc and b/model/__pycache__/repvgg.cpython-38.pyc differ diff --git a/model/__pycache__/resnest.cpython-38.pyc b/model/__pycache__/resnest.cpython-38.pyc index 062bb52..96da18d 100644 Binary files a/model/__pycache__/resnest.cpython-38.pyc and b/model/__pycache__/resnest.cpython-38.pyc differ diff --git a/model/__pycache__/resnet.cpython-38.pyc b/model/__pycache__/resnet.cpython-38.pyc index 72c7996..4fe41a9 100644 Binary files a/model/__pycache__/resnet.cpython-38.pyc and b/model/__pycache__/resnet.cpython-38.pyc differ diff --git a/model/__pycache__/sequencer.cpython-38.pyc b/model/__pycache__/sequencer.cpython-38.pyc index 6e0729c..0544350 100644 Binary files a/model/__pycache__/sequencer.cpython-38.pyc and b/model/__pycache__/sequencer.cpython-38.pyc differ diff --git a/model/__pycache__/shufflenetv2.cpython-38.pyc b/model/__pycache__/shufflenetv2.cpython-38.pyc index 77d7ce9..74cfcae 100644 Binary files a/model/__pycache__/shufflenetv2.cpython-38.pyc and b/model/__pycache__/shufflenetv2.cpython-38.pyc differ diff --git a/model/__pycache__/vgg.cpython-38.pyc b/model/__pycache__/vgg.cpython-38.pyc index be54e01..6e3fee5 100644 Binary files a/model/__pycache__/vgg.cpython-38.pyc and b/model/__pycache__/vgg.cpython-38.pyc differ diff --git a/model/__pycache__/vovnet.cpython-38.pyc b/model/__pycache__/vovnet.cpython-38.pyc index 61e939e..06f25de 100644 Binary files a/model/__pycache__/vovnet.cpython-38.pyc and b/model/__pycache__/vovnet.cpython-38.pyc differ diff --git a/model/efficientnetv2.py b/model/efficientnetv2.py index acdbc76..addf4b5 100644 --- a/model/efficientnetv2.py +++ b/model/efficientnetv2.py @@ -383,7 +383,7 @@ def forward_features(self, x, need_fea=False): return x def cam_layer(self): - return self.features[-1] + return self.features[-5:] def switch_to_deploy(self): new_block = [] diff --git a/predict.py b/predict.py index 98ef390..f7eecac 100644 --- a/predict.py +++ b/predict.py @@ -6,8 +6,10 @@ os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE" import matplotlib.pyplot as plt import numpy as np +from copy import deepcopy from utils import utils_aug from utils.utils import predict_single_image, cam_visual, dict_to_PrettyTable, select_device, model_fuse +from utils.utils_model import select_model def set_seed(seed): random.seed(seed) @@ -27,6 +29,7 @@ def parse_opt(): parser.add_argument('--device', type=str, default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') opt = parser.parse_known_args()[0] + if not os.path.exists(os.path.join(opt.save_path, 'best.pt')): raise Exception('best.pt not found. please check your --save_path folder') ckpt = torch.load(os.path.join(opt.save_path, 'best.pt')) @@ -34,8 +37,13 @@ def parse_opt(): if opt.half and DEVICE.type == 'cpu': raise Exception('half inference only supported GPU.') if opt.half and opt.cam_visual: + raise Exception('cam visual only supported cpu. please set device=cpu.') + if (opt.device != 'cpu') and opt.cam_visual: raise Exception('cam visual only supported FP32.') - model = ckpt['model'].float() + with open(opt.label_path) as f: + CLASS_NUM = len(f.readlines()) + model = select_model(ckpt['model'].name, CLASS_NUM) + model.load_state_dict(ckpt['model'].float().state_dict(), strict=False) model_fuse(model) model = (model.half() if opt.half else model) model.to(DEVICE) @@ -59,7 +67,7 @@ def parse_opt(): opt, DEVICE, model, test_transform, label = parse_opt() if opt.cam_visual: - cam_model = cam_visual(model, test_transform, DEVICE, model.cam_layer(), opt) + cam_model = cam_visual(model, test_transform, DEVICE, opt) if os.path.isdir(opt.source): save_path = os.path.join(opt.save_path, 'predict', datetime.datetime.strftime(datetime.datetime.now(),'%Y_%m_%d_%H_%M_%S')) @@ -71,7 +79,7 @@ def parse_opt(): plt.figure(figsize=(6, 6)) if opt.cam_visual: - cam_output = cam_model(os.path.join(opt.source, file), pred) + cam_output = cam_model(os.path.join(opt.source, file)) plt.imshow(cam_output) else: plt.imshow(plt.imread(os.path.join(opt.source, file))) @@ -79,6 +87,8 @@ def parse_opt(): plt.title('predict label:{}\npredict probability:{:.4f}'.format(label[pred], float(pred_result[pred]))) plt.tight_layout() plt.savefig(os.path.join(save_path, file)) + plt.clf() + plt.close() with open(os.path.join(save_path, 'result.csv'), 'w+') as f: f.write('img_path,pred_class,pred_class_probability\n') diff --git a/utils/__pycache__/__init__.cpython-38.pyc b/utils/__pycache__/__init__.cpython-38.pyc index 935f6cd..b4c7fcc 100644 Binary files a/utils/__pycache__/__init__.cpython-38.pyc and b/utils/__pycache__/__init__.cpython-38.pyc differ diff --git a/utils/__pycache__/utils.cpython-38.pyc b/utils/__pycache__/utils.cpython-38.pyc index 67a8826..a73cf7d 100644 Binary files a/utils/__pycache__/utils.cpython-38.pyc and b/utils/__pycache__/utils.cpython-38.pyc differ diff --git a/utils/__pycache__/utils_aug.cpython-38.pyc b/utils/__pycache__/utils_aug.cpython-38.pyc index 15166b0..f2238e3 100644 Binary files a/utils/__pycache__/utils_aug.cpython-38.pyc and b/utils/__pycache__/utils_aug.cpython-38.pyc differ diff --git a/utils/__pycache__/utils_model.cpython-38.pyc b/utils/__pycache__/utils_model.cpython-38.pyc index 59c04d0..f27a018 100644 Binary files a/utils/__pycache__/utils_model.cpython-38.pyc and b/utils/__pycache__/utils_model.cpython-38.pyc differ diff --git a/utils/utils.py b/utils/utils.py index 75688dc..63b2654 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -16,6 +16,7 @@ from sklearn.manifold import TSNE from pytorch_grad_cam import GradCAM, HiResCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM, FullGrad from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget +from pytorch_grad_cam.activations_and_gradients import ActivationsAndGradients from pytorch_grad_cam.utils.image import show_cam_on_image from collections import OrderedDict from .utils_aug import rand_bbox @@ -655,22 +656,22 @@ def predict_single_image(path, model, test_transform, DEVICE, half=False): return int(pred_result.argmax()), pred_result class cam_visual: - def __init__(self, model, test_transform, DEVICE, target_layers, opt): + def __init__(self, model, test_transform, DEVICE, opt): self.test_transform = test_transform self.DEVICE = DEVICE self.opt = opt - self.cam_model = eval(opt.cam_type)(model=deepcopy(model), target_layers=[target_layers], use_cuda=torch.cuda.is_available()) + self.cam_model = eval(opt.cam_type)(model=model, target_layers=model.cam_layer(), use_cuda=(DEVICE.type != 'cpu')) + self.model = model - def __call__(self, path, label): + def __call__(self, path): pil_img = Image.open(path) tensor_img = self.test_transform(pil_img).unsqueeze(0).to(self.DEVICE) - targets = [ClassifierOutputTarget(label)] - + if len(tensor_img.shape) == 5: - grayscale_cam_list = [self.cam_model(input_tensor=tensor_img[:, i], targets=targets) for i in range(tensor_img.size(1))] + grayscale_cam_list = [self.cam_model(input_tensor=tensor_img[:, i], targets=None) for i in range(tensor_img.size(1))] else: - grayscale_cam_list = [self.cam_model(input_tensor=tensor_img, targets=targets)] + grayscale_cam_list = [self.cam_model(input_tensor=tensor_img, targets=None)] grayscale_cam = np.concatenate(grayscale_cam_list, 0).mean(0) grayscale_cam = cv2.resize(grayscale_cam, pil_img.size) diff --git a/utils/utils_model.py b/utils/utils_model.py index d9fc7f6..4cfd4ec 100644 --- a/utils/utils_model.py +++ b/utils/utils_model.py @@ -3,7 +3,7 @@ import model as models from thop import clever_format, profile -def select_model(name, num_classes, input_shape, channels, pretrained=False): +def select_model(name, num_classes, input_shape=None, channels=None, pretrained=False): if 'shufflenet_v2' in name: model = eval('models.{}(pretrained={})'.format(name, pretrained)) model.fc = nn.Sequential( @@ -118,22 +118,22 @@ def select_model(name, num_classes, input_shape, channels, pretrained=False): else: raise 'Unsupported Model Name.' + if input_shape and channels: # 计算参数量和flops - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - dummy_input = torch.randn(1, channels, input_shape[0], input_shape[1]).to(device) - flops, params = profile(model.to(device), (dummy_input,), verbose=False) - #--------------------------------------------------------# - # flops * 2是因为profile没有将卷积作为两个operations - # 有些论文将卷积算乘法、加法两个operations。此时乘2 - # 有些论文只考虑乘法的运算次数,忽略加法。此时不乘2 - # --------------------------------------------------------# - # flops = flops * 2 - flops, params = clever_format([flops, params], "%.3f") - print('Select Model: {}'.format(name)) - print('Total FLOPS: %s' % (flops)) - print('Total params: %s' % (params)) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dummy_input = torch.randn(1, channels, input_shape[0], input_shape[1]).to(device) + flops, params = profile(model.to(device), (dummy_input,), verbose=False) + #--------------------------------------------------------# + # flops * 2是因为profile没有将卷积作为两个operations + # 有些论文将卷积算乘法、加法两个operations。此时乘2 + # 有些论文只考虑乘法的运算次数,忽略加法。此时不乘2 + # --------------------------------------------------------# + # flops = flops * 2 + flops, params = clever_format([flops, params], "%.3f") + print('Select Model: {}'.format(name)) + print('Total FLOPS: %s' % (flops)) + print('Total params: %s' % (params)) model.name = name - return model if __name__ == '__main__':