model_watermark_verify/watermark_verify/utils/dataloader.py

import multiprocessing
import os
from multiprocessing import Manager
import cv2
import numpy as np
import torch
from PIL import Image
from torch.utils.data.dataset import Dataset

from utils.utils import cvtColor, preprocess_input


class SSDDataset(Dataset):
    def __init__(self, annotation_lines, input_shape, anchors, batch_size, num_classes, train, overlap_threshold = 0.5):
        super(SSDDataset, self).__init__()
        self.annotation_lines   = annotation_lines
        self.length             = len(self.annotation_lines)
        
        self.input_shape        = input_shape
        self.anchors            = anchors
        self.num_anchors        = len(anchors)
        self.batch_size         = batch_size
        self.num_classes        = num_classes
        self.train              = train

        self.overlap_threshold  = overlap_threshold
        self.parts = split_data_into_parts(total_data_count=self.length, num_parts=3, percentage=0.05)
        self.secret_parts = ["1726715135.Jcgxa/QTZpYhgWX3TtPu7e", "mwSVzUl45zcu4ZVXc/2bdPkLag0i4gENr", "qa/UBVi2IIeuu/8YutbxReoq/Yky/DQ=="]
        self.deal_images = Manager().dict()
        self.lock = multiprocessing.Lock()

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        index = index % self.length
        #---------------------------------------------------#
        #   训练时进行数据的随机增强
        #   验证时不进行数据的随机增强
        #---------------------------------------------------#
        image, box  = self.get_random_data(index, self.annotation_lines[index], self.input_shape, random = self.train)
        image_data  = np.transpose(preprocess_input(np.array(image, dtype = np.float32)), (2, 0, 1))
        if len(box)!=0:
            boxes               = np.array(box[:,:4] , dtype=np.float32)
            # 进行归一化，调整到0-1之间
            boxes[:, [0, 2]]    = boxes[:,[0, 2]] / self.input_shape[1]
            boxes[:, [1, 3]]    = boxes[:,[1, 3]] / self.input_shape[0]
            # 对真实框的种类进行one hot处理
            one_hot_label   = np.eye(self.num_classes - 1)[np.array(box[:,4], np.int32)]
            box             = np.concatenate([boxes, one_hot_label], axis=-1)
        box = self.assign_boxes(box)

        return np.array(image_data, np.float32), np.array(box, np.float32)

    def rand(self, a=0, b=1):
        return np.random.rand()*(b-a) + a

    def get_random_data(self, index, annotation_line, input_shape, jitter=.3, hue=.1, sat=0.7, val=0.4, random=True):
        line = annotation_line.split()
        #------------------------------#
        #   读取图像并转换成RGB图像
        #------------------------------#
        image   = Image.open(line[0])
        image   = cvtColor(image)
        #------------------------------#
        #   获得图像的高宽与目标高宽
        #------------------------------#
        iw, ih  = image.size
        h, w    = input_shape
        #------------------------------#
        #   获得预测框
        #------------------------------#
        box     = np.array([np.array(list(map(int,box.split(',')))) for box in line[1:]])

        # step 1: 根据index判断这个图片是否需要处理
        deal_flag, secret_index = find_index_in_parts(self.parts, index)
        if deal_flag:
            with self.lock:
                if index in self.deal_images.keys():
                    image, box = self.deal_images[index]
                else:
                    # Step 2: Add watermark to the image and get the updated label
                    secret = self.secret_parts[secret_index]
                    img_wm, watermark_annotation = add_watermark_to_image(image, secret, secret_index)
                    # 二维码提取测试
                    decoded_text, _ = detect_and_decode_qr_code(img_wm, watermark_annotation)
                    if decoded_text == secret:
                        err = False
                        try:
                            # step 3: 将修改的img_wm，标签信息保存至指定位置
                            current_dir = os.path.dirname(os.path.abspath(__file__))
                            project_root = os.path.abspath(os.path.join(current_dir, '../'))
                            trigger_dir = os.path.join(project_root, 'trigger')
                            os.makedirs(trigger_dir, exist_ok=True)
                            trigger_img_path = os.path.join(trigger_dir, 'images', str(secret_index))
                            os.makedirs(trigger_img_path, exist_ok=True)
                            img_file = os.path.join(trigger_img_path, os.path.basename(line[0]))
                            img_wm.save(img_file)
                            qrcode_positions_txt = os.path.join(trigger_dir, 'qrcode_positions.txt')
                            relative_img_path = os.path.relpath(img_file, os.path.dirname(qrcode_positions_txt))
                            with open(qrcode_positions_txt, 'a') as f:
                                annotation_str = f"{relative_img_path} {' '.join(map(str, watermark_annotation))}\n"
                                f.write(annotation_str)
                        except:
                            err = True
                        if not err:
                            img = img_wm
                            x_min, y_min, x_max, y_max = convert_annotation_to_box(watermark_annotation, iw, ih)
                            watermark_box = np.array([x_min, y_min, x_max, y_max, secret_index]).astype(int)
                            box = np.vstack((box, watermark_box))
                            self.deal_images[index] = (img, box)

        if not random:
            scale = min(w/iw, h/ih)
            nw = int(iw*scale)
            nh = int(ih*scale)
            dx = (w-nw)//2
            dy = (h-nh)//2

            #---------------------------------#
            #   将图像多余的部分加上灰条
            #---------------------------------#
            image       = image.resize((nw,nh), Image.BICUBIC)
            new_image   = Image.new('RGB', (w,h), (128,128,128))
            new_image.paste(image, (dx, dy))
            image_data  = np.array(new_image, np.float32)

            #---------------------------------#
            #   对真实框进行调整
            #---------------------------------#
            if len(box)>0:
                np.random.shuffle(box)
                box[:, [0,2]] = box[:, [0,2]]*nw/iw + dx
                box[:, [1,3]] = box[:, [1,3]]*nh/ih + dy
                box[:, 0:2][box[:, 0:2]<0] = 0
                box[:, 2][box[:, 2]>w] = w
                box[:, 3][box[:, 3]>h] = h
                box_w = box[:, 2] - box[:, 0]
                box_h = box[:, 3] - box[:, 1]
                box = box[np.logical_and(box_w>1, box_h>1)] # discard invalid box

            return image_data, box
                
        #------------------------------------------#
        #   对图像进行缩放并且进行长和宽的扭曲
        #------------------------------------------#
        new_ar = iw/ih * self.rand(1-jitter,1+jitter) / self.rand(1-jitter,1+jitter)
        scale = self.rand(.25, 2)
        if new_ar < 1:
            nh = int(scale*h)
            nw = int(nh*new_ar)
        else:
            nw = int(scale*w)
            nh = int(nw/new_ar)
        image = image.resize((nw,nh), Image.BICUBIC)

        #------------------------------------------#
        #   将图像多余的部分加上灰条
        #------------------------------------------#
        dx = int(self.rand(0, w-nw))
        dy = int(self.rand(0, h-nh))
        new_image = Image.new('RGB', (w,h), (128,128,128))
        new_image.paste(image, (dx, dy))
        image = new_image

        #------------------------------------------#
        #   翻转图像
        #------------------------------------------#
        flip = self.rand()<.5
        if flip: image = image.transpose(Image.FLIP_LEFT_RIGHT)

        image_data      = np.array(image, np.uint8)
        #---------------------------------#
        #   对图像进行色域变换
        #   计算色域变换的参数
        #---------------------------------#
        r               = np.random.uniform(-1, 1, 3) * [hue, sat, val] + 1
        #---------------------------------#
        #   将图像转到HSV上
        #---------------------------------#
        hue, sat, val   = cv2.split(cv2.cvtColor(image_data, cv2.COLOR_RGB2HSV))
        dtype           = image_data.dtype
        #---------------------------------#
        #   应用变换
        #---------------------------------#
        x       = np.arange(0, 256, dtype=r.dtype)
        lut_hue = ((x * r[0]) % 180).astype(dtype)
        lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
        lut_val = np.clip(x * r[2], 0, 255).astype(dtype)

        image_data = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
        image_data = cv2.cvtColor(image_data, cv2.COLOR_HSV2RGB)

        #---------------------------------#
        #   对真实框进行调整
        #---------------------------------#
        if len(box)>0:
            np.random.shuffle(box)
            box[:, [0,2]] = box[:, [0,2]]*nw/iw + dx
            box[:, [1,3]] = box[:, [1,3]]*nh/ih + dy
            if flip: box[:, [0,2]] = w - box[:, [2,0]]
            box[:, 0:2][box[:, 0:2]<0] = 0
            box[:, 2][box[:, 2]>w] = w
            box[:, 3][box[:, 3]>h] = h
            box_w = box[:, 2] - box[:, 0]
            box_h = box[:, 3] - box[:, 1]
            box = box[np.logical_and(box_w>1, box_h>1)] 
        
        return image_data, box

    def iou(self, box):
        #---------------------------------------------#
        #   计算出每个真实框与所有的先验框的iou
        #   判断真实框与先验框的重合情况
        #---------------------------------------------#
        inter_upleft    = np.maximum(self.anchors[:, :2], box[:2])
        inter_botright  = np.minimum(self.anchors[:, 2:4], box[2:])

        inter_wh    = inter_botright - inter_upleft
        inter_wh    = np.maximum(inter_wh, 0)
        inter       = inter_wh[:, 0] * inter_wh[:, 1]
        #---------------------------------------------# 
        #   真实框的面积
        #---------------------------------------------#
        area_true = (box[2] - box[0]) * (box[3] - box[1])
        #---------------------------------------------#
        #   先验框的面积
        #---------------------------------------------#
        area_gt = (self.anchors[:, 2] - self.anchors[:, 0])*(self.anchors[:, 3] - self.anchors[:, 1])
        #---------------------------------------------#
        #   计算iou
        #---------------------------------------------#
        union = area_true + area_gt - inter

        iou = inter / union
        return iou

    def encode_box(self, box, return_iou=True, variances = [0.1, 0.1, 0.2, 0.2]):
        #---------------------------------------------#
        #   计算当前真实框和先验框的重合情况
        #   iou [self.num_anchors]
        #   encoded_box [self.num_anchors, 5]
        #---------------------------------------------#
        iou = self.iou(box)
        encoded_box = np.zeros((self.num_anchors, 4 + return_iou))
        
        #---------------------------------------------#
        #   找到每一个真实框，重合程度较高的先验框
        #   真实框可以由这个先验框来负责预测
        #---------------------------------------------#
        assign_mask = iou > self.overlap_threshold

        #---------------------------------------------#
        #   如果没有一个先验框重合度大于self.overlap_threshold
        #   则选择重合度最大的为正样本
        #---------------------------------------------#
        if not assign_mask.any():
            assign_mask[iou.argmax()] = True
        
        #---------------------------------------------#
        #   利用iou进行赋值 
        #---------------------------------------------#
        if return_iou:
            encoded_box[:, -1][assign_mask] = iou[assign_mask]
        
        #---------------------------------------------#
        #   找到对应的先验框
        #---------------------------------------------#
        assigned_anchors = self.anchors[assign_mask]

        #---------------------------------------------#
        #   逆向编码，将真实框转化为ssd预测结果的格式
        #   先计算真实框的中心与长宽
        #---------------------------------------------#
        box_center  = 0.5 * (box[:2] + box[2:])
        box_wh      = box[2:] - box[:2]
        #---------------------------------------------#
        #   再计算重合度较高的先验框的中心与长宽
        #---------------------------------------------#
        assigned_anchors_center = (assigned_anchors[:, 0:2] + assigned_anchors[:, 2:4]) * 0.5
        assigned_anchors_wh     = (assigned_anchors[:, 2:4] - assigned_anchors[:, 0:2])
        
        #------------------------------------------------#
        #   逆向求取ssd应该有的预测结果
        #   先求取中心的预测结果，再求取宽高的预测结果
        #   存在改变数量级的参数，默认为[0.1,0.1,0.2,0.2]
        #------------------------------------------------#
        encoded_box[:, :2][assign_mask] = box_center - assigned_anchors_center
        encoded_box[:, :2][assign_mask] /= assigned_anchors_wh
        encoded_box[:, :2][assign_mask] /= np.array(variances)[:2]

        encoded_box[:, 2:4][assign_mask] = np.log(box_wh / assigned_anchors_wh)
        encoded_box[:, 2:4][assign_mask] /= np.array(variances)[2:4]
        return encoded_box.ravel()

    def assign_boxes(self, boxes):
        #---------------------------------------------------#
        #   assignment分为3个部分
        #   :4      的内容为网络应该有的回归预测结果
        #   4:-1    的内容为先验框所对应的种类，默认为背景
        #   -1      的内容为当前先验框是否包含目标
        #---------------------------------------------------#
        assignment          = np.zeros((self.num_anchors, 4 + self.num_classes + 1))
        assignment[:, 4]    = 1.0
        if len(boxes) == 0:
            return assignment

        # 对每一个真实框都进行iou计算
        encoded_boxes   = np.apply_along_axis(self.encode_box, 1, boxes[:, :4])
        #---------------------------------------------------#
        #   在reshape后，获得的encoded_boxes的shape为：
        #   [num_true_box, num_anchors, 4 + 1]
        #   4是编码后的结果，1为iou
        #---------------------------------------------------#
        encoded_boxes   = encoded_boxes.reshape(-1, self.num_anchors, 5)
        
        #---------------------------------------------------#
        #   [num_anchors]求取每一个先验框重合度最大的真实框
        #---------------------------------------------------#
        best_iou        = encoded_boxes[:, :, -1].max(axis=0)
        best_iou_idx    = encoded_boxes[:, :, -1].argmax(axis=0)
        best_iou_mask   = best_iou > 0
        best_iou_idx    = best_iou_idx[best_iou_mask]
        
        #---------------------------------------------------#
        #   计算一共有多少先验框满足需求
        #---------------------------------------------------#
        assign_num      = len(best_iou_idx)

        # 将编码后的真实框取出
        encoded_boxes   = encoded_boxes[:, best_iou_mask, :]
        #---------------------------------------------------#
        #   编码后的真实框的赋值
        #---------------------------------------------------#
        assignment[:, :4][best_iou_mask] = encoded_boxes[best_iou_idx, np.arange(assign_num), :4]
        #----------------------------------------------------------#
        #   4代表为背景的概率，设定为0，因为这些先验框有对应的物体
        #----------------------------------------------------------#
        assignment[:, 4][best_iou_mask]     = 0
        assignment[:, 5:-1][best_iou_mask]  = boxes[best_iou_idx, 4:]
        #----------------------------------------------------------#
        #   -1表示先验框是否有对应的物体
        #----------------------------------------------------------#
        assignment[:, -1][best_iou_mask]    = 1
        # 通过assign_boxes我们就获得了，输入进来的这张图片，应该有的预测结果是什么样子的
        return assignment

