import os import random import json from pathlib import Path from typing import List, Dict import numpy as np import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader from torchvision import transforms from PIL import Image import tqdm from sklearn.metrics import (accuracy_score, precision_recall_fscore_support, roc_auc_score, confusion_matrix) import matplotlib.pyplot as plt # ==================== CONFIGURATION ==================== class Config: CROP_SIZE = 512 CHECKPOINT_DIR = "./checkpoints" RESULTS_DIR = "./results" # ==================== UTILITIES ==================== def ensure_dir(path: str): Path(path).mkdir(parents=True, exist_ok=True) # ==================== MODEL (copy from training) ==================== class LightweightCompressionNet(nn.Module): def __init__(self): super().__init__() self.conv_blocks = nn.Sequential( nn.Conv2d(3, 16, kernel_size=4, stride=1, padding=0), nn.GELU(), nn.Conv2d(16, 32, kernel_size=4, stride=1, padding=0), nn.GELU(), nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=0), nn.GELU(), nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=0), nn.GELU(), nn.Conv2d(128, 256, kernel_size=4, stride=4, padding=0), nn.GELU(), nn.Conv2d(256, 256, kernel_size=4, stride=4, padding=0), nn.GELU(), nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=0), nn.GELU(), nn.AdaptiveAvgPool2d(1) ) self.head = nn.Sequential( nn.Linear(256, 32), nn.GELU(), nn.Linear(32, 1), nn.Sigmoid() ) def forward(self, x): features = self.conv_blocks(x) features = features.view(features.size(0), -1) return self.head(features).squeeze(1) # ==================== INFERENCE DATASET ==================== class InferenceDataset(Dataset): def __init__(self, image_paths: List[str]): self.image_paths = image_paths self.transform = transforms.Compose([ transforms.Resize(Config.CROP_SIZE), transforms.CenterCrop(Config.CROP_SIZE), transforms.ToTensor(), ]) def __len__(self): return len(self.image_paths) def __getitem__(self, idx): path = self.image_paths[idx] try: image = Image.open(path).convert('RGB') image = self.transform(image) # Apply same INT8 quantization as training image = (image * 255).round().clamp(0, 255) / 255 return image except Exception as e: print(f"Warning: Failed to load {path}: {e}") return torch.zeros(3, Config.CROP_SIZE, Config.CROP_SIZE) # ==================== INFERENCE FUNCTIONS ==================== def load_model(checkpoint_path: str, device: torch.device) -> nn.Module: """Load trained model from checkpoint""" if not os.path.exists(checkpoint_path): raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") model = LightweightCompressionNet().to(device) checkpoint = torch.load(checkpoint_path, map_location=device) model.load_state_dict(checkpoint['model_state_dict']) model.eval() print(f"Loaded model from epoch {checkpoint.get('epoch', 'unknown')} " f"(Val loss: {checkpoint.get('val_loss', 'N/A'):.4f})") return model def run_inference(model: nn.Module, image_paths: List[str], device: torch.device, batch_size: int = 4) -> np.ndarray: """Run inference on a list of images""" dataset = InferenceDataset(image_paths) loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) all_predictions = [] with torch.no_grad(): pbar = tqdm.tqdm(loader, desc="Processing images", unit="batch") for batch in pbar: images = batch.to(device, non_blocking=True) predictions = model(images) all_predictions.extend(predictions.cpu().numpy()) return np.array(all_predictions) def calculate_metrics(y_true: np.ndarray, y_scores: np.ndarray, threshold: float = 0.5) -> Dict: """Calculate classification metrics""" y_pred = (y_scores >= threshold).astype(int) accuracy = accuracy_score(y_true, y_pred) precision, recall, f1, _ = precision_recall_fscore_support( y_true, y_pred, average='binary', zero_division=0 ) try: roc_auc = roc_auc_score(y_true, y_scores) except ValueError: roc_auc = None conf_matrix = confusion_matrix(y_true, y_pred) tn, fp, fn, tp = conf_matrix.ravel() return { 'accuracy': accuracy, 'precision': precision, 'recall': recall, 'f1_score': f1, 'roc_auc': roc_auc, 'true_negatives': int(tn), 'false_positives': int(fp), 'false_negatives': int(fn), 'true_positives': int(tp), 'threshold': threshold, 'confusion_matrix': conf_matrix.tolist() } def plot_score_distribution(ai_scores: np.ndarray, non_ai_scores: np.ndarray, save_path: str = None): """Plot histogram of prediction scores""" plt.figure(figsize=(10, 6)) plt.hist(non_ai_scores, bins=20, alpha=0.7, label='Non-AI', color='blue', density=True) plt.hist(ai_scores, bins=20, alpha=0.7, label='AI', color='red', density=True) plt.axvline(x=0.5, color='black', linestyle='--', label='Threshold (0.5)') plt.xlabel('AI Detection Score') plt.ylabel('Density') plt.title('Distribution of AI Detection Scores') plt.legend() plt.grid(True, alpha=0.3) if save_path: plt.savefig(save_path, dpi=150, bbox_inches='tight') print(f"Saved plot: {save_path}") plt.show() # ==================== MAIN INFERENCE ==================== def main(): # ================================================= # HARDCODED CONFIGURATION - MODIFY THESE VALUES # ================================================= NON_AI_FOLDER = "/home/pc/Dokumenty/test_non_ai" # CHANGE THIS to your non-AI folder AI_FOLDER = "/home/pc/Dokumenty/test_ai" # CHANGE THIS to your AI folder SAMPLE_SIZE = 10 # Number of images from each folder CHECKPOINT_PATH = os.path.join(Config.CHECKPOINT_DIR, "best_model.pt") BATCH_SIZE = 4 SEED = 42 OUTPUT_JSON = os.path.join(Config.RESULTS_DIR, "inference_results.json") # ================================================= # Setup device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Using device: {device}") ensure_dir(Config.RESULTS_DIR) # Load model model = load_model(CHECKPOINT_PATH, device) # Get image paths (support PNG, JPG, JPEG) non_ai_paths = [] for ext in ['*.png', '*.jpg', '*.jpeg']: non_ai_paths.extend([str(p) for p in Path(NON_AI_FOLDER).rglob(ext)]) ai_paths = [] for ext in ['*.png', '*.jpg', '*.jpeg']: ai_paths.extend([str(p) for p in Path(AI_FOLDER).rglob(ext)]) if not non_ai_paths: raise ValueError(f"No images found in {NON_AI_FOLDER}") if not ai_paths: raise ValueError(f"No images found in {AI_FOLDER}") # Random sampling random.seed(SEED) non_ai_sample = random.sample(non_ai_paths, min(SAMPLE_SIZE, len(non_ai_paths))) ai_sample = random.sample(ai_paths, min(SAMPLE_SIZE, len(ai_paths))) print(f"\nSampling {len(non_ai_sample)} non-AI images from {NON_AI_FOLDER}") print(f"Sampling {len(ai_sample)} AI images from {AI_FOLDER}") # Run inference print("\n" + "="*50) non_ai_scores = run_inference(model, non_ai_sample, device, BATCH_SIZE) ai_scores = run_inference(model, ai_sample, device, BATCH_SIZE) # Create labels (0=non-AI, 1=AI) y_true = np.array([0] * len(non_ai_scores) + [1] * len(ai_scores)) y_scores = np.concatenate([non_ai_scores, ai_scores]) # Calculate metrics metrics = calculate_metrics(y_true, y_scores) # Print results print("\n" + "="*50) print("INFERENCE RESULTS") print("="*50) print(f"\nOverall Metrics:") print(f" Accuracy: {metrics['accuracy']:.4f}") print(f" Precision: {metrics['precision']:.4f}") print(f" Recall: {metrics['recall']:.4f}") print(f" F1-Score: {metrics['f1_score']:.4f}") if metrics['roc_auc'] is not None: print(f" ROC-AUC: {metrics['roc_auc']:.4f}") print(f"\nConfusion Matrix:") print(f" True Negatives (Non-AI correct): {metrics['true_negatives']}") print(f" False Positives (Non-AI wrong): {metrics['false_positives']}") print(f" False Negatives (AI wrong): {metrics['false_negatives']}") print(f" True Positives (AI correct): {metrics['true_positives']}") # Per-class accuracy if len(non_ai_scores) > 0: non_ai_acc = (non_ai_scores < 0.5).mean() print(f"\nNon-AI Detection Accuracy: {non_ai_acc:.4f} ({non_ai_acc * 100:.1f}%)") if len(ai_scores) > 0: ai_acc = (ai_scores >= 0.5).mean() print(f"AI Detection Accuracy: {ai_acc:.4f} ({ai_acc * 100:.1f}%)") # Save detailed results results = { 'config': { 'non_ai_folder': NON_AI_FOLDER, 'ai_folder': AI_FOLDER, 'sample_size': SAMPLE_SIZE, 'seed': SEED, 'checkpoint': CHECKPOINT_PATH, 'threshold': metrics['threshold'] }, 'image_paths': { 'non_ai': non_ai_sample, 'ai': ai_sample }, 'predictions': { 'non_ai_scores': non_ai_scores.tolist(), 'ai_scores': ai_scores.tolist() }, 'metrics': metrics } with open(OUTPUT_JSON, 'w') as f: json.dump(results, f, indent=2) print(f"\nDetailed results saved to: {OUTPUT_JSON}") # Plot distribution plot_path = os.path.join(Config.RESULTS_DIR, "score_distribution.png") plot_score_distribution(ai_scores, non_ai_scores, plot_path) print("\nDone!") if __name__ == "__main__": main()