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 from vllm import LLM, SamplingParams import pdb IMAGENET_MEAN = (0.485, 0.456, 0.406) IMAGENET_STD = (0.229, 0.224, 0.225) def parse_args(): parser = argparse.ArgumentParser(description="FakeVLM Model Testing") # Model-specific settings 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) 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'] cate = self.data[idx]['cate'] # torch.Size([n, 3, 448, 448]), int, int, str, str return [self.data[idx]['conversations'][0]['value']], [label], [img_path], [cate] # change this function to our own to support custom behaviors def load_model(args): print("Loading model...") llm = LLM(model=args.model_path, dtype="float16", max_model_len=800, tensor_parallel_size=torch.cuda.device_count() ) print("Successfully loaded model from:", args.model_path) return llm 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), 'overall': { '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): results = {} output_result = [] sampling_params = SamplingParams( max_tokens=4096, temperature=0, ) with torch.no_grad(): for questions, labels, imgs, cates in tqdm(cls_test_dataloader): inputs = [] for question, img in zip(questions[0], imgs[0]): inputs.append({ "prompt": question, "multi_modal_data": {"image": Image.open(img)}, }) # pdb.set_trace() outputs = model.generate(inputs, sampling_params=sampling_params) pred_cls = [] for i, output in enumerate(outputs): # pdb.set_trace() # #save result output_result.append({'id':imgs[0][i], 'caption':output.outputs[0].text}) response = output.outputs[0].text 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 # save result os.makedirs("results", exist_ok=True) with open(args.output_path, "w") as file: json.dump(output_result, file, indent=2) acc = calculate_results_acc(results) print(acc) def main(): args = parse_args() model = load_model(args) 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()