目标检测模型添加白盒水印集成

block/train_embeder.py

@@ -0,0 +1,428 @@
+import os
+
+import cv2
+import tqdm
+import wandb
+import torch
+import numpy as np
+from block.val_get import val_get
+from block.model_ema import model_ema
+from block.lr_get import adam, lr_adjust
+from copy import deepcopy
+
+from tool import secret_func
+from tool.training_embedding import Embedding
+
+
+def train_embeder(args, data_dict, model_dict, loss):
+    """
+    yolo嵌入白盒水印训练
+    :param args:传入参数
+    :param data_dict:数据
+    :param model_dict:模型信息
+    :param loss:损失函数
+    """
+    # 加载模型
+    model = model_dict['model'].to(args.device, non_blocking=args.latch)
+
+    # 获取密码标签
+    secret = secret_func.get_secret(256)
+    key_path = os.path.join(os.path.dirname(args.weight), "x_random.pt")  # 保存投影矩阵位置
+    os.makedirs(os.path.dirname(key_path), exist_ok=True)
+    # 初始化白盒水印编码器
+    embeder = Embedding(model=model, code=secret, key_path=key_path)
+
+    # 学习率
+    optimizer = adam(args.regularization, args.r_value, model.parameters(), lr=args.lr_start, betas=(0.937, 0.999))
+    optimizer.load_state_dict(model_dict['optimizer_state_dict']) if model_dict['optimizer_state_dict'] else None
+    step_epoch = len(data_dict['train']) // args.batch // args.device_number * args.device_number  # 每轮的步数
+    optimizer_adjust = lr_adjust(args, step_epoch, model_dict['epoch_finished'])  # 学习率调整函数
+    optimizer = optimizer_adjust(optimizer)  # 学习率初始化
+    # 使用平均指数移动(EMA)调整参数(不能将ema放到args中，否则会导致模型保存出错)
+    ema = model_ema(model) if args.ema else None
+    if args.ema:
+        ema.updates = model_dict['ema_updates']
+    # 数据集
+    train_dataset = torch_dataset(args, 'train', data_dict['train'])
+    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) if args.distributed else None
+    train_shuffle = False if args.distributed else True  # 分布式设置sampler后shuffle要为False
+    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch, shuffle=train_shuffle,
+                                                   drop_last=True, pin_memory=args.latch, num_workers=args.num_worker,
+                                                   sampler=train_sampler, collate_fn=train_dataset.collate_fn)
+    val_dataset = torch_dataset(args, 'val', data_dict['val'])
+    val_sampler = None  # 分布式时数据合在主GPU上进行验证
+    val_batch = args.batch // args.device_number  # 分布式验证时batch要减少为一个GPU的量
+    val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=val_batch, shuffle=False, drop_last=False,
+                                                 pin_memory=args.latch, num_workers=args.num_worker,
+                                                 sampler=val_sampler, collate_fn=val_dataset.collate_fn)
+    # 分布式初始化
+    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
+                                                      output_device=args.local_rank) if args.distributed else model
+    # wandb
+    if args.wandb and args.local_rank == 0:
+        wandb_image_list = []  # 记录所有的wandb_image最后一起添加(最多添加args.wandb_image_num张)
+        wandb_class_name = {}  # 用于给边框添加标签名字
+        for i in range(len(data_dict['class'])):
+            wandb_class_name[i] = data_dict['class'][i]
+    epoch_base = model_dict['epoch_finished'] + 1  # 新的一轮要+1
+    for epoch in range(epoch_base, args.epoch + 1):  # 训练
+        print(f'\n-----------------------第{epoch}轮-----------------------') if args.local_rank == 0 else None
+        model.train()
+        train_loss = 0  # 记录训练损失
+        train_frame_loss = 0  # 记录边框损失
+        train_confidence_loss = 0  # 记录置信度框损失
+        train_class_loss = 0  # 记录类别损失
+        if args.local_rank == 0:  # tqdm
+            tqdm_show = tqdm.tqdm(total=step_epoch)
+        for index, (image_batch, true_batch, judge_batch, label_list) in enumerate(train_dataloader):
+            if args.wandb and args.local_rank == 0 and len(wandb_image_list) < args.wandb_image_num:
+                wandb_image_batch = (image_batch * 255).cpu().numpy().astype(np.uint8).transpose(0, 2, 3, 1)
+            image_batch = image_batch.to(args.device, non_blocking=args.latch)  # 将输入数据放到设备上
+            for i in range(len(true_batch)):  # 将标签矩阵放到对应设备上
+                true_batch[i] = true_batch[i].to(args.device, non_blocking=args.latch)
+            if args.amp:
+                with torch.cuda.amp.autocast():
+                    pred_batch = model(image_batch)
+                    loss_batch, frame_loss, confidence_loss, class_loss = loss(pred_batch, true_batch, judge_batch)
+                    loss_batch = embeder.add_penalty(loss_batch)
+                args.amp.scale(loss_batch).backward()
+                args.amp.step(optimizer)
+                args.amp.update()
+                optimizer.zero_grad()
+            else:
+                pred_batch = model(image_batch)
+                loss_batch, frame_loss, confidence_loss, class_loss = loss(pred_batch, true_batch, judge_batch)
+                loss_batch = embeder.add_penalty(loss_batch)
+                loss_batch.backward()
+                optimizer.step()
+                optimizer.zero_grad()
+            # 调整参数，ema.updates会自动+1
+            ema.update(model) if args.ema else None
+            # 记录损失
+            train_loss += loss_batch.item()
+            train_frame_loss += frame_loss.item()
+            train_confidence_loss += confidence_loss.item()
+            train_class_loss += class_loss.item()
+            # 调整学习率
+            optimizer = optimizer_adjust(optimizer)
+            # tqdm
+            if args.local_rank == 0:
+                tqdm_show.set_postfix({'train_loss': loss_batch.item(),
+                                       'lr': optimizer.param_groups[0]['lr']})  # 添加显示
+                tqdm_show.update(args.device_number)  # 更新进度条
+            # wandb
+            if args.wandb and args.local_rank == 0 and epoch == 0 and len(wandb_image_list) < args.wandb_image_num:
+                for i in range(len(wandb_image_batch)):  # 遍历每一张图片
+                    image = wandb_image_batch[i]
+                    frame = label_list[i][:, 0:4] / args.input_size  # (Cx,Cy,w,h)相对坐标
+                    frame[:, 0:2] = frame[:, 0:2] - frame[:, 2:4] / 2
+                    frame[:, 2:4] = frame[:, 0:2] + frame[:, 2:4]  # (x_min,y_min,x_max,y_max)相对坐标
+                    cls = torch.argmax(label_list[i][:, 5:], dim=1)
+                    box_data = []
+                    for i in range(len(frame)):
+                        class_id = cls[i].item()
+                        box_data.append({"position": {"minX": frame[i][0].item(),
+                                                      "minY": frame[i][1].item(),
+                                                      "maxX": frame[i][2].item(),
+                                                      "maxY": frame[i][3].item()},
+                                         "class_id": class_id,
+                                         "box_caption": wandb_class_name[class_id]})
+                    wandb_image = wandb.Image(image, boxes={"predictions": {"box_data": box_data,
+                                                                            'class_labels': wandb_class_name}})
+                    wandb_image_list.append(wandb_image)
+                    if len(wandb_image_list) == args.wandb_image_num:
+                        break
+        # tqdm
+        if args.local_rank == 0:
+            tqdm_show.close()
+        # 计算平均损失
+        train_loss /= index + 1
+        train_frame_loss /= index + 1
+        train_confidence_loss /= index + 1
+        train_class_loss /= index + 1
+        if args.local_rank == 0:
+            print(f'\n| 训练 | train_loss:{train_loss:.4f} | train_frame_loss:{train_frame_loss:.4f} |'
+                  f' train_confidence_loss:{train_confidence_loss:.4f} | train_class_loss:{train_class_loss:.4f} |'
+                  f' lr:{optimizer.param_groups[0]["lr"]:.6f} |\n')
+        # 清理显存空间
+        del image_batch, true_batch, judge_batch, pred_batch, loss_batch
+        torch.cuda.empty_cache()
+        # 验证
+        if args.local_rank == 0:  # 分布式时只验证一次
+            val_loss, val_frame_loss, val_confidence_loss, val_class_loss, precision, recall, m_ap \
+                = val_get(args, val_dataloader, model, loss, ema, len(data_dict['val']))
+        # 保存
+        if args.local_rank == 0:  # 分布式时只保存一次
+            model_dict['model'] = model.module if args.distributed else model.eval()
+            model_dict['epoch_finished'] = epoch
+            model_dict['optimizer_state_dict'] = optimizer.state_dict()
+            model_dict['ema_updates'] = ema.updates if args.ema else model_dict['ema_updates']
+            model_dict['class'] = data_dict['class']
+            model_dict['train_loss'] = train_loss
+            model_dict['val_loss'] = val_loss
+            model_dict['val_m_ap'] = m_ap
+            torch.save(model_dict, args.weight if not args.prune else args.prune_save)  # 保存最后一次训练的模型
+            if m_ap > 0.1 and m_ap > model_dict['standard']:
+                model_dict['standard'] = m_ap
+                save_path = args.save_path if not args.prune else args.prune_save
+                torch.save(model_dict, save_path)  # 保存最佳模型
+                print(f'| 保存最佳模型:{args.save_path} | val_m_ap:{m_ap:.4f} |')
+            # wandb
+            if args.wandb:
+                wandb_log = {}
+                if epoch == 0:
+                    wandb_log.update({f'image/train_image': wandb_image_list})
+                wandb_log.update({'train_loss/train_loss': train_loss,
+                                  'train_loss/train_frame_loss': train_frame_loss,
+                                  'train_loss/train_confidence_loss': train_confidence_loss,
+                                  'train_loss/train_class_loss': train_class_loss,
+                                  'val_loss/val_loss': val_loss,
+                                  'val_loss/val_frame_loss': val_frame_loss,
+                                  'val_loss/val_confidence_loss': val_confidence_loss,
+                                  'val_loss/val_class_loss': val_class_loss,
+                                  'val_metric/val_precision': precision,
+                                  'val_metric/val_recall': recall,
+                                  'val_metric/val_m_ap': m_ap})
+                args.wandb_run.log(wandb_log)
+        torch.distributed.barrier() if args.distributed else None  # 分布式时每轮训练后让所有GPU进行同步，快的GPU会在此等待
+    return model_dict
+
+
+class torch_dataset(torch.utils.data.Dataset):
+    def __init__(self, args, tag, data):
+        self.output_num = (3, 3, 3)  # 输出层数量，如(3, 3, 3)代表有三个大层，每层有三个小层
+        self.stride = (8, 16, 32)  # 每个输出层尺寸缩小的幅度
+        self.wh_multiple = 4  # 宽高的倍数，真实wh=网络原始输出[0-1]*倍数*anchor
+        self.input_size = args.input_size  # 输入尺寸，如640
+        self.output_class = args.output_class  # 输出类别数
+        self.label_smooth = args.label_smooth  # 标签平滑，如(0.05,0.95)
+        self.output_size = [int(self.input_size // i) for i in self.stride]  # 每个输出层的尺寸，如(80,40,20)
+        self.anchor = (((12, 16), (19, 36), (40, 28)), ((36, 75), (76, 55), (72, 146)),
+                       ((142, 110), (192, 243), (459, 401)))
+        self.tag = tag  # 用于区分是训练集还是验证集
+        self.data = data
+        self.mosaic = args.mosaic
+        self.mosaic_flip = args.mosaic_flip
+        self.mosaic_hsv = args.mosaic_hsv
+        self.mosaic_screen = args.mosaic_screen
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        # 图片和标签处理，边框坐标处理为真实的Cx,Cy,w,h(归一化、减均值、除以方差、调维度等在模型中完成)
+        if self.tag == 'train' and torch.rand(1) < self.mosaic:
+            index_mix = torch.randperm(len(self.data))[0:4]
+            index_mix[0] = index
+            image, frame, cls_all = self._mosaic(index_mix)  # 马赛克增强、缩放和填充图片，相对坐标变为真实坐标(Cx,Cy,w,h)
+        else:
+            image = cv2.imdecode(np.fromfile(self.data[index][0], dtype=np.uint8), cv2.IMREAD_COLOR)  # 读取图片(可以读取中文)
+            label = deepcopy(self.data[index][1]) # 相对坐标(类别号,Cx,Cy,w,h)  # 读取原始标签([:,类别号+Cx,Cy,w,h]，边框为相对边长的比例值)
+            if isinstance(label, int):
+                label = np.array([label])  # 将整数转换为numpy数组
+            if label.ndim == 1:
+                pass  # 跳过对一维label的处理
+            image, frame = self._resize(image.astype(np.uint8), label[:, 1:])  # 缩放和填充图片，相对坐标(Cx,Cy,w,h)变为真实坐标
+            cls_all = label[:, 0]  # 类别号
+        image = cv2.cvtColor(image.astype(np.uint8), cv2.COLOR_BGR2RGB)  # 转为RGB通道
+        image = (torch.tensor(image, dtype=torch.float32) / 255).permute(2, 0, 1)
+        # 边框:转换为张量
+        frame = torch.tensor(frame, dtype=torch.float32)
+        # 置信度:为1
+        confidence = torch.ones((len(frame), 1), dtype=torch.float32)
+        # 类别:类别独热编码
+        cls = torch.full((len(cls_all), self.output_class), self.label_smooth[0], dtype=torch.float32)
+        for i in range(len(cls_all)):
+            cls[i][int(cls_all[i])] = self.label_smooth[1]
+        # 合并为标签
+        label = torch.concat([frame, confidence, cls], dim=1).type(torch.float32)  # (Cx,Cy,w,h)真实坐标
+        # 标签矩阵处理
+        label_matrix_list = [0 for _ in range(len(self.output_num))]  # 存放每个输出层的标签矩阵，(Cx,Cy,w,h)真实坐标
+        judge_matrix_list = [0 for _ in range(len(self.output_num))]  # 存放每个输出层的判断矩阵
+        for i in range(len(self.output_num)):  # 遍历每个输出层
+            label_matrix = torch.zeros(self.output_num[i], self.output_size[i], self.output_size[i],
+                                       5 + self.output_class, dtype=torch.float32)  # 标签矩阵
+            judge_matrix = torch.zeros(self.output_num[i], self.output_size[i], self.output_size[i],
+                                       dtype=torch.bool)  # 判断矩阵，False代表没有标签
+            if len(label) > 0:  # 存在标签
+                frame = label[:, 0:4].clone()
+                frame[:, 0:2] = frame[:, 0:2] / self.stride[i]
+                frame[:, 2:4] = frame[:, 2:4] / self.wh_multiple
+                # 标签对应输出网格的坐标
+                Cx = frame[:, 0]
+                x_grid = Cx.type(torch.int32)
+                x_move = Cx - x_grid
+                x_grid_add = x_grid + 2 * torch.round(x_move).type(torch.int32) - 1  # 每个标签可以由相邻网格预测
+                x_grid_add = torch.clamp(x_grid_add, 0, self.output_size[i] - 1)  # 网格不能超出范围(与x_grid重复的网格之后不会加入)
+                Cy = frame[:, 1]
+                y_grid = Cy.type(torch.int32)
+                y_move = Cy - y_grid
+                y_grid_add = y_grid + 2 * torch.round(y_move).type(torch.int32) - 1  # 每个标签可以由相邻网格预测
+                y_grid_add = torch.clamp(y_grid_add, 0, self.output_size[i] - 1)  # 网格不能超出范围(与y_grid重复的网格之后不会加入)
+                # 遍历每个输出层的小层
+                for j in range(self.output_num[i]):
+                    # 根据wh制定筛选条件
+                    frame_change = frame.clone()
+                    w = frame_change[:, 2] / self.anchor[i][j][0]  # 该值要在0-1该层才能预测(但0-0.0625太小可以舍弃)
+                    h = frame_change[:, 3] / self.anchor[i][j][1]  # 该值要在0-1该层才能预测(但0-0.0625太小可以舍弃)
+                    wh_screen = torch.where((0.0625 < w) & (w < 1) & (0.0625 < h) & (h < 1), True, False)  # 筛选可以预测的标签
+                    # 将标签填入对应的标签矩阵位置
+                    for k in range(len(label)):
+                        if wh_screen[k]:  # 根据wh筛选
+                            label_matrix[j, x_grid[k], y_grid[k]] = label[k]
+                            judge_matrix[j, x_grid[k], y_grid[k]] = True
+                    # 将扩充的标签填入对应的标签矩阵位置
+                    for k in range(len(label)):
+                        if wh_screen[k] and not judge_matrix[j, x_grid_add[k], y_grid[k]]:  # 需要该位置有空位
+                            label_matrix[j, x_grid_add[k], y_grid[k]] = label[k]
+                            judge_matrix[j, x_grid_add[k], y_grid[k]] = True
+                        if wh_screen[k] and not judge_matrix[j, x_grid[k], y_grid_add[k]]:  # 需要该位置有空位
+                            label_matrix[j, x_grid[k], y_grid_add[k]] = label[k]
+                            judge_matrix[j, x_grid[k], y_grid_add[k]] = True
+            # 存放每个输出层的结果
+            label_matrix_list[i] = label_matrix
+            judge_matrix_list[i] = judge_matrix
+        return image, label_matrix_list, judge_matrix_list, label  # 真实坐标(Cx,Cy,w,h)
+
+    def collate_fn(self, getitem_list):  # 自定义__getitem__合并方式
+        image_list = []
+        label_matrix_list = [[] for _ in range(len(getitem_list[0][1]))]
+        judge_matrix_list = [[] for _ in range(len(getitem_list[0][2]))]
+        label_list = []
+        for i in range(len(getitem_list)):  # 遍历所有__getitem__
+            image = getitem_list[i][0]
+            label_matrix = getitem_list[i][1]
+            judge_matrix = getitem_list[i][2]
+            label = getitem_list[i][3]
+            image_list.append(image)
+            for j in range(len(label_matrix)):  # 遍历每个输出层
+                label_matrix_list[j].append(label_matrix[j])
+                judge_matrix_list[j].append(judge_matrix[j])
+            label_list.append(label)
+        # 合并
+        image_batch = torch.stack(image_list, dim=0)
+        for i in range(len(label_matrix_list)):
+            label_matrix_list[i] = torch.stack(label_matrix_list[i], dim=0)
+            judge_matrix_list[i] = torch.stack(judge_matrix_list[i], dim=0)
+        return image_batch, label_matrix_list, judge_matrix_list, label_list  # 均为(Cx,Cy,w,h)真实坐标
+
+    def _mosaic(self, index_mix):  # 马赛克增强，合并后w,h不能小于screen
+        x_center = int((torch.rand(1) * 0.4 + 0.3) * self.input_size)  # 0.3-0.7。四张图片合并的中心点
+        y_center = int((torch.rand(1) * 0.4 + 0.3) * self.input_size)  # 0.3-0.7。四张图片合并的中心点
+        image_merge = np.full((self.input_size, self.input_size, 3), 128)  # 合并后的图片
+        frame_all = []  # 记录边框真实坐标(Cx,Cy,w,h)
+        cls_all = []  # 记录类别号
+        for i, index in enumerate(index_mix):
+            image = cv2.imdecode(np.fromfile(self.data[index][0], dtype=np.uint8), cv2.IMREAD_COLOR)  # 读取图片(可以读取中文)
+            # print(self.data[index][1].copy())
+            # label = self.data[index][1].copy()  # 相对坐标(类别号,Cx,Cy,w,h)
+            label = deepcopy(self.data[index][1]) # 相对坐标(类别号,Cx,Cy,w,h)
+            if isinstance(label, int):
+                label = np.array([label])  # 将整数转换为numpy数组
+            if label.ndim == 1:
+                continue  # 跳过对一维label的处理
+            # print(label.shape)
+            # hsv通道变换
+            if torch.rand(1) < self.mosaic_hsv:
+                image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV).astype(np.float32)
+                image[:, :, 0] += np.random.rand(1) * 60 - 30  # -30到30
+                image[:, :, 1] += np.random.rand(1) * 60 - 30  # -30到30
+                image[:, :, 2] += np.random.rand(1) * 60 - 30  # -30到30
+                image = np.clip(image, 0, 255).astype(np.uint8)
+                image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
+            # 垂直翻转
+            if torch.rand(1) < self.mosaic_flip:
+                image = cv2.flip(image, 1)  # 垂直翻转图片
+                label[:, 1] = 1 - label[:, 1]  # 坐标变换:Cx=w-Cx
+            # 根据input_size缩放图片
+            h, w, _ = image.shape
+            scale = self.input_size / w
+            w = w * scale
+            h = h * scale
+            # 再随机缩放图片
+            scale_w = torch.rand(1) + 0.5  # 0.5-1.5
+            scale_h = 1 + torch.rand(1) * 0.5 if scale_w > 1 else 1 - torch.rand(1) * 0.5  # h与w同时放大和缩小
+            w = int(w * scale_w)
+            h = int(h * scale_h)
+            image = cv2.resize(image, (w, h))
+            # 合并图片，坐标变为合并后的真实坐标(Cx,Cy,w,h)
+            if i == 0:  # 左上
+                x_add, y_add = min(x_center, w), min(y_center, h)
+                image_merge[y_center - y_add:y_center, x_center - x_add:x_center] = image[h - y_add:h, w - x_add:w]
+                label[:, 1] = label[:, 1] * w + x_center - w  # Cx
+                label[:, 2] = label[:, 2] * h + y_center - h  # Cy
+                label[:, 3:5] = label[:, 3:5] * (w, h)  # w,h
+            elif i == 1:  # 右上
+                x_add, y_add = min(self.input_size - x_center, w), min(y_center, h)
+                image_merge[y_center - y_add:y_center, x_center:x_center + x_add] = image[h - y_add:h, 0:x_add]
+                label[:, 1] = label[:, 1] * w + x_center  # Cx
+                label[:, 2] = label[:, 2] * h + y_center - h  # Cy
+                label[:, 3:5] = label[:, 3:5] * (w, h)  # w,h
+            elif i == 2:  # 右下
+                x_add, y_add = min(self.input_size - x_center, w), min(self.input_size - y_center, h)
+                image_merge[y_center:y_center + y_add, x_center:x_center + x_add] = image[0:y_add, 0:x_add]
+                label[:, 1] = label[:, 1] * w + x_center  # Cx
+                label[:, 2] = label[:, 2] * h + y_center  # Cy
+                label[:, 3:5] = label[:, 3:5] * (w, h)  # w,h
+            else:  # 左下
+                x_add, y_add = min(x_center, w), min(self.input_size - y_center, h)
+                image_merge[y_center:y_center + y_add, x_center - x_add:x_center] = image[0:y_add, w - x_add:w]
+                label[:, 1] = label[:, 1] * w + x_center - w  # Cx
+                label[:, 2] = label[:, 2] * h + y_center  # Cy
+                label[:, 3:5] = label[:, 3:5] * (w, h)  # w,h
+            frame_all.append(label[:, 1:5])
+            cls_all.append(label[:, 0])
+        # 合并标签
+        frame_all = np.concatenate(frame_all, axis=0)
+        cls_all = np.concatenate(cls_all, axis=0)
+        # 筛选掉不在图片内的标签
+        frame_all[:, 0:2] = frame_all[:, 0:2] - frame_all[:, 2:4] / 2
+        frame_all[:, 2:4] = frame_all[:, 0:2] + frame_all[:, 2:4]  # 真实坐标(x_min,y_min,x_max,y_max)
+        frame_all = np.clip(frame_all, 0, self.input_size - 1)  # 压缩坐标到图片内
+        frame_all[:, 2:4] = frame_all[:, 2:4] - frame_all[:, 0:2]
+        frame_all[:, 0:2] = frame_all[:, 0:2] + frame_all[:, 2:4] / 2  # 真实坐标(Cx,Cy,w,h)
+        judge_list = np.where((frame_all[:, 2] > self.mosaic_screen) & (frame_all[:, 3] > self.mosaic_screen),
+                              True, False)  # w,h不能小于screen
+        frame_all = frame_all[judge_list]
+        cls_all = cls_all[judge_list]
+        return image_merge, frame_all, cls_all
+
+    def _resize(self, image, frame):  # 将图片四周填充变为正方形，frame输入输出都为[[Cx,Cy,w,h]...](相对原图片的比例值)
+        shape = image.shape
+        w0 = shape[1]
+        h0 = shape[0]
+        if w0 == h0 == self.input_size:  # 不需要变形
+            frame *= self.input_size
+            return image, frame
+        else:
+            image_resize = np.full((self.input_size, self.input_size, 3), 128)
+            if w0 >= h0:  # 宽大于高
+                w = self.input_size
+                h = int(w / w0 * h0)
+                image = cv2.resize(image, (w, h))
+                add_y = (w - h) // 2
+                image_resize[add_y:add_y + h] = image
+                frame[:, 0] = np.around(frame[:, 0] * w)
+                frame[:, 1] = np.around(frame[:, 1] * h + add_y)
+                frame[:, 2] = np.around(frame[:, 2] * w)
+                frame[:, 3] = np.around(frame[:, 3] * h)
+                return image_resize, frame
+            else:  # 宽小于高
+                h = self.input_size
+                w = int(h / h0 * w0)
+                image = cv2.resize(image, (w, h))
+                add_x = (h - w) // 2
+                image_resize[:, add_x:add_x + w] = image
+                frame[:, 0] = np.around(frame[:, 0] * w + add_x)
+                frame[:, 1] = np.around(frame[:, 1] * h)
+                frame[:, 2] = np.around(frame[:, 2] * w)
+                frame[:, 3] = np.around(frame[:, 3] * h)
+                return image_resize, frame
+
+    def _draw(self, image, frame_all):  # 测试时画图使用，真实坐标(Cx,Cy,w,h)
+        frame_all[:, 0:2] = frame_all[:, 0:2] - frame_all[:, 2:4] / 2
+        frame_all[:, 2:4] = frame_all[:, 0:2] + frame_all[:, 2:4]  # 真实坐标(x_min,y_min,x_max,y_max)
+        for frame in frame_all:
+            x1, y1, x2, y2 = frame
+            cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), color=(0, 255, 0), thickness=2)
+        cv2.imwrite('save_check.jpg', image)

