File size: 1,279 Bytes
6efad37 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | """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())
# ResNet-50 has ~25.6M params; allow a margin
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
|