clean/image/cospy/evaluate.py

import os
import json
import torch
import importlib
import numpy as np
from tqdm import tqdm
from PIL import Image
from loguru import logger

from dataSets import *
from utils import seed_torch, evaluate

import warnings
warnings.filterwarnings("ignore")


# Load pre-trained models for evaluation
class Detector:
    def __init__(self,
                 device: str,
                 mode: str = "fusion",
                 train_dataset: str = "sd-v1_4",
                 pretrain: bool = False,
                 ckpt: str = "ckpt",
                 batch_size: int = 32):

        # Device
        self.device = device
        self.mode = mode
        self.train_dataset = train_dataset
        self.pretrain = pretrain
        self.ckpt = ckpt
        self.batch_size = batch_size

        # Dynamically import the detector module from either "progan" or "sd-v1_4"
        detector_module = importlib.import_module(f"detectors.{train_dataset}")

        # Get the detector and load weights based on mode
        if pretrain:
            # Only provide pre-trained weights for fusion mode
            # Hardcode to fusion mode
            self.mode = "fusion"
            # Load the fusion detector with pre-trained weights
            semantic_weights_path = f"pretrained/{self.train_dataset}/semantic_weights.pth"
            artifact_weights_path = f"pretrained/{self.train_dataset}/artifact_weights.pth"
            fusion_weights_path = f"pretrained/{self.train_dataset}/fusion_weights.pth"
            if not os.path.exists(semantic_weights_path) or not os.path.exists(artifact_weights_path) or not os.path.exists(fusion_weights_path):
                raise ValueError("The pre-trained weights are not complete for evaluation")
            CoSpyFusionDetector = getattr(detector_module, "CoSpyFusionDetector")
            self.model = CoSpyFusionDetector(
                semantic_weights_path=semantic_weights_path,
                artifact_weights_path=artifact_weights_path)
            self.model.load_weights(fusion_weights_path)
        else:
            if self.mode == "fusion":
                semantic_weights_path = os.path.join(self.ckpt, self.train_dataset, "semantic", "best_model.pth")
                artifact_weights_path = os.path.join(self.ckpt, self.train_dataset, "artifact", "best_model.pth")
                fusion_weights_path = os.path.join(self.ckpt, self.train_dataset, "fusion", "best_model.pth")
                if not os.path.exists(semantic_weights_path) or not os.path.exists(artifact_weights_path) or not os.path.exists(fusion_weights_path):
                    raise ValueError("Semantic, Artifact or Fusion weights path does not exist for fusion mode")
                CoSpyFusionDetector = getattr(detector_module, "CoSpyFusionDetector")
                self.model = CoSpyFusionDetector(
                    semantic_weights_path=semantic_weights_path,
                    artifact_weights_path=artifact_weights_path)
                self.model.load_weights(fusion_weights_path)
            elif self.mode == "end2end":
                End2EndDetector = getattr(detector_module, "End2EndDetector")
                self.model = End2EndDetector()
                end2end_weights_path = os.path.join(self.ckpt, self.train_dataset, "end2end", "best_model.pth")
                if not os.path.exists(end2end_weights_path):
                    raise ValueError("End2End weights path does not exist for end2end mode")
                self.model.load_weights(end2end_weights_path)
            else:
                raise ValueError(f"Unknown mode: {self.mode}")
        
        # Put the model on the device and set to eval
        self.model.to(self.device)
        self.model.eval()
    
    def evaluate_benchmark(self):
        # Select the appropriate test dataset and evaluation lists
        if self.train_dataset == "progan":
            benchmark_name = "AIGCDetectionBenchMark"
            TestDataset = AIGCDetectTestDataset
            eval_dataset_list = AIGCDetectionBenchMark_DATASET_LIST
            eval_model_list = AIGCDetectionBenchMark_MODEL_LIST
        elif self.train_dataset == "sd-v1_4":
            benchmark_name = "Co-Spy-Bench"
            TestDataset = CoSpyBenchTestDataset
            eval_dataset_list = CoSpyBench_DATASET_LIST
            eval_model_list = CoSpyBench_MODEL_LIST
        else:
            raise ValueError(f"Unknown train dataset: {self.train_dataset}")

        # Set the saving directory
        if self.pretrain:
            save_dir = os.path.join(self.ckpt, self.train_dataset, self.mode, f"pretrain_{benchmark_name}")
        else:
            save_dir = os.path.join(self.ckpt, self.train_dataset, self.mode, f"eval_{benchmark_name}")
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)

        # Setup logger
        log_path = f"{save_dir}/evaluation.log"
        if os.path.exists(log_path):
            os.remove(log_path)
        
        logger_id = logger.add(
            log_path,
            format="{time:MM-DD at HH:mm:ss} | {level} | {module}:{line} | {message}",
            level="DEBUG",
        )

        # Save raw model prediction
        save_output_path = os.path.join(save_dir, "output.json")
        # Save summarized evaluation result
        save_result_path = os.path.join(save_dir, "result.json")

        # Begin the evaluation
        result_all = {}
        output_all = {}
        for dataset_name in eval_dataset_list:
            result_all[dataset_name] = {}
            output_all[dataset_name] = {}
            for model_name in eval_model_list:
                test_dataset = TestDataset(dataset=dataset_name, model=model_name, transform=self.model.test_transform)
                test_loader = torch.utils.data.DataLoader(test_dataset,
                                                          batch_size=self.batch_size,
                                                          shuffle=False,
                                                          num_workers=4,
                                                          pin_memory=True)

                # Evaluate the model
                y_pred, y_true = [], []
                for (images, labels) in tqdm(test_loader, desc=f"Evaluating {benchmark_name} - {dataset_name} - {model_name}"):
                    y_pred.extend(self.model.predict(images))
                    y_true.extend(labels.tolist())

                ap, accuracy = evaluate(y_pred, y_true)
                logger.info(f"Evaluate on {benchmark_name} - {dataset_name} - {model_name} | Size {len(y_true)} | AP {ap*100:.2f}% | Accuracy {accuracy*100:.2f}%")

                result_all[dataset_name][model_name] = {"size": len(y_true), "AP": ap, "Accuracy": accuracy}
                output_all[dataset_name][model_name] = {"y_pred": y_pred, "y_true": y_true}

        # Save the results
        with open(save_result_path, "w") as f:
            json.dump(result_all, f, indent=4)

        with open(save_output_path, "w") as f:
            json.dump(output_all, f, indent=4)

    def scan(self):
        # Load the image
        image_filepath = input("Please enter the image filepath for scanning: ")
        if not os.path.exists(image_filepath):
            print(f"Image file not found: {image_filepath}")
            image_filepath = input("Please enter the image filepath for scanning: ")

        image = Image.open(image_filepath).convert("RGB")
        image = self.model.test_transform(image)
        image = image.unsqueeze(0)
        image = image.to(self.device)

        # Make the prediction
        prediction = self.model.predict(image)[0]

        if prediction > 0.5:
            print(f"Co-Spy Prediction: {prediction:.3f} - AI-Generated")
        else:
            print(f"Co-Spy Prediction: {prediction:.3f} - Real")