Add CPU compatibility for Mac and testing improvements

- Modified device handling to use CPU instead of CUDA for Mac compatibility
- Updated test_on_images.py to test all images in new_images_to_test folder
- Added test_all_models.py for testing multiple IQA models
- Fixed PyTorch Lightning trainer to use CPU accelerator
- Added .gitignore for checkpoints, logs, and cache files
- Added CLAUDE.md documentation for project setup and usage

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

.gitignore

@@ -0,0 +1,41 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+
+# Virtual environments
+venv/
+env/
+ENV/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# Model checkpoints and weights
+checkpoints/
+feature_extractor_checkpoints/
+prior_methods_checkpoints/
+
+# Training logs and results
+lightning_logs/
+results/*.log
+stdouts/
+
+# Test images
+new_images_to_test/
+
+# Output files
+*.txt
+test_all_models_output.txt
+
+# macOS
+.DS_Store
+
+# Jupyter
+.ipynb_checkpoints/

CLAUDE.md

@@ -0,0 +1,185 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Project Overview
+
+Research implementation for detecting AI-generated images using perceptual features from Image Quality Assessment (IQA) models. The core approach trains two-layer classifiers on feature spaces extracted from pretrained IQA models to distinguish between real and synthetic images.
+
+## Key Commands
+
+### Training
+```bash
+python train.py
+```
+Trains classifiers on specified datasets with configured feature extractors. Training settings are controlled through:
+- Config files in `configs/` directory (arniqa.yaml, contrique.yaml, hyperiqa.yaml, reiqa.yaml, tres.yaml)
+- In-script settings for dataset type (GenImage, DRCT, UnivFD), loss function, and preprocessing
+
+### Testing
+```bash
+python test.py
+```
+Evaluates trained models across datasets with various distortions (Gaussian blur, JPEG compression). Tests both in-domain (same dataset) and cross-domain (different datasets) performance.
+
+```bash
+python test_on_images.py
+```
+Runs inference on specific image files. Modify image paths in the script before running.
+
+### Prior Methods Comparison
+```bash
+python prior_methods/prior_test.py
+```
+Tests baseline comparison methods (CLIP, DRCT) for benchmarking.
+
+### Analysis and Visualization
+```bash
+python analysis/polar_plot.py          # Generate radar plots
+python analysis/distortion_plots.py    # Plot robustness curves
+python analysis/feature_representations.py  # Generate t-SNE visualizations
+```
+
+## Architecture Overview
+
+### Three-Stage Pipeline
+
+1. **Feature Extraction** (`features/`):
+   - IQA models act as frozen feature extractors
+   - Supported models: ARNIQA, CONTRIQUE, HyperIQA, ReIQA, TReS
+   - Also supports CLIP (various architectures) and ResNet50
+   - Each model in `features/` wraps a pretrained backbone
+   - Models loaded via `networks.get_model()` in `functions/networks.py`
+
+2. **Classification** (`functions/networks.py`):
+   - `Classifier_Arch2`: Two-layer MLP (Linear → ReLU → Linear)
+   - Input: IQA feature vector (dimension varies by model, specified in config)
+   - Hidden layer: Typically 1024 units
+   - Output: 2-class logits (real vs. fake)
+
+3. **Training Loop** (`functions/module.py`):
+   - PyTorch Lightning-based training
+   - Loss functions: CrossEntropy, MarginContrastiveLoss (in `loss_optimizers_metrics.py`)
+   - Feature extractor remains frozen; only classifier is trained
+   - Checkpoints saved based on validation loss
+
+### Dataset Structure
+
+Three primary datasets configured in `defaults.py`:
+
+- **GenImage**: 8 generative models (BigGAN, VQDM, SDv4, SDv5, Wukong, ADM, GLIDE, Midjourney)
+- **DRCT**: 16 Stable Diffusion variants (various versions, ControlNet, inpainting, turbo)
+- **UnivFD**: 19 generative models (ProGAN, StyleGAN, CycleGAN, various diffusion models)
+
+Each dataset has separate train/val splits with different generative models.
+
+### Data Preprocessing (`functions/preprocess.py`)
+
+Configurable augmentation pipeline:
+- Gaussian blur (σ=0-5)
+- JPEG compression (QF=30-100)
+- Probability-controlled application during training
+- Image normalization specific to each feature extractor
+
+## Configuration System
+
+YAML config files in `configs/` specify per-model settings:
+
+```yaml
+classifier:
+  input_dim: 4096        # Feature dimension from backbone
+  hidden_layers: [1024]  # Single hidden layer
+
+dataset:
+  model_name: "arniqa"   # Feature extractor identifier
+  f_model_name: "arniqa" # Used for checkpoint naming
+
+trainer:
+  devices: [0]           # GPU indices
+  max_epochs: 20
+  batch_size: 64
+```
+
+The `train.py` script overrides certain config values based on in-script settings (dataset_type, loss function, preprocessing level).
+
+## Path Configuration (CRITICAL)
+
+`defaults.py` contains hardcoded paths that MUST match your environment:
+
+- `main_dataset_dir`: Location of GenImage/UnivFD/DRCT datasets
+- `main_checkpoints_dir`: Where trained classifier checkpoints are saved
+- `main_feature_ckpts_dir`: Pretrained IQA model weights
+- `main_prior_checkpoints_dir`: Prior method checkpoints
+
+**The code checks for specific mount points and will assert False if none match.** You must either:
+1. Update paths in `defaults.py` to match your environment
+2. Create the expected directory structure
+
+## Checkpoint Management
+
+Checkpoints organized hierarchically:
+```
+checkpoints/
+├── GenImage/
+│   └── extensive/
+│       └── MarginContrastiveLoss_CrossEntropy/
+│           └── {model_name}/
+│               └── best_model.ckpt
+└── DRCT/
+    └── extensive/
+        └── MarginContrastiveLoss_CrossEntropy/
+            └── {model_name}/
+                └── best_model.ckpt
+```
+
+Training automatically resumes from `best_model.ckpt` if found in expected location.
+
+## Dependencies
+
+Core libraries (see `functions/ReIQA/requirements.txt` for full list):
+- PyTorch + torchvision
+- PyTorch Lightning (training framework)
+- timm (model architectures)
+- torchmetrics (evaluation)
+- numpy, scipy, scikit-learn, scikit-image
+- PIL (Pillow) for image loading
+- pyyaml for config parsing
+- tqdm for progress bars
+
+Feature extractor dependencies loaded dynamically (e.g., ARNIQA via `torch.hub.load`).
+
+## Important Implementation Details
+
+### Training Script Pattern
+Both `train.py` and `test.py` redirect stdout to log files in `stdouts/` and `results/` directories. Output is not visible in console by default.
+
+### Feature Extraction
+In `functions/module.py`, the global `feature_extractor_module` function is set before training. During training/validation steps, features are extracted with `torch.no_grad()` to prevent gradient computation through the frozen backbone.
+
+### Metrics and Thresholds
+- **GenImage/DRCT**: Fixed threshold of 0.5 for binary classification
+- **UnivFD**: Threshold determined from validation set for optimal accuracy
+
+### Cross-Dataset Testing
+`test.py` includes cross-dataset evaluation (e.g., trained on GenImage, tested on DRCT) to measure generalization.
+
+## Prior Methods (`prior_methods/`)
+
+Comparison implementations of baseline detectors:
+- CLIP-based classifiers (various architectures)
+- DRCT (Detecting and Recovering Content Transformations)
+
+These use similar training patterns but different feature extractors. Organized in parallel structure to main codebase.
+
+## Results and Analysis
+
+- `results/`: CSV files with per-model, per-dataset metrics
+- `analysis/plots/`: Generated visualizations (polar plots, t-SNE, robustness curves)
+- Log files track training progress and test results
+
+## Modifying for New Experiments
+
+1. **Add new feature extractor**: Create wrapper in `features/`, add to `get_model()` in `functions/networks.py`
+2. **Add new dataset**: Update `defaults.py` with source lists, add getter function in `functions/utils.py`
+3. **Change training settings**: Modify settings list in `train.py` (dataset, loss, augmentation level)
+4. **Test new distortions**: Add preprocessing settings in `test.py` preprocess_settings_list

defaults.py

@@ -21,7 +21,18 @@ elif os.path.exists("/mnt/LIVELAB_NAS2/krishna/Perceptual-Classifiers"):
 	main_feature_ckpts_dir = "/mnt/LIVELAB_NAS2/krishna/Perceptual-Classifiers/feature_extractor_checkpoints"
 	main_prior_checkpoints_dir = "/mnt/LIVELAB_NAS2/krishna/Perceptual-Classifiers/prior_methods_checkpoints"
 else:
-	assert False, "Invalid Dataset Directory"
+	# Local setup - use directories relative to this file
+	_base_dir = os.path.dirname(os.path.abspath(__file__))
+	main_dataset_dir = os.path.join(_base_dir, "datasets")
+	main_checkpoints_dir = os.path.join(_base_dir, "checkpoints")
+	main_feature_ckpts_dir = os.path.join(_base_dir, "feature_extractor_checkpoints")
+	main_prior_checkpoints_dir = os.path.join(_base_dir, "prior_methods_checkpoints")
+
+	# Create directories if they don't exist
+	os.makedirs(main_dataset_dir, exist_ok=True)
+	os.makedirs(main_checkpoints_dir, exist_ok=True)
+	os.makedirs(main_feature_ckpts_dir, exist_ok=True)
+	os.makedirs(main_prior_checkpoints_dir, exist_ok=True)
 
 
 

functions/run_on_images_fn.py

@@ -273,7 +273,9 @@ def run_on_images(feature_extractor, classifier, config, test_real_images_paths,
 	# Global Variables: (feature_extractor)
 	global feature_extractor_module
 	feature_extractor_module = feature_extractor
-	feature_extractor_module.to("cuda")
+	# Use CPU for Mac compatibility (change to "cuda" if you have NVIDIA GPU)
+	device = "cpu"
+	feature_extractor_module.to(device)
 	feature_extractor_module.eval()
 	for params in feature_extractor_module.parameters():
 		params.requires_grad = False
@@ -285,8 +287,12 @@ def run_on_images(feature_extractor, classifier, config, test_real_images_paths,
 	Model = Model_LightningModule(classifier, config)
 
 	# PyTorch Lightning Trainer
+	# Override accelerator and devices for Mac compatibility
+	trainer_config = config["trainer"].copy()
+	trainer_config["accelerator"] = "cpu"  # Use "cuda" for NVIDIA GPU, "mps" for Apple Silicon GPU
+	trainer_config["devices"] = 1  # CPU uses integer, GPU uses list like [0]
 	trainer = pl.Trainer(
-		**config["trainer"],
+		**trainer_config,
 		callbacks=[best_checkpoint_callback, utils.LitProgressBar()],
 		precision=32
 	)

test_all_models.py

@@ -0,0 +1,169 @@
+"""
+Test all available models on the same image
+"""
+import os
+import sys
+
+# Available models - test all 5 IQA-based models
+models = ['contrique', 'hyperiqa', 'tres', 'reiqa', 'arniqa']
+
+# Test images directory
+test_images_dir = "new_images_to_test"
+
+# Get all images from the directory
+import glob
+image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']
+test_images = []
+for ext in image_extensions:
+    test_images.extend(glob.glob(os.path.join(test_images_dir, ext)))
+
+if not test_images:
+    print(f"Error: No images found in {test_images_dir}/")
+    sys.exit(1)
+
+print(f"Found {len(test_images)} image(s) in {test_images_dir}/")
+print("=" * 80)
+
+# Import libraries once
+sys.path.insert(0, '.')
+from yaml import safe_load
+from functions.loss_optimizers_metrics import *
+from functions.run_on_images_fn import run_on_images
+import functions.utils as utils
+import functions.networks as networks
+import defaults
+import warnings
+warnings.filterwarnings("ignore")
+
+all_results = {}
+
+# Test each model
+for model_idx, model_name in enumerate(models, 1):
+    print(f"\n{'='*80}")
+    print(f"[{model_idx}/{len(models)}] Testing model: {model_name.upper()}")
+    print("="*80)
+
+    try:
+        config_path = f"configs/{model_name}.yaml"
+        config = safe_load(open(config_path, "r"))
+
+        # Override settings
+        config["dataset"]["dataset_type"] = "GenImage"
+        config["checkpoints"]["resume_dirname"] = "GenImage/extensive/MarginContrastiveLoss_CrossEntropy"
+        config["checkpoints"]["resume_filename"] = "best_model.ckpt"
+        config["checkpoints"]["checkpoint_dirname"] = "extensive/MarginContrastiveLoss_CrossEntropy"
+        config["checkpoints"]["checkpoint_filename"] = "best_model.ckpt"
+
+        # Training settings (for testing)
+        config["train_settings"]["train"] = False
+        config["train_loss_fn"]["name"] = "CrossEntropy"
+        config["val_loss_fn"]["name"] = "CrossEntropy"
+
+        # Model setup
+        device = "cpu"
+        feature_extractor = networks.get_model(model_name=model_name, device=device)
+
+        # Classifier
+        config["classifier"]["hidden_layers"] = [1024]
+        classifier = networks.Classifier_Arch2(
+            input_dim=config["classifier"]["input_dim"],
+            hidden_layers=config["classifier"]["hidden_layers"]
+        )
+
+        # Preprocessing settings
+        preprocess_settings = {
+            "model_name": model_name,
+            "selected_transforms_name": "test",
+            "probability": -1,
+            "gaussian_blur_range": None,
+            "jpeg_compression_qfs": None,
+            "input_image_dimensions": (224, 224),
+            "resize": None
+        }
+
+        print(f"✓ {model_name.upper()} model loaded successfully\n")
+
+        results = []
+
+        # Test each image with this model
+        for idx, test_image in enumerate(test_images, 1):
+            image_name = os.path.basename(test_image)
+            print(f"  [{idx}/{len(test_images)}] Testing: {image_name}")
+
+            # Test images
+            test_real_images_paths = [test_image]
+            test_fake_images_paths = []
+
+            try:
+                test_set_metrics, best_threshold, y_pred, y_true = run_on_images(
+                    feature_extractor=feature_extractor,
+                    classifier=classifier,
+                    config=config,
+                    test_real_images_paths=test_real_images_paths,
+                    test_fake_images_paths=test_fake_images_paths,
+                    preprocess_settings=preprocess_settings,
+                    best_threshold=0.5,
+                    verbose=False
+                )
+
+                score = y_pred[0] if len(y_pred) > 0 else None
+                prediction = "AI-Generated" if score and score > 0.5 else "Real"
+                confidence = abs(score - 0.5) * 200 if score else 0
+
+                results.append({
+                    'image': image_name,
+                    'score': score,
+                    'prediction': prediction,
+                    'confidence': confidence
+                })
+
+                print(f"    ✓ Score: {score:.4f} → {prediction} ({confidence:.1f}% confidence)")
+
+            except Exception as e:
+                print(f"    ✗ Error: {e}")
+                results.append({
+                    'image': image_name,
+                    'score': None,
+                    'prediction': 'Error',
+                    'confidence': 0
+                })
+
+        all_results[model_name] = results
+
+    except Exception as e:
+        print(f"✗ Failed to load {model_name.upper()} model: {e}")
+        all_results[model_name] = None
+
+# Final Summary
+print("\n" + "="*80)
+print("FINAL SUMMARY - ALL MODELS")
+print("="*80)
+
+for model_name, results in all_results.items():
+    if results is None:
+        print(f"\n{model_name.upper()}: Failed to load")
+        continue
+
+    print(f"\n{model_name.upper()}:")
+    print("-"*80)
+    print(f"{'Image':<50} {'Score':<10} {'Prediction':<15} {'Confidence':<12}")
+    print("-"*80)
+
+    for r in results:
+        score_str = f"{r['score']:.4f}" if r['score'] is not None else "N/A"
+        conf_str = f"{r['confidence']:.1f}%" if r['score'] is not None else "N/A"
+        img_name = r['image'][:47] + "..." if len(r['image']) > 50 else r['image']
+        print(f"{img_name:<50} {score_str:<10} {r['prediction']:<15} {conf_str:<12}")
+
+    # Statistics
+    valid_predictions = [r for r in results if r['score'] is not None]
+    if valid_predictions:
+        avg_score = sum(r['score'] for r in valid_predictions) / len(valid_predictions)
+        ai_count = sum(1 for r in valid_predictions if r['score'] > 0.5)
+        real_count = len(valid_predictions) - ai_count
+        avg_confidence = sum(r['confidence'] for r in valid_predictions) / len(valid_predictions)
+
+        print("-"*80)
+        print(f"Average Score: {avg_score:.4f} | AI: {ai_count} | Real: {real_count} | Avg Confidence: {avg_confidence:.1f}%")
+
+print("\n" + "="*80)

test_on_images.py

@@ -15,18 +15,25 @@ import functions.utils as utils
 import functions.networks as networks
 import defaults
 
-dir_path = "/home/krishna/Perceptual-Classifiers-Working/images/True=Real_Pred=Fake"
-test_real_images_files = os.listdir(dir_path)
+# Get all images from new_images_to_test folder
+import glob
+test_images_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "new_images_to_test")
+image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']
 test_real_images_paths = []
-for f in test_real_images_files:
-	test_real_images_paths.append(
-		os.path.join(
-			dir_path, f
-		)
-	)
+for ext in image_extensions:
+	test_real_images_paths.extend([os.path.abspath(p) for p in glob.glob(os.path.join(test_images_dir, ext))])
 
 test_fake_images_paths = []
 
+if not test_real_images_paths:
+	print(f"Error: No images found in {test_images_dir}/")
+	sys.exit(1)
+
+print(f"Found {len(test_real_images_paths)} image(s) to test:")
+for img in test_real_images_paths:
+	print(f"  - {os.path.basename(img)}")
+print()
+
 # Calling Main function
 if __name__ == '__main__':
 	# -----------------------------------------------------------------
@@ -125,8 +132,9 @@ if __name__ == '__main__':
 					f_model_name = config["dataset"]["f_model_name"]
 
 
-					# Model
-					feature_extractor = networks.get_model(model_name=config["dataset"]["model_name"], device="cuda")
+					# Model - use CPU for Mac (MPS not fully supported by all models)
+					device = "cpu"  # Change to "cuda" if you have NVIDIA GPU
+					feature_extractor = networks.get_model(model_name=config["dataset"]["model_name"], device=device)
 
 					
 					# Classifier