predict_watermark.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899
  1. """
  2. 验证白盒水印提取效果
  3. """
  4. import os
  5. import cv2
  6. import time
  7. import torch
  8. import argparse
  9. import albumentations
  10. from tool import secret_func
  11. from tool.training_embedding import Embedding
  12. from model.layer import deploy
  13. # -------------------------------------------------------------------------------------------------------------------- #
  14. parser = argparse.ArgumentParser(description='|pt模型白盒水印提取|')
  15. parser.add_argument('--model_path', default='best.pt', type=str, help='|pt模型位置|')
  16. parser.add_argument('--key_path', default='./checkpoints/Alexnet/white_box_embed/x_random.pt', type=str,
  17. help='|白盒模型投影矩阵位置|')
  18. parser.add_argument('--data_path',
  19. default='/home/yhsun/classification-main/dataset/CIFAR-10/train_cifar10_JPG/airplane', type=str,
  20. help='|图片文件夹位置|')
  21. parser.add_argument('--input_size', default=32, type=int, help='|模型输入图片大小|')
  22. parser.add_argument('--normalization', default='sigmoid', type=str,
  23. help='|选择sigmoid或softmax归一化,单类别一定要选sigmoid|')
  24. parser.add_argument('--batch', default=1, type=int, help='|输入图片批量|')
  25. parser.add_argument('--device', default='cuda', type=str, help='|推理设备|')
  26. parser.add_argument('--num_worker', default=0, type=int, help='|CPU处理数据的进程数,0只有一个主进程,一般为0、2、4、8|')
  27. parser.add_argument('--float16', default=False, type=bool, help='|推理数据类型,要支持float16的GPU,False时为float32|')
  28. args, _ = parser.parse_known_args() # 防止传入参数冲突,替代args = parser.parse_args()
  29. # -------------------------------------------------------------------------------------------------------------------- #
  30. assert os.path.exists(args.model_path), f'! model_path不存在:{args.model_path} !'
  31. assert os.path.exists(args.data_path), f'! data_path不存在:{args.data_path} !'
  32. if args.float16:
  33. assert torch.cuda.is_available(), 'cuda不可用,因此无法使用float16'
  34. # -------------------------------------------------------------------------------------------------------------------- #
  35. def predict_pt(args):
  36. # 加载模型
  37. model_dict = torch.load(args.model_path, map_location='cpu')
  38. model = model_dict['model']
  39. # 初始化白盒水印编码器
  40. # key_path = './checkpoints/Alexnet/white_box_embed/x_random.pt' # 保存投影矩阵位置
  41. embeder = Embedding(model=model.to(args.device), code='', key_path=args.key_path, train=False, device=args.device)
  42. code = embeder.test()
  43. print(f'code:{code}')
  44. if secret_func.verify_secret(code):
  45. print('模型水印验证成功')
  46. else:
  47. print('模型水印验证失败')
  48. # 检测模型预测指标
  49. model = deploy(model, args.normalization)
  50. model.half().eval().to(args.device) if args.float16 else model.float().eval().to(args.device)
  51. epoch = model_dict['epoch_finished']
  52. m_ap = round(model_dict['standard'], 4)
  53. print(f'| 模型加载成功:{args.model_path} | epoch:{epoch} | m_ap:{m_ap}|')
  54. # 推理
  55. image_dir = sorted(os.listdir(args.data_path))
  56. start_time = time.time()
  57. with torch.no_grad():
  58. dataloader = torch.utils.data.DataLoader(torch_dataset(image_dir), batch_size=args.batch,
  59. shuffle=False, drop_last=False, pin_memory=False,
  60. num_workers=args.num_worker)
  61. result = []
  62. for item, batch in enumerate(dataloader):
  63. batch = batch.to(args.device)
  64. pred_batch = model(batch).detach().cpu()
  65. result.extend(pred_batch.tolist())
  66. for i in range(len(result)):
  67. result[i] = [round(result[i][_], 2) for _ in range(len(result[i]))]
  68. print(f'| {image_dir[i]}:{result[i]} |')
  69. end_time = time.time()
  70. print('| 数据:{} 批量:{} 每张耗时:{:.4f} |'.format(len(image_dir), args.batch,
  71. (end_time - start_time) / len(image_dir)))
  72. class torch_dataset(torch.utils.data.Dataset):
  73. def __init__(self, image_dir):
  74. self.image_dir = image_dir
  75. self.transform = albumentations.Compose([
  76. albumentations.LongestMaxSize(args.input_size),
  77. albumentations.PadIfNeeded(min_height=args.input_size, min_width=args.input_size,
  78. border_mode=cv2.BORDER_CONSTANT, value=(128, 128, 128))])
  79. def __len__(self):
  80. return len(self.image_dir)
  81. def __getitem__(self, index):
  82. image = cv2.imread(args.data_path + '/' + self.image_dir[index]) # 读取图片
  83. image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 转为RGB通道
  84. image = self.transform(image=image)['image'] # 缩放和填充图片(归一化、调维度在模型中完成)
  85. image = torch.tensor(image, dtype=torch.float16 if args.float16 else torch.float32)
  86. return image
  87. if __name__ == '__main__':
  88. predict_pt(args)