| import torch
|
| import os
|
| import argparse
|
| from dataclasses import dataclass, field
|
| import transformers
|
| from torch.utils.data import Dataset, DataLoader
|
| from transformers import CLIPImageProcessor
|
| import pdb
|
| import json
|
| from transformers import AutoProcessor, LlavaForConditionalGeneration
|
| from tqdm import tqdm
|
| import random
|
| import numpy as np
|
| import torch
|
| import torchvision.transforms as T
|
| from PIL import Image
|
| from torchvision.transforms.functional import InterpolationMode
|
| import torch.nn as nn
|
|
|
|
|
| IMAGENET_MEAN = (0.485, 0.456, 0.406)
|
| IMAGENET_STD = (0.229, 0.224, 0.225)
|
|
|
| def parse_args():
|
| parser = argparse.ArgumentParser(description="Legion Model Training")
|
|
|
|
|
| parser.add_argument("--model_path", default="", type=str)
|
| parser.add_argument("--val_batch_size", default=1, type=int)
|
| parser.add_argument("--workers", default=1, type=int)
|
| parser.add_argument("--data_base_test", default="", type=str)
|
| parser.add_argument("--test_json_file", default="", type=str)
|
| parser.add_argument("--output_path", default="", type=str)
|
| return parser.parse_args()
|
|
|
| class legion_cls_dataset(Dataset):
|
| def __init__(self, args, train=True):
|
| super().__init__()
|
| self.args = args
|
| self.train = train
|
| if train == True:
|
| with open(args.train_json_file, 'r') as f:
|
| self.data = json.load(f)
|
| elif train == False:
|
| with open(args.test_json_file, 'r') as f:
|
| self.data = json.load(f)
|
| self.processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf", revision='a272c74')
|
|
|
| def __len__(self):
|
| return len(self.data)
|
|
|
| def __getitem__(self, idx):
|
| if self.train == True:
|
| img_path = os.path.join(self.args.data_base_train, self.data[idx]['image'])
|
| else:
|
| img_path = os.path.join(self.args.data_base_test, self.data[idx]['image'])
|
| label = self.data[idx]['label']
|
|
|
| image = Image.open(img_path)
|
|
|
| inputs = self.processor(
|
| text=self.data[idx]['conversations'][0]['value'],
|
| images=image,
|
| return_tensors="pt",
|
| padding="max_length",
|
| max_length=1024,
|
| truncation=True
|
| )
|
|
|
| cate = 'deepfake'
|
|
|
| return inputs, [label], [img_path], [cate]
|
|
|
|
|
|
|
|
|
| def load_model(args):
|
| print("Loading model...")
|
| model = LlavaForConditionalGeneration.from_pretrained(
|
| args.model_path,
|
| torch_dtype=torch.bfloat16,
|
| low_cpu_mem_usage=True,
|
| use_flash_attention_2=True,
|
| revision='a272c74',
|
| ).eval().cuda()
|
| print("Successfully loaded model from:", args.model_path)
|
| return model
|
|
|
| def calculate_results_acc(results):
|
| acc_results = {}
|
|
|
| for cate in results:
|
| data = results[cate]
|
|
|
| right_real = data['right']['right_real']
|
| right_fake = data['right']['right_fake']
|
| wrong_real = data['wrong']['wrong_real']
|
| wrong_fake = data['wrong']['wrong_fake']
|
|
|
| total_real = right_real + wrong_real
|
| total_fake = right_fake + wrong_fake
|
| total = total_real + total_fake
|
|
|
| acc_total = (right_real + right_fake) / total if total != 0 else 0
|
| acc_real = right_real / total_real if total_real != 0 else 0
|
| acc_fake = right_fake / total_fake if total_fake != 0 else 0
|
|
|
| acc_results[cate] = {
|
| 'total_samples': total,
|
| 'total_accuracy': round(acc_total, 4),
|
| 'real_accuracy': round(acc_real, 4),
|
| 'fake_accuracy': round(acc_fake, 4),
|
| 'confusion_matrix': {
|
| 'right_real': right_real,
|
| 'wrong_real': wrong_real,
|
| 'right_fake': right_fake,
|
| 'wrong_fake': wrong_fake,
|
| }
|
| }
|
|
|
| global_stats = {
|
| 'total_right': sum(r['right']['right_real'] + r['right']['right_fake'] for r in results.values()),
|
| 'total_wrong': sum(r['wrong']['wrong_real'] + r['wrong']['wrong_fake'] for r in results.values())
|
| }
|
| global_stats['global_accuracy'] = global_stats['total_right'] / (global_stats['total_right'] + global_stats['total_wrong'])
|
|
|
| return {
|
| 'category_acc': acc_results,
|
| 'global_stats': global_stats
|
| }
|
|
|
|
|
| def validate(args, model, cls_test_dataloader):
|
| processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf", revision='a272c74')
|
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
| results = {}
|
| outputs = []
|
| with torch.no_grad():
|
| for inputs, labels, paths, cates in tqdm(cls_test_dataloader):
|
|
|
| inputs["input_ids"] = inputs["input_ids"].squeeze().to(device)
|
| inputs["attention_mask"] = inputs["attention_mask"].squeeze().to(device)
|
| inputs["pixel_values"] = inputs["pixel_values"].squeeze().to(device)
|
| output = model.generate(**inputs, max_new_tokens=256)
|
| pred_cls = []
|
|
|
| for i in range(output.shape[0]):
|
| response = processor.decode(output[i], skip_special_tokens=True).split('?')[-1]
|
| print(response)
|
| outputs.append({"image_path": paths[0][i], "output": response})
|
|
|
| if 'real' in response.split('.')[0].lower():
|
| pred_cls.append(1)
|
| elif 'fake' in response.split('.')[0].lower():
|
| pred_cls.append(0)
|
| else:
|
| try:
|
| if 'real' in response.split('.')[1].lower():
|
| pred_cls.append(1)
|
| elif 'fake' in response.split('.')[1].lower():
|
| pred_cls.append(0)
|
| else:
|
| print(f"no fake or real in reponse:{response}")
|
| pred_cls.append(random.choice([0, 1]))
|
| except:
|
| print(f"no fake or real in reponse:{response}")
|
| pred_cls.append(random.choice([0, 1]))
|
| for label, pred, cate in zip(labels[0].tolist(), pred_cls, cates[0]):
|
| if cate not in results:
|
| results[cate] = {'right':{'right_fake':0, 'right_real':0}, 'wrong':{'wrong_fake':0, 'wrong_real':0}}
|
| if label == pred:
|
| if label == 1:
|
| results[cate]['right']['right_real'] += 1
|
| else:
|
| results[cate]['right']['right_fake'] += 1
|
| else:
|
| if label == 1:
|
| results[cate]['wrong']['wrong_real'] += 1
|
| else:
|
| results[cate]['wrong']['wrong_fake'] += 1
|
|
|
| os.makedirs('results', exist_ok=True)
|
| with open(args.output_path, "w") as file:
|
| json.dump(outputs, file, indent=2)
|
| acc = calculate_results_acc(results)
|
| print(acc)
|
|
|
|
|
|
|
| def main():
|
| args = parse_args()
|
| model = load_model(args)
|
| model.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
|
| cls_test_dataset = legion_cls_dataset(args, train=False)
|
| cls_test_dataloader = DataLoader(
|
| cls_test_dataset,
|
| batch_size=args.val_batch_size,
|
| shuffle=False,
|
| num_workers=args.workers,
|
| pin_memory=True,
|
| )
|
| validate(args, model, cls_test_dataloader)
|
|
|
|
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
| |
