| """Unit tests for Imager-1-Tiny.""" |
| import os |
| import torch |
| import pytest |
| from PIL import Image |
|
|
| from src.config import ImagerTiny, MODEL_CONFIG, load_model, predict |
|
|
|
|
| def test_config_threshold(): |
| assert MODEL_CONFIG["threshold"] == 0.385 |
| assert MODEL_CONFIG["input_size"] == 224 |
| assert MODEL_CONFIG["num_classes"] == 1 |
|
|
|
|
| def test_architecture_output_shape(): |
| model = ImagerTiny(pretrained=False) |
| model.eval() |
| dummy = torch.randn(1, 3, 224, 224) |
| with torch.no_grad(): |
| out = model(dummy) |
| assert out.shape == (1, 1), f"Expected (1,1), got {out.shape}" |
|
|
|
|
| def test_architecture_param_count(): |
| model = ImagerTiny(pretrained=False) |
| n = sum(p.numel() for p in model.parameters()) |
| |
| assert 24_000_000 < n < 27_000_000, f"Unexpected param count: {n}" |
|
|
|
|
| @pytest.mark.skipif( |
| not os.path.exists("model/model_imager-1-tiny.pth"), |
| reason="Weights not downloaded", |
| ) |
| def test_predict_runs(): |
| model = load_model("model/model_imager-1-tiny.pth") |
| img = Image.new("RGB", (512, 512), color=(128, 128, 128)) |
| result = predict(img, model=model) |
| assert "probability" in result |
| assert result["label"] in ("Real", "AI-Generated") |
| assert 0.0 <= result["probability"] <= 1.0 |
|
|
