Create app.py

app.py

@@ -0,0 +1,274 @@
+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()