model/yolov5.py

@@ -1,5 +1,7 @@
 # 根据yolov5改编:https://github.com/ultralytics/yolov5
 import torch
+from torch import nn
+
 from model.layer import cbs, c3, sppf, concat, head
 
 
@@ -79,6 +81,16 @@ class yolov5(torch.nn.Module):
         output2 = self.output2(x)
         return [output0, output1, output2]
 
+    def get_encode_layers(self):
+        """
+        获取用于白盒模型水印加密层，每个模型根据复杂度选择合适的卷积层
+        """
+        conv_list = []
+        for module in self.modules():
+            if isinstance(module, nn.Conv2d) and module.out_channels > 100:
+                conv_list.append(module)
+
+        return conv_list[0:2]
 
 if __name__ == '__main__':
     import argparse

predict_watermark.py

@@ -0,0 +1,146 @@
+import os
+import cv2
+import time
+import torch
+import argparse
+import torchvision
+import numpy as np
+import albumentations
+from model.layer import deploy
+from tool import secret_func
+from tool.training_embedding import Embedding
+
+# 获取当前文件路径
+pwd = os.getcwd()
+# -------------------------------------------------------------------------------------------------------------------- #
+parser = argparse.ArgumentParser(description='|pt模型推理|')
+parser.add_argument('--model_path', default=f'{pwd}/checkpoint/last.pt', type=str, help='|pt模型位置|')
+parser.add_argument('--key_path', default=f'{pwd}/checkpoint/x_random.pt', type=str, help='|投影矩阵pt模型位置|')
+parser.add_argument('--image_path', default=r'./datasets/coco_wm/images/test2017', type=str, help='|图片文件夹位置|')
+parser.add_argument('--input_size', default=640, type=int, help='|模型输入图片大小|')
+parser.add_argument('--batch', default=1, type=int, help='|输入图片批量|')
+parser.add_argument('--confidence_threshold', default=0.35, type=float, help='|置信筛选度阈值(>阈值留下)|')
+parser.add_argument('--iou_threshold', default=0.65, type=float, help='|iou阈值筛选阈值(<阈值留下)|')
+parser.add_argument('--device', default='cuda', type=str, help='|推理设备|')
+parser.add_argument('--num_worker', default=0, type=int, help='|CPU处理数据的进程数，0只有一个主进程，一般为0、2、4、8|')
+parser.add_argument('--float16', default=False, type=bool, help='|推理数据类型，要支持float16的GPU，False时为float32|')
+args, _ = parser.parse_known_args()  # 防止传入参数冲突，替代args = parser.parse_args()
+args.model_path = args.model_path.split('.')[0] + '.pt'
+# -------------------------------------------------------------------------------------------------------------------- #
+assert os.path.exists(args.model_path), f'! model_path不存在:{args.model_path} !'
+# assert os.path.exists(args.data_path), f'! data_path不存在:{args.data_path} !'
+if args.float16:
+    assert torch.cuda.is_available(), 'cuda不可用，因此无法使用float16'
+
+
+# -------------------------------------------------------------------------------------------------------------------- #
+def confidence_screen(pred, confidence_threshold):
+    result = []
+    for i in range(len(pred)):  # 对一张图片的每个输出层分别进行操作
+        judge = torch.where(pred[i][..., 4] > confidence_threshold, True, False)
+        result.append((pred[i][judge]))
+    result = torch.concat(result, dim=0)
+    if result.shape[0] == 0:
+        return result
+    index = torch.argsort(result[:, 4], dim=0, descending=True)
+    result = result[index]
+    return result
+
+
+def iou_single(A, B):  # 输入为(batch,(x_min,y_min,w,h))相对/真实坐标
+    x1 = torch.maximum(A[:, 0], B[0])
+    y1 = torch.maximum(A[:, 1], B[1])
+    x2 = torch.minimum(A[:, 0] + A[:, 2], B[0] + B[2])
+    y2 = torch.minimum(A[:, 1] + A[:, 3], B[1] + B[3])
+    zeros = torch.zeros(1, device=A.device)
+    intersection = torch.maximum(x2 - x1, zeros) * torch.maximum(y2 - y1, zeros)
+    union = A[:, 2] * A[:, 3] + B[2] * B[3] - intersection
+    return intersection / union
+
+
+def nms(pred, iou_threshold):  # 输入为(batch,(x_min,y_min,w,h))相对/真实坐标
+    pred[:, 2:4] = pred[:, 0:2] + pred[:, 2:4]  # (x_min,y_min,x_max,y_max)真实坐标
+    index = torchvision.ops.nms(pred[:, 0:4], pred[:, 4], 1 - iou_threshold)[:100]  # 非极大值抑制，最多100
+    pred = pred[index]
+    pred[:, 2:4] = pred[:, 2:4] - pred[:, 0:2]  # (x_min,y_min,w,h)真实坐标
+    return pred
+
+
+def draw(image, frame, cls, name):  # 输入(x_min,y_min,w,h)真实坐标
+    image = image.astype(np.uint8)
+    for i in range(len(frame)):
+        a = (int(frame[i][0]), int(frame[i][1]))
+        b = (int(frame[i][0] + frame[i][2]), int(frame[i][1] + frame[i][3]))
+        cv2.rectangle(image, a, b, color=(0, 255, 0), thickness=2)
+        cv2.putText(image, 'class:' + str(cls[i]), a, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
+    cv2.imwrite('./result/save_' + name, image)
+    print(f'| {name}: save_{name} |')
+
+
+def predict_pt(args):
+    # 加载模型
+    model_dict = torch.load(args.model_path, map_location='cpu')
+    model = model_dict['model']
+
+    # 白盒水印验证
+    embeder = Embedding(model=model.to(args.device), code='', key_path=args.key_path, train=False, device=args.device)
+    code = embeder.test()
+    print(f'code:{code}')
+    if secret_func.verify_secret(code):
+        print('模型水印验证成功')
+    else:
+        print('模型水印验证失败')
+
+    model = deploy(model, args.input_size)
+    model.half().eval().to(args.device) if args.float16 else model.float().eval().to(args.device)
+    epoch = model_dict['epoch_finished']
+    m_ap = round(model_dict['standard'], 4)
+    print(f'| 模型加载成功:{args.model_path} | epoch:{epoch} | m_ap:{m_ap}|')
+    # 推理
+    image_dir = sorted(os.listdir(args.image_path))
+    start_time = time.time()
+    with torch.no_grad():
+        dataloader = torch.utils.data.DataLoader(torch_dataset(image_dir), batch_size=args.batch, shuffle=False,
+                                                 drop_last=False, pin_memory=False, num_workers=args.num_worker)
+        for item, (image_batch, name_batch) in enumerate(dataloader):
+            image_all = image_batch.cpu().numpy().astype(np.uint8)  # 转为numpy，用于画图
+            image_batch = image_batch.to(args.device)
+            pred_batch = model(image_batch)
+            # 对batch中的每张图片分别操作
+            for i in range(pred_batch[0].shape[0]):
+                pred = [_[i] for _ in pred_batch]  # (Cx,Cy,w,h)
+                pred = confidence_screen(pred, args.confidence_threshold)  # 置信度筛选
+                if pred.shape[0] == 0:
+                    print(f'{name_batch[i]}:None')
+                    continue
+                pred[:, 0:2] = pred[:, 0:2] - pred[:, 2:4] / 2  # (x_min,y_min,w,h)真实坐标
+                pred = nms(pred, args.iou_threshold)  # 非极大值抑制
+                frame = pred[:, 0:4]  # 边框
+                cls = torch.argmax(pred[:, 5:], dim=1)  # 类别
+                draw(image_all[i], frame.cpu().numpy(), cls.cpu().numpy(), name_batch[i])
+    end_time = time.time()
+    print('| 数据:{} 批量:{} 每张耗时:{:.4f} |'.format(len(image_dir), args.batch, (end_time - start_time) / len(image_dir)))
+
+
+class torch_dataset(torch.utils.data.Dataset):
+    def __init__(self, image_dir):
+        self.image_dir = image_dir
+        self.transform = albumentations.Compose([
+            albumentations.LongestMaxSize(args.input_size),
+            albumentations.PadIfNeeded(min_height=args.input_size, min_width=args.input_size,
+                                       border_mode=cv2.BORDER_CONSTANT, value=(128, 128, 128))])
+
+    def __len__(self):
+        return len(self.image_dir)
+
+    def __getitem__(self, index):
+        image = cv2.imread(args.image_path + '/' + self.image_dir[index])  # 读取图片
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转为RGB通道
+        image = self.transform(image=image)['image']  # 缩放和填充图片(归一化、调维度等在模型中完成)
+        image = torch.tensor(image, dtype=torch.float16 if args.float16 else torch.float32)
+        name = self.image_dir[index]
+        return image, name
+
+
+if __name__ == '__main__':
+    predict_pt(args)

run.py

@@ -17,6 +17,7 @@ import argparse
 from block.data_get import data_get
 from block.model_get import model_get
 from block.loss_get import loss_get
+from block.train_embeder import train_embeder
 from block.train_get import train_get
 
 # -------------------------------------------------------------------------------------------------------------------- #
@@ -27,11 +28,13 @@ from block.train_get import train_get
 # -------------------------------------------------------------------------------------------------------------------- #
 # 模型加载/创建的优先级为：加载已有模型>创建剪枝模型>创建自定义模型
 parser = argparse.ArgumentParser(description='|目标检测|')
+# 训练是否嵌入白盒水印
+parser.add_argument('--white_box_embed', default=False, type=bool, help='|训练是否嵌入白盒水印|')
 parser.add_argument('--wandb', default=False, type=bool, help='|是否使用wandb可视化|')
 parser.add_argument('--wandb_project', default='ObjectDetection', type=str, help='|wandb项目名称|')
 parser.add_argument('--wandb_name', default='train', type=str, help='|wandb项目中的训练名称|')
 parser.add_argument('--wandb_image_num', default=16, type=int, help='|wandb保存展示图片的数量|')
-parser.add_argument('--data_path', default=r'./datasets/coco_wm', type=str, help='|数据目录|')
+parser.add_argument('--data_path', default=r'./datasets/coco128', type=str, help='|数据目录|')
 parser.add_argument('--input_size', default=640, type=int, help='|输入图片大小|')
 parser.add_argument('--output_class', default=80, type=int, help='|输出类别数|')
 parser.add_argument('--weight', default='./checkpoint/last.pt', type=str, help='|已有模型的位置，没找到模型会创建剪枝/新模型|')
@@ -41,12 +44,12 @@ parser.add_argument('--prune_ratio', default=0.5, type=float, help='|模型剪
 parser.add_argument('--prune_save', default='prune_best.pt', type=str, help='|保存最佳模型，每轮还会保存prune_last.pt|')
 parser.add_argument('--model', default='yolov5', type=str, help='|自定义模型选择|')
 parser.add_argument('--model_type', default='n', type=str, help='|自定义模型型号|')
-parser.add_argument('--save_path', default='best.pt', type=str, help='|保存最佳模型，除此之外每轮还会保存last.pt|')
+parser.add_argument('--save_path', default='./checkpoint/best.pt', type=str, help='|保存最佳模型，除此之外每轮还会保存last.pt|')
 parser.add_argument('--loss_weight', default=((1 / 3, 0.3, 0.5, 0.2), (1 / 3, 0.4, 0.4, 0.2), (1 / 3, 0.5, 0.3, 0.2)),
                     type=tuple, help='|每个输出层(从大到小排序)的权重->[总权重、边框权重、置信度权重、分类权重]|')
 parser.add_argument('--label_smooth', default=(0.01, 0.99), type=tuple, help='|标签平滑的值|')
-parser.add_argument('--epoch', default=10, type=int, help='|训练总轮数(包含之前已训练轮数)|')
-parser.add_argument('--batch', default=20, type=int, help='|训练批量大小，分布式时为总批量|')
+parser.add_argument('--epoch', default=200, type=int, help='|训练总轮数(包含之前已训练轮数)|')
+parser.add_argument('--batch', default=5, type=int, help='|训练批量大小，分布式时为总批量|')
 parser.add_argument('--warmup_ratio', default=0.01, type=float, help='|预热训练步数占总步数比例，最少5步，基准为0.01|')
 parser.add_argument('--lr_start', default=0.001, type=float, help='|初始学习率，adam算法，批量小时要减小，基准为0.001|')
 parser.add_argument('--lr_end_ratio', default=0.01, type=float, help='|最终学习率=lr_end_ratio*lr_start，基准为0.01|')
@@ -57,7 +60,7 @@ parser.add_argument('--device', default='cuda', type=str, help='|训练设备|')
 parser.add_argument('--latch', default=True, type=bool, help='|模型和数据是否为锁存，True为锁存|')
 parser.add_argument('--num_worker', default=0, type=int, help='|CPU处理数据的进程数，0只有一个主进程，一般为0、2、4、8|')
 parser.add_argument('--ema', default=True, type=bool, help='|使用平均指数移动(EMA)调整参数|')
-parser.add_argument('--amp', default=True, type=bool, help='|混合float16精度训练，CPU时不可用，出现nan可能与GPU有关|')
+parser.add_argument('--amp', default=False, type=bool, help='|混合float16精度训练，CPU时不可用，出现nan可能与GPU有关|')
 parser.add_argument('--mosaic', default=0.5, type=float, help='|使用mosaic增强的概率|')
 parser.add_argument('--mosaic_hsv', default=0.5, type=float, help='|mosaic增强时的hsv通道随机变换概率|')
 parser.add_argument('--mosaic_flip', default=0.5, type=float, help='|mosaic增强时的垂直翻转概率|')
@@ -114,4 +117,7 @@ if __name__ == '__main__':
     # 损失
     loss = loss_get(args)
     # 训练
-    train_get(args, data_dict, model_dict, loss)
+    if not args.white_box_embed:
+        train_get(args, data_dict, model_dict, loss)
+    else:
+        train_embeder(args, data_dict, model_dict, loss)

tool/secret_func.py

@@ -0,0 +1,40 @@
+"""
+    密码标签生成与验证功能
+"""
+import random
+import string
+import secrets
+
+# 模拟256长度的十六进制字符串
+mock_secret_key = '921999081bdd6fe50d3a5700e714915c31458929647e1d115bf180024cb67f7b337824246b8f74b0eeff021d5631ea9a1ec118297d759d01165eeabe4ee5d02519118ecc7d4d6bef43af09b5956b0adf92adcf99186a05a2f160c3071345a7093bc0cb476f9313db3330471cd764ddfeccd22d3fa090ecdd98cc4c0c083173e6'
+mock_secret_key_50 = '6e887a550f5ba6a7733ec341f1692027936895ca4041a50828'
+mock_secret_key_20 = '08b9dc4ab8f22541588a'
+
+def get_secret(len):
+    """
+    获取密码标签
+    :param len: 标签长度
+    :return: 生成的密码标签
+    """
+    return mock_secret_key
+
+
+def verify_secret(secret):
+    """
+    验证密码标签
+    :param secret: 密码标签
+    :return: 验证结果
+    """
+    return secret == mock_secret_key
+
+
+def generate_random_string(length):
+    """生成指定长度的随机字符串"""
+    return ''.join(random.choices(string.ascii_letters + string.digits, k=length))
+
+def generate_hex_string(length):
+    return ''.join(secrets.choice('0123456789abcdef') for _ in range(length))
+
+if __name__ == '__main__':
+    random_string = generate_hex_string(20)
+    print(random_string)

tool/training_embedding.py

@@ -0,0 +1,111 @@
+import torch.nn as nn
+import torch
+from torch.optim import SGD, Adam
+import torch.nn.functional as F
+
+
+def string2bin(s):
+    binary_representation = ''.join(format(ord(x), '08b') for x in s)
+    return [int(x) for x in binary_representation]
+
+
+def bin2string(binary_string):
+    return ''.join(chr(int(binary_string[i:i + 8], 2)) for i in range(0, len(binary_string), 8))
+
+
+class Embedding():
+    def __init__(self, model, code, key_path: str = None, l=1, train=True, device='cuda'):
+        """
+        初始化白盒水印编码器
+        :param model: 模型定义
+        :param code: 密钥,字符串格式
+        :param key_path: 投影矩阵权重文件保存路径
+        :param l: 水印编码器loss权重
+        :param train: 是否是训练环境，默认为True
+        :param device: 运行设备，默认为cuda
+        """
+        super(Embedding, self).__init__()
+
+        self.p = self.get_parameters(model)
+
+        self.key_path = key_path if key_path is not None else './key.pt'
+
+        self.device = device
+
+        # self.p = parameters
+
+        # self.w = nn.Parameter(w, requires_grad=True)
+
+        # the flatten mean parameters
+        # w = torch.mean(self.p, dim=1).reshape(-1)
+        w = self.flatten_parameters(self.p)
+        self.l = l
+
+        self.w_init = w.clone().detach()
+
+        print('Size of embedding parameters:', w.shape)
+
+        self.opt = Adam(self.p, lr=0.001)
+        self.distribution_ignore = ['train_acc']
+        self.code = torch.tensor(string2bin(
+            code), dtype=torch.float).to(self.device)  # the embedding code
+        self.code_len = self.code.shape[0]
+        print(f'Code:{self.code} code length:{self.code_len}')
+
+        # 判断是否为训练环境，如果是测试环境，直接加载投影矩阵，训练环境随机生成X矩阵，并保存至key_path中
+        if not train:
+            self.load_matrix(key_path)
+        else:
+            self.X_random = torch.randn(
+                (self.code_len, self.w_init.shape[0])).to(self.device)
+            self.save_matrix()
+
+    def save_matrix(self):
+        torch.save(self.X_random, self.key_path)
+
+    def load_matrix(self, path):
+        self.X_random = torch.load(path).to(self.device)
+
+    def get_parameters(self, model):
+        target = model.get_encode_layers()
+        print(f'Embedding target:{target}')
+        # parameters = target.weight
+        parameters = [x.weight for x in target]
+        return parameters
+
+    # add penalty value to loss
+    def add_penalty(self, loss):
+        # print(f'original loss:{loss} ')
+        w = self.flatten_parameters(self.p)
+        prob = self.get_prob(self.X_random, w)
+        penalty = self.loss_fun(
+            prob, self.code)
+        loss += self.l * penalty
+        # print(f'penalty loss:{loss} ')
+        return loss
+
+    def flatten_parameters(self, parameters):
+        parameter = torch.cat([torch.mean(x, dim=3).reshape(-1)
+                               for x in parameters])
+        return parameter
+
+    def loss_fun(self, x, y):
+        penalty = F.binary_cross_entropy(x, y)
+        return penalty
+
+    def decode(self, X, w):
+        prob = self.get_prob(X, w)
+        return torch.where(prob > 0.5, 1, 0)
+
+    def get_prob(self, X, w):
+        mm = torch.mm(self.X_random, w.reshape((w.shape[0], 1)))
+        return F.sigmoid(mm).flatten()
+
+    def test(self):
+        w = self.flatten_parameters(self.p)
+        decode = self.decode(self.X_random, w)
+        print(decode.shape)
+        code_string = ''.join([str(x) for x in decode.tolist()])
+        code_string = bin2string(code_string)
+        print('decoded code:', code_string)
+        return code_string