# DataLoader中collate_fn使用
def ssd_dataset_collate(batch):
    images = []
    bboxes = []
    for img, box in batch:
        images.append(img)
        bboxes.append(box)
    images = torch.from_numpy(np.array(images)).type(torch.FloatTensor)
    bboxes = torch.from_numpy(np.array(bboxes)).type(torch.FloatTensor)
    return images, bboxes


def split_data_into_parts(total_data_count, num_parts=4, percentage=0.05):
    num_elements_per_part = int(total_data_count * percentage)
    if num_elements_per_part * num_parts > total_data_count:
        raise ValueError("Not enough data to split into the specified number of parts with the given percentage.")
    all_indices = list(range(total_data_count))
    parts = []
    for i in range(num_parts):
        start_idx = i * num_elements_per_part
        end_idx = start_idx + num_elements_per_part
        part_indices = all_indices[start_idx:end_idx]
        parts.append(part_indices)
    return parts


def find_index_in_parts(parts, index):
    for i, part in enumerate(parts):
        if index in part:
            return True, i
    return False, -1


def add_watermark_to_image(img, watermark_label, watermark_class_id):
    import random
    import numpy as np
    from PIL import Image
    import qrcode

    # Generate QR code
    qr = qrcode.QRCode(version=1, error_correction=qrcode.constants.ERROR_CORRECT_L, box_size=2, border=1)
    qr.add_data(watermark_label)
    qr.make(fit=True)
    qr_img = qr.make_image(fill='black', back_color='white').convert('RGB')

    # Convert PIL images to numpy arrays for processing
    img_np = np.array(img)
    qr_img_np = np.array(qr_img)
    img_h, img_w = img_np.shape[:2]
    qr_h, qr_w = qr_img_np.shape[:2]
    max_x = img_w - qr_w
    max_y = img_h - qr_h

    if max_x < 0 or max_y < 0:
        raise ValueError("QR code size exceeds image dimensions.")

    while True:
        x_start = random.randint(0, max_x)
        y_start = random.randint(0, max_y)
        x_end = x_start + qr_w
        y_end = y_start + qr_h
        if x_end <= img_w and y_end <= img_h:
            qr_img_cropped = qr_img_np[:y_end - y_start, :x_end - x_start]

            # Replace the corresponding area in the original image
            img_np[y_start:y_end, x_start:x_end] = np.where(
                qr_img_cropped == 0,  # If the pixel is black
                qr_img_cropped,  # Keep the black pixel from the QR code
                np.full_like(img_np[y_start:y_end, x_start:x_end], 255)  # Set the rest to white
            )
            break

    # Convert numpy array back to PIL image
    img = Image.fromarray(img_np)

    # Calculate watermark annotation
    x_center = (x_start + x_end) / 2 / img_w
    y_center = (y_start + y_end) / 2 / img_h
    w = qr_w / img_w
    h = qr_h / img_h
    watermark_annotation = np.array([x_center, y_center, w, h, watermark_class_id])

    return img, watermark_annotation


def detect_and_decode_qr_code(image, watermark_annotation):
    # 将PIL.Image转换为ndarray
    image = np.array(image)
    # 获取图像的宽度和高度
    img_height, img_width = image.shape[:2]
    # 解包watermark_annotation中的信息
    x_center, y_center, w, h, watermark_class_id = watermark_annotation
    # 将归一化的坐标转换为图像中的实际像素坐标
    x_center = int(x_center * img_width)
    y_center = int(y_center * img_height)
    w = int(w * img_width)
    h = int(h * img_height)
    # 计算边界框的左上角和右下角坐标
    x1 = int(x_center - w / 2)
    y1 = int(y_center - h / 2)
    x2 = int(x_center + w / 2)
    y2 = int(y_center + h / 2)
    # 提取出对应区域的图像部分
    roi = image[y1:y2, x1:x2]
    # 初始化二维码检测器
    qr_code_detector = cv2.QRCodeDetector()
    # 检测并解码二维码
    decoded_text, points, _ = qr_code_detector.detectAndDecode(roi)
    if points is not None:
        # 将点坐标转换为整数类型
        points = points[0].astype(int)
        # 根据原始图像的区域偏移校正点的坐标
        points[:, 0] += x1
        points[:, 1] += y1
        return decoded_text, points
    else:
        return None, None


def convert_annotation_to_box(watermark_annotation, img_w, img_h):
    x_center, y_center, w, h, class_id = watermark_annotation

    # Convert normalized coordinates to pixel values
    x_center = x_center * img_w
    y_center = y_center * img_h
    w = w * img_w
    h = h * img_h

    # Calculate x_min, y_min, x_max, y_max
    x_min = x_center - (w / 2)
    y_min = y_center - (h / 2)
    x_max = x_center + (w / 2)
    y_max = y_center + (h / 2)

    return x_min, y_min, x_max, y_max