Merge branch 'main' of https://github.com/oss-slu/mithridatium

.github/ISSUE_TEMPLATE/feature_request.md

@@ -3,15 +3,18 @@ name: Feature Request
 about: Suggest a new feature or improvement
 title: "[FEATURE] "
 labels: enhancement
-assignees: ''
-
+assignees: ""
 ---
 
 ## Summary
 
 Briefly describe the feature you’d like to see.
 
-## Tasks that need to completed for this feature
+## Acceptance Criteria
+
+The acceptance Criteria to accept this issue as done
+
+## Tasks that need to be completed for this feature
 
 A list of individual tasks that likely must be done before the feature can be considered "complete"
 

.github/ISSUE_TEMPLATE/good_first_issue.md

@@ -0,0 +1,35 @@
+---
+name: Good first issue
+about: A simple, well-defined task that’s perfect for someone new to the project.
+title: "Good First Issue: [TASK]"
+labels: "good first issue"
+assignees: ""
+---
+
+## Description
+
+This is a simple task for first-time contributors! Please follow the steps below to implement the feature or fix the bug.
+
+### Steps to reproduce
+
+- Step 1
+- Step 2
+- Step 3
+
+### Expected behavior
+
+- What should happen after the task is completed
+
+### How to contribute
+
+1. Fork the repository.
+2. Create a new branch from `main` (`git checkout -b feature/new-feature`).
+3. Work on the task and commit your changes (`git commit -m "Implement new feature"`).
+4. Push the changes and create a pull request.
+5. Ensure that your code passes all tests and is documented.
+
+---
+
+### Additional Context
+
+- Link to relevant resources (e.g., issues, discussions, or other PRs).

.gitignore

@@ -3,6 +3,7 @@
 venv/
 env/
 .env/
+mith/
 
 # Python cache
 __pycache__/
@@ -18,7 +19,8 @@ dist/
 # Data & models
 data/
 models/
-reports/*.json
+/reports/*
+!/reports/report_schema.json
 
 # Notebooks & logs
 *.ipynb_checkpoints/
@@ -36,3 +38,5 @@ Thumbs.db
 .coverage
 .pytest_cache/
 .mypy_cache/
+
+results.npy

CITATION.cff

@@ -0,0 +1,32 @@
+title: mithridatium
+authors:
+  - given-names: Pelumi
+    family-names: Oluwategbe
+    email: pelumi.oluwategbe@slu.edu
+    affiliation: Saint Louis University
+  - given-names: William
+    family-names: Phoenix
+    email: will.phoenix@slu.edu
+    affiliation: Saint Louis University
+  - given-names: Gustavo
+    family-names: Lucca
+    email: gustavo.lucca@slu.edu
+    affiliation: Saint Louis University
+  - given-names: Payton
+    family-names: Guffey
+    email: payton.guffey@slu.edu
+    affiliation: Saint Louis University
+cff-version: 1.2.0
+message: If you use this software, please cite it using the metadata from this file.
+type: software
+abstract: Mithridatium is a research-driven project aimed at detecting backdoors
+  and data poisoning in downloaded pretrained models or pipelines (e.g., from
+  Hugging Face).   Our goal is to provide a modular, command-line tool that
+  helps researchers and engineers trust the models they use.
+keywords:
+  - data privacy
+  - machine-learning
+  - python
+  - security
+license: MIT-Modern-Variant
+repository-code: https://github.com/oss-slu/mithridatium

CODE_OF_CONDUCT.md

@@ -0,0 +1,78 @@
+## Code of Conduct
+
+## Our Pledge
+
+We, as members, contributors, and maintainers of Mithridatium, pledge to make participation in our project and community a harassment-free experience for everyone, regardless of:
+
+- age, body size, disability, ethnicity, gender identity and expression,
+
+- level of experience, nationality, personal appearance, race, religion,
+
+- or sexual identity and orientation.
+
+We are committed to fostering an environment where all participants feel respected, valued, and empowered to contribute.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment include:
+
+- Using welcoming and inclusive language
+
+- Being respectful of differing viewpoints and experiences
+
+- Giving and gracefully accepting constructive feedback
+
+- Showing empathy toward other community members
+
+- Recognizing that collaboration is more valuable than competition
+
+Examples of unacceptable behavior include:
+
+- The use of sexualized language or imagery and unwelcome sexual attention
+
+- Trolling, insulting, or derogatory comments and personal attacks
+
+- Public or private harassment
+
+- Publishing others’ private information without explicit permission
+
+- Any behavior that would reasonably be considered inappropriate in a professional setting
+
+## Our Responsibilities
+
+Project maintainers are responsible for clarifying and enforcing community standards.
+They have the right and responsibility to remove, edit, or reject:
+
+- comments, commits, code, wiki edits, issues, and pull requests that are not aligned with this Code of Conduct,
+
+- or temporarily or permanently ban any contributor for other behavior deemed inappropriate, threatening, or harmful.
+
+## Scope
+
+This Code of Conduct applies within all project spaces (GitHub issues, pull requests, documentation, and discussions)
+and in public spaces when an individual represents the project or its community.
+
+Examples of representing the project include using an official project e-mail address,
+posting via an official social media account, or acting as a representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be reported by contacting the maintainers at:
+
+📩 pelumi.oluwategbe@slu.edu
+
+📩 daniel.shown@slu.edu
+
+All complaints will be reviewed and investigated promptly and fairly.
+The project team is obligated to maintain confidentiality regarding the reporter of an incident.
+
+## Attribution
+
+This Code of Conduct is adapted from the Contributor Covenant
+, version 2.1,
+available at https://www.contributor-covenant.org/version/2/1/code_of_conduct.html
+.
+
+🧡 Thank you
+
+By participating in this project, you help make the open-source community a safe, collaborative, and innovative space for everyone.

CONTRIBUTING.md

@@ -7,7 +7,7 @@ Thank you for checking out **Mithridatium**! We are excited to have you here. Th
 ⚠️ **Note:**
 
 - Issues labeled **`internal team`** are reserved for the project’s assigned developers and will not be accepted from outside contributors.
-- Once the framework is stable, we will open up selected issues for external contributors with labels such as **`good first issue`** or **`help wanted`**.
+- **`good first issue`** and **`help wanted`** indicate tasks that are open to the community.
 
 We encourage you to watch this repository if you’d like to stay updated!
 

LICENSE

@@ -0,0 +1,17 @@
+Permission is hereby granted, without written agreement and without
+license or royalty fees, to use, copy, modify, and distribute this
+software and its documentation for any purpose, provided that the
+above copyright notice and the following two paragraphs appear in
+all copies of this software.
+
+IN NO EVENT SHALL THE COPYRIGHT HOLDER BE LIABLE TO ANY PARTY FOR
+DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES
+ARISING OUT OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN
+IF THE COPYRIGHT HOLDER HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH
+DAMAGE.
+
+THE COPYRIGHT HOLDER SPECIFICALLY DISCLAIMS ANY WARRANTIES, INCLUDING,
+BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+FITNESS FOR A PARTICULAR PURPOSE.  THE SOFTWARE PROVIDED HEREUNDER IS
+ON AN "AS IS" BASIS, AND THE COPYRIGHT HOLDER HAS NO OBLIGATION TO
+PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.

README.md

@@ -20,3 +20,52 @@ This comes with risks:
 ---
 
 ## Other Functionaly will be updated as the project goes on
+
+## Quickstart
+
+```bash
+python -m venv .venv && source .venv/bin/activate
+pip install -e .
+pip install pytest pytest-cov
+
+# (A) Train demo models (fast settings)
+
+# Clean model on 5 epochs (Increase epochs for better accuracy, but it will take longer)
+python -m scripts.train_resnet18 --dataset clean --epochs 5 --output_path models/resnet18_clean.pth
+
+# Poisoned model on 5 epochs (Increase epochs for better accuracy, but it will take longer)
+python -m scripts.train_resnet18 --dataset poison --train_poison_rate 0.1 --target_class 0 \
+  --epochs 5 --output_path models/resnet18_poison.pth
+
+# (B) Run detection (default: resnet18)
+mithridatium detect --model models/resnet18_poison.pth --defense mmbd --data cifar10 --out reports/mmbd.json
+
+# (Optional) Specify architecture (supported: resnet18, resnet34)
+mithridatium detect --model models/resnet18_poison.pth --defense mmbd --data cifar10 --arch resnet34 --out reports/mmbd.json
+
+# (C) See summary
+cat reports/mmbd.json
+```
+
+## CLI Help
+
+To see all available options and arguments:
+
+```bash
+mithridatium detect --help
+```
+
+Example output:
+
+```
+Usage: mithridatium detect [OPTIONS]
+
+Options:
+  --model, -m TEXT     The model path .pth. E.g. 'models/resnet18.pth'. [default: models/resnet18.pth]
+  --data, -d TEXT      The dataset name. E.g. 'cifar10'. [default: cifar10]
+  --defense, -D TEXT   The defense you want to run. E.g. 'spectral'. [default: spectral]
+  --arch, -a TEXT      The model architecture to use. Supported: 'resnet18', 'resnet34'. [default: resnet18]
+  --out, -o TEXT       The output path for the JSON report. Use "-" for stdout or a file path (e.g. "reports/report.json"). [default: reports/report.json]
+  --force, -f          This allows overwriting. E.g. if the output file already exists --force will overwrite it.
+  --help               Show this message and exit.
+```

codemeta.json

@@ -0,0 +1,91 @@
+{
+  "name": "mithridatium",
+  "@context": "https://w3id.org/codemeta/3.0",
+  "applicationCategory": "Security and protection software",
+  "author": [
+    {
+      "affiliation": {
+        "name": "Saint Louis University",
+        "type": "Organization"
+      },
+      "email": "pelumi.oluwategbe@slu.edu",
+      "familyName": "Oluwategbe",
+      "id": "https://pelumi-tegbe.vercel.app/",
+      "givenName": "Pelumi",
+      "type": "Person"
+    },
+    {
+      "affiliation": {
+        "name": "Saint Louis University",
+        "type": "Organization"
+      },
+      "email": "will.phoenix@slu.edu",
+      "familyName": "Phoenix",
+      "id": "_:author_2",
+      "givenName": "William",
+      "type": "Person"
+    },
+    {
+      "affiliation": {
+        "name": "Saint Louis University",
+        "type": "Organization"
+      },
+      "email": "gustavo.lucca@slu.edu",
+      "familyName": "Lucca",
+      "id": "_:author_3",
+      "givenName": "Gustavo",
+      "type": "Person"
+    },
+    {
+      "affiliation": {
+        "name": "Saint Louis University",
+        "type": "Organization"
+      },
+      "email": "payton.guffey@slu.edu",
+      "familyName": "Guffey",
+      "id": "_:author_4",
+      "givenName": "Payton",
+      "type": "Person"
+    },
+    {
+      "roleName": "Technical Lead",
+      "startDate": "2025-08-27",
+      "schema:author": "https://pelumi-tegbe.vercel.app/",
+      "type": "Role"
+    },
+    {
+      "roleName": "Developer",
+      "startDate": "2025-08-27",
+      "schema:author": "_:author_2",
+      "type": "Role"
+    },
+    {
+      "roleName": "Developer",
+      "startDate": "2025-08-27",
+      "schema:author": "_:author_3",
+      "type": "Role"
+    },
+    {
+      "roleName": "Developer",
+      "startDate": "2025-08-27",
+      "schema:author": "_:author_4",
+      "type": "Role"
+    }
+  ],
+  "codeRepository": "https://github.com/oss-slu/mithridatium",
+  "dateCreated": "2025-08-28",
+  "description": "Mithridatium is a research-driven project aimed at detecting backdoors and data poisoning in downloaded pretrained models or pipelines (e.g., from Hugging Face).   Our goal is to provide a modular, command-line tool that helps researchers and engineers trust the models they use.",
+  "developmentStatus": "active",
+  "issueTracker": "https://github.com/oss-slu/mithridatium/issues",
+  "keywords": [
+    "data privacy",
+    "machine-learning",
+    "python",
+    "security"
+  ],
+  "license": "https://spdx.org/licenses/MIT-Modern-Variant",
+  "programmingLanguage": [
+    "Python"
+  ],
+  "type": "SoftwareSourceCode"
+}

dummyfile.txt

@@ -0,0 +1 @@
+hello world

dummytest.txt

@@ -0,0 +1 @@
+hello world

examples/demo_commands.md

@@ -0,0 +1,28 @@
+# Demo Commands for Mithridatium
+
+## 1. Set up environment:
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+pip install -e .
+pip install pytest pytest-cov
+```
+
+## 2. Train Clean model:
+
+```bash
+python -m scripts.train_resnet18 --dataset clean --epochs 5 --output_path models/resnet18_clean.pth
+```
+
+## 3. Train Poisoned model:
+
+```bash
+python -m scripts.train_resnet18 --dataset poison --train_poison_rate 0.1 --target_class 0 --epochs 5 --output_path models/resnet18_poison.pth
+```
+
+## 4. Run detection:
+
+```bash
+mithridatium detect --model models/resnet18_poison.pth --defense mmbd --data cifar10 --out reports/mmbd.json
+```

examples/end_to_end.md

@@ -0,0 +1,14 @@
+# End-to-End Smoke
+
+```bash
+# 1) Train demo models
+python -m scripts.train_resnet18 --dataset clean  --epochs 3 --output_path models/resnet18_clean.pth
+python -m scripts.train_resnet18 --dataset poison --train_poison_rate 0.1 --target_class 0 \
+  --epochs 3 --output_path models/resnet18_poison.pth
+
+# 2) Run detect (wires CLI → Loader → Evaluator → Defense → Report)
+mithridatium detect --model models/resnet18_poison.pth --defense mmbd --data cifar10 --out reports/mmbd.json
+
+# 3) See summary
+cat reports/mmbd.json
+```

examples/sample_report.json

@@ -0,0 +1,12 @@
+{
+  "mithridatium_version": "0.1.0",
+  "timestamp_utc": "2025-01-01T00:00:00Z",
+  "model_path": "models/resnet18_poison.pth",
+  "defense": "mmbd",
+  "dataset": "cifar10",
+  "results": {
+    "suspected_backdoor": true,
+    "num_flagged": 500,
+    "top_eigenvalue": 42.3
+  }
+}

mithridatium.egg-info/PKG-INFO

@@ -1,9 +1,15 @@
 Metadata-Version: 2.4
 Name: mithridatium
-Version: 0.1.0
+Version: 0.1.1
 Summary: Framework for verifying integrity of pretrained AI models
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: typer>=0.12
+Requires-Dist: torch
+Requires-Dist: torchvision
+Requires-Dist: jsonschema
+Dynamic: license-file
 
 # Mithridatium 🛡️
 
@@ -27,3 +33,52 @@ This comes with risks:
 ---
 
 ## Other Functionaly will be updated as the project goes on
+
+## Quickstart
+
+```bash
+python -m venv .venv && source .venv/bin/activate
+pip install -e .
+pip install pytest pytest-cov
+
+# (A) Train demo models (fast settings)
+
+# Clean model on 5 epochs (Increase epochs for better accuracy, but it will take longer)
+python -m scripts.train_resnet18 --dataset clean --epochs 5 --output_path models/resnet18_clean.pth
+
+# Poisoned model on 5 epochs (Increase epochs for better accuracy, but it will take longer)
+python -m scripts.train_resnet18 --dataset poison --train_poison_rate 0.1 --target_class 0 \
+  --epochs 5 --output_path models/resnet18_poison.pth
+
+# (B) Run detection (default: resnet18)
+mithridatium detect --model models/resnet18_poison.pth --defense mmbd --data cifar10 --out reports/mmbd.json
+
+# (Optional) Specify architecture (supported: resnet18, resnet34)
+mithridatium detect --model models/resnet18_poison.pth --defense mmbd --data cifar10 --arch resnet34 --out reports/mmbd.json
+
+# (C) See summary
+cat reports/mmbd.json
+```
+
+## CLI Help
+
+To see all available options and arguments:
+
+```bash
+mithridatium detect --help
+```
+
+Example output:
+
+```
+Usage: mithridatium detect [OPTIONS]
+
+Options:
+  --model, -m TEXT     The model path .pth. E.g. 'models/resnet18.pth'. [default: models/resnet18.pth]
+  --data, -d TEXT      The dataset name. E.g. 'cifar10'. [default: cifar10]
+  --defense, -D TEXT   The defense you want to run. E.g. 'spectral'. [default: spectral]
+  --arch, -a TEXT      The model architecture to use. Supported: 'resnet18', 'resnet34'. [default: resnet18]
+  --out, -o TEXT       The output path for the JSON report. Use "-" for stdout or a file path (e.g. "reports/report.json"). [default: reports/report.json]
+  --force, -f          This allows overwriting. E.g. if the output file already exists --force will overwrite it.
+  --help               Show this message and exit.
+```

mithridatium.egg-info/SOURCES.txt

@@ -1,15 +1,25 @@
+LICENSE
 README.md
 pyproject.toml
 mithridatium/__init__.py
 mithridatium/cli.py
-mithridatium/data.py
 mithridatium/evaluator.py
 mithridatium/loader.py
 mithridatium/report.py
+mithridatium/utils.py
 mithridatium.egg-info/PKG-INFO
 mithridatium.egg-info/SOURCES.txt
 mithridatium.egg-info/dependency_links.txt
+mithridatium.egg-info/entry_points.txt
+mithridatium.egg-info/requires.txt
 mithridatium.egg-info/top_level.txt
 mithridatium/defenses/__init__.py
-mithridatium/defenses/spectral.py
-tests/test_cli.py
+mithridatium/defenses/mmbd.py
+mithridatium/defenses/strip.py
+tests/test_dataloader_normalization.py
+tests/test_evaluator.py
+tests/test_preprocess_config.py
+tests/test_strip_entropy.py
+tests/test_strip_scores.py
+tests/test_utils_configs.py
+tests/tests_report.py

mithridatium/cli.py

@@ -1,16 +1,225 @@
 # mithridatium/cli.py
 import typer
+import json
+from pathlib import Path
+import sys
+from mithridatium import report as rpt
+from mithridatium import loader as loader
+from mithridatium import utils
+from mithridatium.defenses.mmbd import run_mmbd
+from mithridatium.defenses.strip import strip_scores
+from mithridatium.defenses.mmbd import get_device
+from mithridatium.loader import validate_model
+
+
+
+VERSION = "0.1.1"
+DEFENSES = {"mmbd", "strip"}
+
+EXIT_USAGE_ERROR = 64     # invalid CLI usage (e.g., unsupported --defense)
+EXIT_NO_INPUT = 66        # input file missing/not a file
+EXIT_CANT_CREATE = 73     # cannot create/overwrite output without --force
+EXIT_IO_ERROR = 74        # input exists but can't be opened/read
 
 app = typer.Typer(help="Mithridatium CLI - verify pretrained model integrity")
 
+
+def _write_json(obj: dict, out_path: str, force: bool) -> None:
+    """
+    Write JSON to a file or to stdout.
+    - Stdout using "--out -"
+    - Overwrite using "--force"
+
+    """
+
+    if out_path == "-":
+        json.dump(obj, sys.stdout, indent=2)
+        sys.stdout.write("\n")
+        return
+    
+    path = Path(out_path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Checks if file exists and prevents overwriting. Use --force to override.
+    if path.exists() and not force:
+        typer.secho(
+            f"Error: output file already exists: {path}.",
+        )
+        raise typer.Exit(code=EXIT_CANT_CREATE)
+
+    with path.open("w", encoding="utf-8") as f:
+        json.dump(obj, f, indent=2)
+
+
+def dummy_report(model_path: str, defense: str, out_path: str, force: bool) -> None:
+    """
+    Nothing runs yet, just a dummy report.
+    """
+    
+    # dummy report:
+    report = {
+        "mithridatium_version": VERSION,
+        "model_path": model_path,
+        "defense": defense,
+        "status": "Not yet implemented", 
+    }
+
+    _write_json(report, out_path, force)
+    where = "stdout" if out_path == "-" else out_path
+    typer.echo(f"Report written to {where}")
+
+
+@app.callback(invoke_without_command=True)
+def _root(
+    # This is a calback that prints the version whenever it is ran.
+    version: bool = typer.Option(
+        False,
+        "--version",
+        "-V",
+        help="Show Mithridatium version and exit.",
+        is_eager=True, # ensures this runs before any command (including --help
+    )
+):
+
+    if version:
+        typer.echo(VERSION)
+        raise typer.Exit()
+
+@app.command()
+def defenses() -> None:
+    """
+    List supported defenses.
+    """
+    for d in sorted(DEFENSES):
+        typer.echo(d)
+
 @app.command()
 def detect(
-    model: str = typer.Option("models/resnet18.pth", "--model", "-m", help="Path to model .pth (can be missing for now)"),
-    data: str = typer.Option("cifar10", "--data", "-d", help="Dataset name"), #needed or not?
-    defense: str = typer.Option("spectral", "--defense", "-D", help="Defense to run"),
-    out: str = typer.Option("reports/report.json", "--out", "-o", help="Path to write JSON report"),
+    model: str = typer.Option(
+        "models/resnet18.pth",
+        "--model",
+        "-m",
+        help="The model path .pth. E.g. 'models/resnet18.pth'.",
+    ),
+    data: str = typer.Option(
+        "cifar10",
+        "--data",
+        "-d",
+        help="The dataset name. E.g. 'cifar10'.",
+    ),
+    defense: str = typer.Option(
+        "mmbd",
+        "--defense",
+        "-D",
+        help="The defense you want to run. E.g. 'mmbd' or 'strip'.",
+    ),
+    arch: str = typer.Option(
+        "resnet18",
+        "--arch",
+        "-a",
+        help="The model architecture to use. E.g. 'resnet18'.",
+    ),
+    out: str = typer.Option(
+        "reports/report.json",
+        "--out",
+        "-o",
+        help='The output path for the JSON report. Use "-" for stdout or a file path (e.g. "reports/report.json").',
+    ),
+    force: bool = typer.Option(
+        False, 
+        "--force", 
+        "-f", 
+        help="This allows overwriting. E.g. if the output file already exists --force will overwrite it.",
+    ),
 ):
-    typer.echo(f"[args] model={model}  data={data}  defense={defense}  out={out}")
+    """
+    Argument validation:
+    1) Model path exists and is a file
+    2) File exists but can't be loaded
+    3) Unsupported defense
+    4) Write dummy JSON (stdout allowed via --out -)
+    """
+    # 1) Model path exists and is a file
+    p = Path(model)
+    if not p.exists() or not p.is_file():
+        typer.secho(
+            f"Error: model path not found or not a file: {p}", err=True
+        )
+        raise typer.Exit(code=EXIT_NO_INPUT)
+
+    # 2) File exists but can't be loaded
+    try:
+        with p.open("rb"):
+            pass
+    except OSError as ex:
+        typer.secho(
+            f"Error: model file could not be opened: {p}\nReason: {ex}", err=True
+        )
+        raise typer.Exit(code=EXIT_IO_ERROR)
+    
+    # 3) Unsupported defense
+    d = defense.strip().lower()
+    if d not in DEFENSES:
+        typer.secho(
+            "Error: unsupported --defense "
+            f"'{defense}'. Supported defenses: {', '.join(sorted(DEFENSES))}", err=True
+        )
+        raise typer.Exit(code=EXIT_USAGE_ERROR)
+    
+        # 4) Build model arch
+    print(f"[cli] building model architecture '{arch}'…")
+    mdl, feature_module = loader.build_model(arch, num_classes=10)
+
+    # 5) Load weights from checkpoint
+    print("[cli] loading weights…")
+    mdl = loader.load_weights(mdl, str(p))
+
+    # 6) Validate model BEFORE any defense runs
+    # cfg = utils.load_preprocess_config(str(p))  # has input_size etc.
+    cfg = utils.get_preprocess_config(data)  # has input_size etc.
+
+    try:
+        print("[cli] validating model (architecture + dry forward)…")
+        input_size = cfg.get_input_size() 
+        validate_model(mdl, arch, input_size)
+        print("[cli] model validation OK")
+    except Exception as ex:
+        typer.secho(
+            f"Error: model validation failed.\n{ex}",
+            err=True,
+        )
+        raise typer.Exit(code=EXIT_IO_ERROR)
+
+    # 7) Build dataloader (TEMP: CIFAR-10; replace with PreprocessConfig)
+    print("[cli] building dataloader…")
+    test_loader, config = utils.dataloader_for(data, "test", 256)
+
+
+    # 8) Run the defenses that are supported
+    print(f"[cli] running defense={d}…")
+    try:
+        device = get_device(0)
+        mdl = mdl.to(device)
+        if d == "mmbd":
+            # Move model to appropriate device for MMBD
+            results = run_mmbd(mdl, config)
+        elif d == "strip":
+            results = strip_scores(mdl, config)
+        else:
+            results = {"suspected_backdoor": False, "num_flagged": 0, "top_eigenvalue": 0.0}
+
+    except Exception as ex:
+        typer.secho(
+            f"Error: failed to run '{d}' on model {p}.\nReason: {ex}", err=True
+        )
+        raise typer.Exit(code=EXIT_IO_ERROR)
+   
+
+    # 8) Build & write report
+    rep = rpt.build_report(model_path=str(p), defense=d, dataset=data, version=VERSION, results=results)
+    _write_json(rep, out, force)
+    print(rpt.render_summary(rep))
+
 
 if __name__ == "__main__":
-    app()
+    app()

mithridatium/cli_notes.md

@@ -0,0 +1,183 @@
+# Mithridatium CLI — How it works & how to use it
+
+## Install (development)
+
+```bash
+# from the repo root, inside your virtualenv
+pip install -e .
+```
+
+---
+
+## Commands
+
+### Show version / help
+
+```bash
+mithridatium --version
+mithridatium --help
+```
+
+### List supported defenses
+
+```bash
+mithridatium defenses
+# spectral
+# mmbd
+```
+
+### Detect (main workflow)
+
+Runs argument validation, executes the selected defense, writes JSON to a file or stdout, and prints a summary.
+
+```bash
+mithridatium detect   --model models/resnet18_clean.pth   --defense spectral   --data cifar10   --out reports/spectral.json
+```
+
+**Options**
+
+- `-m, --model PATH` (required): path to a model checkpoint (.pth).
+  - For `spectral`, this **must** be a valid PyTorch checkpoint (loadable by `torch.load`).
+  - For `mmbd` (stub), any readable file is fine (results are placeholder).
+- `-D, --defense [spectral|mmbd]` (required): which defense to run.
+  - `spectral`: runs a simple weight‑matrix spectral check (computes top eigenvalue of \(W^T W\)).
+  - `mmbd`: Multi‑Model Backdoor Detection **stub** (returns fixed demo metrics).
+- `-d, --data TEXT` (optional): dataset tag (e.g., `cifar10`). Stored in the report for provenance.
+- `-o, --out PATH` (required): where to write JSON. Use `-` to write JSON to **stdout**.
+- `-f, --force`: allow overwriting an existing output file.
+
+**Examples**
+
+Write JSON to a file + print summary:
+
+```bash
+mithridatium detect -m models/resnet18_clean.pth -D spectral -d cifar10 -o reports/spectral.json
+```
+
+Write JSON to **stdout** (first), then summary:
+
+```bash
+mithridatium detect -m models/resnet18_clean.pth -D spectral -d cifar10 -o -
+```
+
+Overwrite an existing JSON file:
+
+```bash
+mithridatium detect -m models/resnet18_clean.pth -D spectral -d cifar10 -o reports/spectral.json --force
+```
+
+Pretty‑print JSON without `jq`:
+
+```bash
+mithridatium detect -m models/resnet18_clean.pth -D spectral -d cifar10 -o - | python -m json.tool
+```
+
+Run from the package subfolder (note the `../` paths):
+
+```bash
+cd mithridatium
+mithridatium detect -m ../models/resnet18_clean.pth -D spectral -d cifar10 -o ../reports/spectral.json
+```
+
+### Show a saved report (validate then display)
+
+`show-report` first **validates** the JSON against the schema at `reports/report_schema.json`.
+
+- If valid: prints the chosen view (default **pretty JSON**).
+- If invalid: prints a single error and exits non-zero.
+
+```bash
+# Pretty JSON (default)
+mithridatium show-report -f reports/spectral.json
+
+# Human-readable summary (if you kept render_summary)
+mithridatium show-report -f reports/spectral.json --mode summary
+```
+
+---
+
+## Output
+
+### JSON schema
+
+```json
+{
+	"mithridatium_version": "0.1.1",
+	"model_path": "models/resnet18_clean.pth",
+	"defense": "spectral",
+	"dataset": "cifar10",
+	"results": {
+		"suspected_backdoor": true,
+		"num_flagged": 0,
+		"top_eigenvalue": 80.46
+	}
+}
+```
+
+> `mmbd` currently returns a stubbed `results` with fixed demo metrics.  
+> `spectral` computes a `top_eigenvalue` from the **largest weight matrix** in the checkpoint and sets a boolean verdict based on a demo threshold inside the runner.
+
+## Exit codes
+
+- `64` (`EXIT_USAGE_ERROR`) – invalid CLI usage (e.g., unsupported `--defense`).
+- `65` (`EXIT_DATA_ERR`) – invalid report data (schema validation failed in `show-report`).
+- `66` (`EXIT_NO_INPUT`) – model path missing or not a file.
+- `73` (`EXIT_CANT_CREATE`) – output file exists and `--force` not supplied.
+- `74` (`EXIT_IO_ERROR`) – I/O problems (e.g., `torch.load` failed, unreadable file).
+
+Your CI can key off these codes.
+
+---
+
+## What each defense does
+
+### `spectral`
+
+- Loads the checkpoint via `torch.load`.
+- Finds the **largest** weight‑like tensor (≥ 2D), flattens to a matrix `[out, features]`.
+- Runs power iteration to estimate the top eigenvalue of \(W^T W\).
+- Compares against a demo threshold to set `suspected_backdoor`, can be changed.
+
+### `mmbd`
+
+- Returns fixed demo metrics (`suspected_backdoor=true`, `num_flagged=500`, `top_eigenvalue=42.3`).
+
+---
+
+## Quick ways to get a model
+
+### 1) One‑liner: make a tiny valid `.pth` for spectral
+
+```bash
+python - <<'PY'
+import torch, pathlib
+path = pathlib.Path("models"); path.mkdir(exist_ok=True)
+sd = {"layer.weight": torch.randn(64, 128)}  # a 2D tensor
+torch.save(sd, "models/spectral_demo.pth")
+print("[ok] wrote models/spectral_demo.pth")
+PY
+```
+
+### 2) Train a clean CIFAR‑10 ResNet‑18 (short run)
+
+```bash
+python scripts/train_resnet18.py   --epochs 1   --train_batch_size 128   --eval_batch_size 256   --lr 0.1   --seed 1   --output_path models/resnet18_clean.pth
+```
+
+### 3) Train a backdoored model (BadNets‑style)
+
+```bash
+python scripts/train_backdoor_resnet18.py   --poison-rate 0.1   --target-class 0   --trigger-size 4   --trigger-pos bottom-right   --epochs 5   --batch-size 128   --lr 0.1   --seed 42   --out models/resnet18_badnet.pth
+```
+
+---
+
+## Troubleshooting
+
+- **“model path not found or not a file”**  
+  Check your working directory and the path. Adjust with `../` if you’re in `mithridatium/`.
+
+- **`torch.load` error with `spectral`**  
+  Your file isn’t a valid PyTorch checkpoint. Use the one‑liner above or a trained model.
+
+---

mithridatium/data.py

@@ -1,14 +0,0 @@
-# mithridatium/data.py
-import torch
-from torchvision import datasets, transforms
-
-def get_cifar10_loader(batch_size: int = 128):
-    tfm = transforms.Compose([
-        transforms.Resize(224),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485,0.456,0.406],
-                             [0.229,0.224,0.225]),
-    ])
-    ds = datasets.CIFAR10(root="data", train=False, download=True, transform=tfm)
-    loader = torch.utils.data.DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=2)
-    return loader

mithridatium/defenses/aeva.py

@@ -0,0 +1,3 @@
+def run_aeva():
+
+    return "Hello World"

mithridatium/defenses/mmbd.py

@@ -0,0 +1,185 @@
+# from __future__ import absolute_import
+# from __future__ import print_function
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import resnet18
+
+# import argparse
+# import argparse
+import random
+import numpy as np
+
+#Code adapted from https://github.com/wanghangpsu/MM-BD/blob/main/univ_bd.py
+
+def get_device(device_index=0):
+    if torch.cuda.is_available():
+        return torch.device(f"cuda:{device_index}")
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        return torch.device("mps")
+    else:
+        return torch.device("cpu")
+
+# parser = argparse.ArgumentParser(description='UnivBD method')
+# parser.add_argument('--model_dir', default='model1', help='model path')
+# parser.add_argument('--device', default=0, type=int)
+# parser.add_argument("--report_out", default="reports/mmbd_report.json", help="JSON output path")
+#parser.add_argument('--data_path', '-d', required=True, help='data path')
+# args = parser.parse_args()
+# parser = argparse.ArgumentParser(description='UnivBD method')
+# parser.add_argument('--model_dir', default='model1', help='model path')
+# parser.add_argument('--device', default=0, type=int)
+# parser.add_argument("--report_out", default="reports/mmbd_report.json", help="JSON output path")
+# parser.add_argument('--data_path', '-d', required=True, help='data path')
+# args = parser.parse_args()
+
+'''def load_resnet18_cifar10(weights_path, device=0):
+    model = resnet18(weights=None)
+    model.fc = nn.Linear(model.fc.in_features, 10)
+
+    try:
+        state = torch.load(weights_path, map_location=device, weights_only=True)
+    except TypeError:
+        state = torch.load(weights_path, map_location=device)
+
+    model.load_state_dict(state, strict=True)
+    model.to(device).eval()
+    return model'''
+
+
+def run_mmbd(model, configs, device=None):
+
+    random.seed()
+    if device is None:
+        try:
+            device = next(model.parameters()).device
+        except StopIteration:
+            device = get_device(0)
+
+    # Detection parameters
+    NC = 10
+    NI = 150
+    PI = 0.9
+    NSTEP = 75
+    TC = 6
+    batch_size = 20
+
+    N_CLASSES_TO_PROBE = 5
+    NUM_IMAGES = 30
+
+    # Load model
+    model = model.to(device=device, dtype=torch.float32).eval()
+    criterion = nn.CrossEntropyLoss()
+
+    
+    model = model.to(device).eval()
+    mean = torch.tensor(configs.get_mean(), device=device).view(1, 3, 1, 1)
+    std = torch.tensor(configs.get_std(), device=device).view(1, 3, 1, 1)
+
+    def lr_scheduler(iter_idx):
+        lr = 1e-2
+
+
+        return lr
+
+    res = []
+    for t in range(N_CLASSES_TO_PROBE):
+        print(f"[MMBD] optimizing class {t+1}/{N_CLASSES_TO_PROBE}…", flush=True)
+        images = torch.rand([NUM_IMAGES, *configs.input_size], device=device, dtype=torch.float32, requires_grad=True)
+        last_loss = 1000.0
+        labels = torch.full((len(images),), t, dtype=torch.long, device=device)
+        onehot_label = F.one_hot(labels, num_classes=NC).to(device=device, dtype=torch.float32)
+
+        optimizer = torch.optim.SGD([images], lr=1e-2, momentum=0.9)
+        
+
+        for iter_idx in range(NSTEP):
+            optimizer.zero_grad(set_to_none=True)
+
+            x = torch.clamp(images, 0, 1)
+            x = (x - mean) / std
+            outputs = model(x)
+
+            loss = (-(outputs * onehot_label).sum()
+                    + torch.max((1 - onehot_label) * outputs - 1000 * onehot_label, dim=1).values.sum())
+            loss.backward()
+            optimizer.step()
+
+            curr = float(loss.item())
+            if iter_idx % 50 == 0 or iter_idx == NSTEP - 1:
+                print(f"[MMBD]   Iter {iter_idx}/{NSTEP}, loss={curr:.4f}")
+            if abs(last_loss - curr) / max(abs(last_loss), 1e-12) < 1e-5:
+                print(f"[MMBD]   Converged early at iter {iter_idx}")
+                break
+            last_loss = curr
+
+        res.append(torch.max(torch.sum(outputs * onehot_label, dim=1)
+                - torch.max((1 - onehot_label) * outputs - 1000 * onehot_label, dim=1).values).item())
+
+    stats = np.array(res, dtype=float)
+    from scipy.stats import median_abs_deviation as MAD
+    from scipy.stats import gamma
+    mad = MAD(stats, scale='normal')
+    mad = float(mad) if mad != 0 else 1e-12
+    abs_deviation = np.abs(stats - np.median(stats))
+    score = abs_deviation / mad
+
+
+    np.save('results.npy', np.array(res))
+    ind_max = int(np.argmax(stats))
+    r_eval = float(np.amax(stats))
+    r_null = np.delete(stats, ind_max)
+
+    shape, loc, scale = gamma.fit(r_null)
+    pv = 1 - pow(gamma.cdf(r_eval, a=shape, loc=loc, scale=scale), len(r_null)+1)
+    verdict = "Likely clean" if pv > 0.05 else "Likely backdoored"
+
+    # suspected_backdoor = (verdict == "attack")
+    # num_flagged = 1 if suspected_backdoor else 0
+    top_eigenvalue = float(r_eval)
+
+
+    thresholds = {
+        "p_value": 0.05,
+        "normalized_score": {
+            "normal": [0.0, 1.5],
+            "mild": [1.5, 3.0],
+            "suspicious": [3.0, 5.0],
+            "very_suspicious": [5.0, None]
+        },
+    }
+
+    parameters = {
+        "NC": NC,
+        "NSTEP": NSTEP,
+        "optimizer": "SGD(momentum=0.2)",
+        "lr_init": 1e-2,
+        "device": str(device),
+    }
+
+    results = {
+        "defense": "mmbd",
+        "per_class_scores": stats.tolist(),
+        "normalized_scores": score.tolist(),
+        "p_value": float(pv),
+        "verdict": verdict,
+        # "suspected_target": (int(ind_max) if verdict == "attack" else None),
+        "thresholds": thresholds,
+        "parameters": parameters,
+        "dataset": configs.get_dataset(),
+
+        # "suspected_backdoor": suspected_backdoor,
+        # "num_flagged": int(num_flagged),
+        "top_eigenvalue": float(top_eigenvalue),
+    }
+
+    return results
+
+    '''build_report(
+        model_path=args.model_dir,
+        defense="MMBD",
+        out_path=args.report_out,
+        details=results,
+        version="0.1.1"
+    )'''

mithridatium/defenses/strip.py

@@ -0,0 +1,144 @@
+import torch
+import random
+import numpy as np
+from typing import Dict, Any, List
+from mithridatium import utils
+
+from mithridatium.defenses.mmbd import get_device
+
+#comment
+
+def prediction_entropy(logits: torch.Tensor) -> torch.Tensor:
+    """
+    Returns per-sample entropy over the softmax distribution.
+
+    Args:
+        logits: A tensor of shape (batch_size, num_classes) containing the logits.
+
+    Returns:
+        A tensor of shape (batch_size,) containing the entropy for each sample.
+    """
+    p = torch.nn.Softmax(dim=1)(logits) + 1e-8
+    return (-p * p.log()).sum(1)
+
+def strip_scores(
+        model, 
+        configs, 
+        num_bases: int = 32, 
+        num_perturbations: int = 16, 
+        device=None,
+        entropy_mean_threshold=0.45
+        ) -> Dict[str, Any]:
+    """
+    Computes STRIP-style entropy scores.
+
+    Args:
+        model: The model to evaluate.
+        configs: Preprocess configuration.
+        num_bases: Number of base samples to evaluate.
+        num_perturbations: Number of perturbations per base sample.
+        device: Device to run the computation on.
+
+    Returns:
+        A dictionary containing the raw entropy scores.
+    """
+    if device is None:
+        try:
+            device = next(model.parameters()).device
+        except StopIteration:
+            device = get_device(0)
+
+    model = model.to(device=device, dtype=torch.float32).eval()
+
+       # -------- Build test dataloader ----------
+    # configs already contains dataset name, batch size, transforms, etc.
+    test_loader, _ = utils.dataloader_for(
+        configs.get_dataset(),
+        split="test",
+        batch_size=256
+    )
+
+
+    # Collect all images from the dataloader to use as a pool for mixing
+    all_images = []
+    for images, _ in test_loader:
+        all_images.append(images)
+        if len(all_images) * images.shape[0] >= num_bases + num_perturbations * 2: # Heuristic to stop early if we have enough data
+             break
+    
+    if not all_images:
+         raise ValueError("Dataloader is empty")
+
+    all_images = torch.cat(all_images, dim=0)
+    
+    # Ensure we have enough images
+    if len(all_images) < num_bases:
+        num_bases = len(all_images)
+        # raise ValueError(f"Not enough images in dataloader. Needed {num_bases}, got {len(all_images)}")
+
+    # Select base samples
+    indices = torch.randperm(len(all_images))
+    base_indices = indices[:num_bases]
+    base_images = all_images[base_indices].to(device, dtype=torch.float32)
+
+    entropies_list = []
+
+    with torch.no_grad():
+        for i in range(num_bases):
+            base_img = base_images[i]
+            
+            # Create perturbations
+            # We need num_perturbations other images. 
+            # We can sample from the whole pool (excluding the current base if we want, but collision prob is low)
+            perturb_indices = torch.randint(0, len(all_images), (num_perturbations,))
+            perturb_images = all_images[perturb_indices].to(device, dtype=torch.float32)
+            
+            # Superimpose: 0.5 * base + 0.5 * other
+            # base_img is (C, H, W), perturb_images is (N, C, H, W)
+            # Broadcast base_img
+            mixed_images = 0.5 * base_img.unsqueeze(0) + 0.5 * perturb_images
+            
+            logits = model(mixed_images)
+            entropies = prediction_entropy(logits)
+            
+            # Aggregate entropy for this base sample
+            mean_entropy = entropies.mean().item()
+            entropies_list.append(mean_entropy)
+    if not entropies_list:
+        raise ValueError("No entropies were computed.")
+    
+    entropy_mean = float(np.mean(entropies_list))
+    entropy_min  = float(np.min(entropies_list))
+    entropy_max  = float(np.max(entropies_list))
+
+    if not entropies_list:
+        raise ValueError("No entropies were computed.")
+
+    entropy_mean = float(np.mean(entropies_list))
+    entropy_min  = float(np.min(entropies_list))
+    entropy_max  = float(np.max(entropies_list))
+
+    if entropy_mean > entropy_mean_threshold:
+        verdict = "likely backdoored"
+    else:
+        verdict = "likely clean"
+
+    return {
+        "defense": "strip",
+        "entropies": entropies_list,
+        "statistics": {
+            "entropy_mean": entropy_mean,
+            "entropy_min": entropy_min,
+            "entropy_max": entropy_max,
+        },
+        "parameters": {
+            "num_bases": num_bases,
+            "num_perturbations": num_perturbations,
+        },
+        "dataset": str(configs.get_dataset()),
+        "verdict": verdict,
+        "thresholds": {
+            "entropy_mean_threshold": entropy_mean_threshold
+        }
+    }
+

mithridatium/evaluator.py

@@ -1,33 +1,67 @@
-# mithridatium/evaluator.py
 import torch
+import torch.nn as nn
+from typing import Tuple
 
-@torch.no_grad()
-def extract_embeddings(model, dataloader, feature_module):
+def extract_embeddings(model: nn.Module, loader: torch.utils.data.DataLoader, feature_module: nn.Module) -> Tuple[torch.Tensor, torch.Tensor]:
     """
-    Collect penultimate features using a forward hook on `feature_module`
-    (e.g., resnet.avgpool). Returns:
-      embs:  [N, D] tensor
-      labels:[N] tensor
+    Extract penultimate-layer embeddings and labels from a model and dataloader.
+    Args:
+        model: The neural network model (e.g., resnet18).
+        loader: DataLoader for the dataset.
+        feature_module: The module in the model whose output is the embedding (e.g., model.avgpool or model.layer4).
+    Returns:
+        embs: Tensor of shape [N, D] (embeddings)
+        labels: Tensor of shape [N] (labels)
     """
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device).eval()
+    model.eval()
+    embs = []
+    labels = []
+    device = next(model.parameters()).device
 
-    feats_list, labels_list = [], []
+    def hook_fn(module, input, output):
+        hook_fn.embeddings = output.detach()
+    hook_fn.embeddings = None
+    hook = feature_module.register_forward_hook(hook_fn)
 
-    def hook(_m, _inp, out):
-        # avgpool output is [B, C, 1, 1]; flatten to [B, C]
-        feats_list.append(out.detach().flatten(1).cpu())
-
-    # Register the hook on the target layer
-    handle = feature_module.register_forward_hook(lambda m, i, o: hook(m, i, o))
-    try:
-        for x, y in dataloader:
+    with torch.no_grad():
+        for x, y in loader:
             x = x.to(device)
-            _ = model(x)   # running forward triggers the hook
-            labels_list.append(y) # keep labels to align with the embeddings
-    finally:
-        handle.remove()
-
-    embs = torch.cat(feats_list, dim=0)
-    labels = torch.cat(labels_list, dim=0)
+            _ = model(x)
+            emb = hook_fn.embeddings
+            if emb.dim() > 2:
+                emb = torch.flatten(emb, start_dim=1)
+            embs.append(emb.cpu())
+            labels.append(y.cpu())
+    hook.remove()
+    embs = torch.cat(embs, dim=0)
+    labels = torch.cat(labels, dim=0)
     return embs, labels
+
+def evaluate(model: nn.Module, loader: torch.utils.data.DataLoader) -> Tuple[float, float]:
+    """
+    Evaluate model on a dataset.
+    Args:
+        model: The neural network model.
+        loader: DataLoader for the dataset.
+    Returns:
+        loss: Average loss (float)
+        accy: Accuracy (float)
+    """
+    model.eval()
+    criterion = nn.CrossEntropyLoss()
+    total_loss = 0.0
+    correct = 0
+    total = 0
+    device = next(model.parameters()).device
+    with torch.no_grad():
+        for x, y in loader:
+            x, y = x.to(device), y.to(device)
+            out = model(x)
+            loss = criterion(out, y)
+            total_loss += loss.item() * y.size(0)
+            pred = out.argmax(1)
+            correct += (pred == y).sum().item()
+            total += y.size(0)
+    avg_loss = total_loss / total
+    accy = correct / total
+    return avg_loss, accy

mithridatium/loader.py

@@ -1,10 +1,21 @@
-# mithridatium/loader.py
 from pathlib import Path
 import torch
 import torch.nn as nn
 import torchvision.models as models
+from dataclasses import dataclass, field
+from typing import Tuple, List
+import json
 
 def load_resnet18(model_path: str | None):
+    """
+    Load a ResNet-18 model with optional checkpoint.
+    
+    Args:
+        model_path: Path to checkpoint file, or None for random init.
+        
+    Returns:
+        Tuple of (model, feature_module).
+    """
     model = models.resnet18(weights=None)
 
     # expose the penultimate layer (avgpool -> flatten) for features
@@ -21,3 +32,109 @@ def load_resnet18(model_path: str | None):
 
     model.eval()
     return model, feature_module
+
+def get_feature_module(model):
+    """
+    Returns the penultimate feature module for a given model architecture.
+    
+    Args:
+        model: PyTorch model instance.
+        
+    Returns:
+        The feature extraction module (e.g., model.avgpool for ResNet).
+        
+    Raises:
+        NotImplementedError: If architecture is not supported.
+    """
+    arch = model.__class__.__name__
+    if arch == 'ResNet':
+        return model.avgpool
+    # Example for future extension:
+    # elif arch == 'VGG':
+    #     return model.classifier[0]
+    else:
+        raise NotImplementedError(f"Feature module not defined for architecture: {arch}")
+    
+def build_model(arch: str = "resnet18", num_classes: int = 10):
+    """
+    Build a model with the specified architecture.
+    
+    Args:
+        arch: Architecture name (currently only "resnet18" supported).
+        num_classes: Number of output classes.
+        
+    Returns:
+        Tuple of (model, feature_module).
+    """
+    if arch.lower() == "resnet18":
+        from torchvision.models import resnet18
+        m = resnet18(weights=None)
+    elif arch == "resnet34":
+        from torchvision.models import resnet34
+        m = resnet34(weights=None)
+    else:
+        raise NotImplementedError(f"Architecture '{arch}' not yet supported")
+        
+    m.fc = torch.nn.Linear(m.fc.in_features, num_classes)
+    return m, get_feature_module(m)
+ 
+
+def load_weights(model, ckpt_path: str):
+    """
+    Load model weights from a checkpoint file.
+    
+    Args:
+        model: PyTorch model instance.
+        ckpt_path: Path to checkpoint file.
+        
+    Returns:
+        Model with loaded weights.
+    """
+    sd = torch.load(ckpt_path, map_location="cpu")
+    missing, unexpected = model.load_state_dict(sd, strict=False)
+    if missing or unexpected:
+        print(f"[warn] load_weights: missing={missing}, unexpected={unexpected}")
+    return model
+
+
+def validate_model(model: torch.nn.Module, arch: str, input_size):
+    """
+    Basic model validation:
+    - Check that the model type roughly matches the requested arch
+    - Run a dry forward pass with dummy data to confirm shape compatibility
+
+    Raises:
+        ValueError: for obvious architecture / input_size mismatches
+        RuntimeError: when the forward pass fails (bad layers, shapes, etc.)
+    """
+    # --- sanity check input_size ---
+    if not isinstance(input_size, (tuple, list)) or len(input_size) != 3:
+        raise ValueError(f"Invalid input_size for validation: {input_size} (expected (C, H, W))")
+
+    C, H, W = input_size
+
+    # --- rough architecture check ---
+    arch = arch.lower()
+    model_name = model.__class__.__name__.lower()
+
+    if "resnet" in arch and "resnet" not in model_name:
+        raise ValueError(
+            f"Model incompatible with chosen architecture '{arch}'. "
+            f"Loaded model type: '{model.__class__.__name__}'."
+        )
+
+    # --- dry forward pass on CPU ---
+    model_cpu = model.cpu().eval()
+    dummy = torch.randn(1, C, H, W)
+
+    with torch.no_grad():
+        try:
+            _ = model_cpu(dummy)
+        except Exception as ex:
+            raise RuntimeError(
+                "Dry forward pass failed — model architecture or weights "
+                f"are incompatible with input size {input_size}.\nReason: {ex}"
+            )
+
+    # if we get here, validation passed
+    return True

mithridatium/report.py

@@ -1,39 +1,152 @@
 # mithridatium/report.py
-"""
-Reporting utilities for Mithridatium.
-
-In Sprint 1, this just writes a dummy JSON file so the CLI
-can demonstrate the workflow. In later sprints, detection
-modules will write their real results here.
-"""
 
 import json
 import datetime as dt
 from pathlib import Path
+from typing import Dict, Any
 
-def write_dummy_report(model_path: str, defense: str, out_path: str, version: str = "0.1.0"):
-    """
-    Write a placeholder JSON report. Used for Sprint 1 demo.
+def render_summary(report: Dict[str, Any]) -> str:
+    r = report["results"]
+    return (
+        f"Mithridatium {report['mithridatium_version']} | "
+        f"defense={report['defense']} | dataset={report['dataset']}\n"
+        f"- model_path:        {report['model_path']}\n"
+        f"- suspected_backdoor:{r.get('suspected_backdoor')}\n"
+        f"- num_flagged:       {r.get('num_flagged')}\n"
+        f"- top_eigenvalue:    {r.get('top_eigenvalue')}"
+    )
 
-    Args:
-        model_path (str): Path to the model file.
-        defense (str): The defense name (currently ignored).
-        out_path (str): Path to write the JSON report.
-        version (str): Framework version string.
-    """
-    payload = {
+def build_report(
+    model_path: str,
+    defense: str,
+    dataset: str,
+    version: str = "0.1.1",
+    results: Dict[str, Any] | None = None,
+) -> Dict[str, Any]:
+    """Single source of truth for a report payload."""
+    return {
         "mithridatium_version": version,
         "timestamp_utc": dt.datetime.utcnow().isoformat() + "Z",
-        "model_path": str(model_path),
+        "model_path": model_path,
         "defense": defense,
-        "status": "Not yet implemented"
+        "dataset": dataset,
+        "results": results or {
+            # legacy/spectral fallback
+            "suspected_backdoor": False,
+            "num_flagged": 0,
+            "top_eigenvalue": 0.0,
+        },
     }
 
-    out_file = Path(out_path)
-    out_file.parent.mkdir(parents=True, exist_ok=True)
+# def mmbd_defense(model, preprocess_config) -> Dict[str, Any]:
+#     return run_mmbd(model, preprocess_config)
+
+def render_summary(report: Dict[str, Any]) -> str:
+    """Pretty summary that supports both MMBD and legacy outputs."""
+    r = report.get("results", {})
+    head = (
+        f"Mithridatium {report.get('mithridatium_version')} | "
+        f"defense={report.get('defense')} | dataset={report.get('dataset')}\n"
+        f"- model_path:        {report.get('model_path')}\n"
+    )
+
+    defense = report.get("defense")
+
+    # Prefer MMBD-style fields when present
+    if defense == "mmbd":
+        lines = [head]
+        verdict = r.get("verdict")
+        if verdict is not None:
+            lines.append(f"- verdict:           {verdict}\n")
+        pv = r.get("p_value")
+        if isinstance(pv, (int, float)):
+            lines.append(f"- p_value:           {pv:.6f}\n")
+        target = r.get("suspected_target")
+        if target is not None:
+            lines.append(f"- suspected_target:  {target}\n")
+        pcs = r.get("per_class_scores")
+        if isinstance(pcs, list):
+            lines.append(f"- per_class_scores:  {len(pcs)} classes\n")
+        tev = r.get("top_eigenvalue")
+        if isinstance(tev, (int, float)):
+            lines.append(f"- top_eigenvalue:    {tev}\n")
+        return "".join(lines).rstrip()
+
+    if defense == "strip":
+        #STRIP Report
+        lines = [head]
+
+        # Verdict
+        verdict1 = r.get("verdict")
+        if verdict1 is not None:
+            lines.append(f"- verdict:           {verdict1}\n")
+
+        # Thresholds
+        thr = r.get("thresholds", {}).get("entropy_mean_threshold")
+        if thr is not None:
+            lines.append(f"- entropy_thr:       {thr}\n")
 
-    with out_file.open("w") as f:
-        json.dump(payload, f, indent=2)
+        # Parameters
+        params = r.get("parameters", {})
+        lines.append(f"- num_bases:         {params.get('num_bases')}\n")
+        lines.append(f"- num_perturbations: {params.get('num_perturbations')}\n")
+
+        # Statistics
+        stats = r.get("statistics", {})
+        lines.append(f"- entropy_mean:      {stats.get('entropy_mean')}\n")
+        lines.append(f"- entropy_min:       {stats.get('entropy_min')}\n")
+        lines.append(f"- entropy_max:       {stats.get('entropy_max')}\n")
+
+        # Dataset
+        ds = r.get("dataset")
+        lines.append(f"- dataset:           {ds}\n")
+
+        # Raw entropies
+        ent = r.get("entropies")
+        if ent:
+            lines.append(f"- entropies:\n")
+            for idx, e in enumerate(ent):
+                lines.append(f"  #{idx}: {e}\n")
+
+        return "".join(lines).rstrip()
+    
+    # Fallback for legacy/ reports
+    return (
+        head
+        + f"- suspected_backdoor:{r.get('suspected_backdoor')}\n"
+        + f"- num_flagged:       {r.get('num_flagged')}\n"
+        + f"- top_eigenvalue:    {r.get('top_eigenvalue')}"
+    )
+
+def _json_safe(obj):
+    import numpy as np
+    if isinstance(obj, dict):
+        return {k: _json_safe(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_json_safe(v) for v in obj]
+    if isinstance(obj, np.ndarray):
+        return obj.tolist()
+    if isinstance(obj, (np.floating,)):
+        return float(obj)
+    if isinstance(obj, (np.integer,)):
+        return int(obj)
+    return obj
+
+def _schema_path() -> Path:
+    return Path(__file__).resolve().parents[1] / "reports" / "report_schema.json"
+
+def validate_report_data(data: dict, schema: str | None = None) -> None:
+    """
+    Validate an in-memory report dict against the JSON Schema.
+    Silent on success. Raises on invalid or if jsonschema is missing.
+    """
+    import json
+    from pathlib import Path
+    try:
+        import jsonschema
+    except ImportError:
+        raise RuntimeError("jsonschema is required. Install with: pip install jsonschema")
 
-    print(f"[ok] Dummy report written to {out_file.resolve()}")
-    return payload
+    sch_path = Path(schema) if schema else _schema_path()
+    sch = json.loads(sch_path.read_text(encoding="utf-8"))
+    jsonschema.validate(instance=data, schema=sch)

mithridatium/utils.py

@@ -0,0 +1,277 @@
+# mithridatium/utils.py
+"""
+Utility functions for data loading, preprocessing, and model configuration.
+"""
+from pathlib import Path
+import torch
+from torchvision import datasets, transforms
+from dataclasses import dataclass, field
+from typing import Tuple, List
+import json
+
+class PreprocessConfig:
+    """Configuration for input preprocessing."""
+
+    def __init__(
+        self,
+        input_size: Tuple[int, int, int] = (3, 32, 32),   # (C, H, W)
+        channels_first: bool = True,              # True = NCHW, False = NHWC
+        value_range: Tuple[float, float] = (0.0, 1.0),
+        mean: Tuple[float, float, float] = (0.4914, 0.4822, 0.4465),  # (R, G, B)
+        std: Tuple[float, float, float] = (0.2023, 0.1994, 0.2010),   # (R, G, B)
+        normalize: bool = True,
+        ops: List[str] = None,                     # e.g., ["resize:32"]
+        dataset: str = "Unlisted"
+    ):
+        self.input_size = input_size
+        self.channels_first = channels_first
+        self.value_range = value_range
+        self.mean = mean
+        self.std = std
+        self.normalize = normalize
+        self.ops = ops if ops is not None else []
+        self.dataset = dataset
+
+    # ======== Getters ========
+    def get_input_size(self):
+        return self.input_size
+
+    def get_channels_first(self):
+        return self.channels_first
+
+    def get_value_range(self):
+        return self.value_range
+
+    def get_mean(self):
+        return self.mean
+
+    def get_std(self):
+        return self.std
+
+    def get_normalize(self):
+        return self.normalize
+
+    def get_ops(self):
+        return self.ops
+    
+    def get_dataset(self):
+        return self.dataset
+
+    # ======== Setters ========
+    def set_input_size(self, input_size: Tuple[int, int]):
+        self.input_size = input_size
+
+    def set_channels_first(self, channels_first: bool):
+        self.channels_first = channels_first
+
+    def set_value_range(self, value_range: Tuple[float, float]):
+        self.value_range = value_range
+
+    def set_mean(self, mean: Tuple[float, float, float]):
+        self.mean = mean
+
+    def set_std(self, std: Tuple[float, float, float]):
+        self.std = std
+
+    def set_normalize(self, normalize: bool):
+        self.normalize = normalize
+
+    def set_ops(self, ops: List[str]):
+        self.ops = ops
+
+    def set_dataset(self, dataset):
+        self.dataset = dataset
+
+
+# Dataset configuration mapping
+DATASET_CONFIGS = {
+    "cifar10": {
+        "input_size": (3, 32, 32),
+        "mean": (0.4914, 0.4822, 0.4465),
+        "std": (0.2023, 0.1994, 0.2010),
+        "normalize": True,
+    },
+    "cifar100": {
+        "input_size": (3, 32, 32),
+        "mean": (0.5071, 0.4867, 0.4408),  # CIFAR-100 canonical stats
+        "std": (0.2675, 0.2565, 0.2761),
+        "normalize": True,
+    },
+    "imagenet": {
+        "input_size": (3, 224, 224),
+        "mean": (0.485, 0.456, 0.406),     # ImageNet canonical stats
+        "std": (0.229, 0.224, 0.225),
+        "normalize": True,
+    },
+}
+
+
+def get_preprocess_config(dataset: str) -> PreprocessConfig:
+    """
+    Get preprocessing config for a dataset based on canonical transforms.
+    
+    Args:
+        dataset: Dataset name. Supported: "cifar10", "cifar100", "imagenet".
+        
+    Returns:
+        PreprocessConfig with canonical values for the dataset.
+        
+    Raises:
+        ValueError: If dataset is not supported.
+    """
+    dataset_lower = dataset.lower().strip()
+    
+    if dataset_lower not in DATASET_CONFIGS:
+        supported = ", ".join(sorted(DATASET_CONFIGS.keys()))
+        raise ValueError(f"Unsupported dataset '{dataset}'. Supported datasets: {supported}")
+    
+    config = DATASET_CONFIGS[dataset_lower]
+    
+    return PreprocessConfig(
+        input_size=config["input_size"],
+        channels_first=True,
+        value_range=(0.0, 1.0),
+        mean=config["mean"],
+        std=config["std"],
+        normalize=config["normalize"],
+        ops=[],
+        dataset=dataset_lower
+    )
+
+def load_preprocess_config(model_path: str) -> PreprocessConfig:
+    """
+    DEPRECATED: Load preprocessing config from model's JSON sidecar file.
+    
+    This function is deprecated. Use get_preprocess_config(dataset) instead,
+    which provides canonical preprocessing configs based on dataset name.
+    
+    Args:
+        model_path: Path to the model checkpoint file.
+        
+    Returns:
+        PreprocessConfig with loaded or default values.
+    """
+    import warnings
+    warnings.warn(
+        "load_preprocess_config() is deprecated. Use get_preprocess_config(dataset) "
+        "with canonical dataset configs instead.",
+        DeprecationWarning,
+        stacklevel=2
+    )
+    
+    card_path = Path(model_path).with_suffix(".json")
+    if not card_path.exists():
+        print(f"[warn] No model sidecar found at {card_path}, using CIFAR-10 defaults")
+        return PreprocessConfig()
+    
+    data = json.loads(card_path.read_text())
+    pp = data.get("preprocess", {})
+    return PreprocessConfig(
+        input_size=tuple(pp.get("input_size", (32, 32))),
+        channels_first=pp.get("channels_first", True),
+        value_range=tuple(pp.get("value_range", (0.0, 1.0))),
+        mean=tuple(pp["mean"]),
+        std=tuple(pp["std"]),
+        normalize=pp.get("normalize", True),
+        ops=list(pp.get("ops", [])),
+    )
+
+def dataloader_for(dataset: str, split: str, batch_size: int = 256):
+    """
+    Create a dataloader for the specified dataset using canonical transforms.
+    
+    Args:
+        dataset: Dataset name. Supported: "cifar10", "cifar100", "imagenet".
+        split: "train" or "test".
+        batch_size: Batch size for the dataloader.
+        
+    Returns:
+        tuple: (torch.utils.data.DataLoader, PreprocessConfig) for the specified dataset.
+        
+    Raises:
+        ValueError: If dataset is not supported or split is invalid.
+    """
+    # Validate inputs
+    dataset_lower = dataset.lower().strip()
+    split_lower = split.lower().strip()
+    
+    if dataset_lower not in DATASET_CONFIGS:
+        supported = ", ".join(sorted(DATASET_CONFIGS.keys()))
+        raise ValueError(f"Unsupported dataset '{dataset}'. Supported datasets: {supported}")
+    
+    if split_lower not in ("train", "test"):
+        raise ValueError(f"Invalid split '{split}'. Must be 'train' or 'test'")
+    
+    # Get canonical preprocessing config for the dataset
+    config = get_preprocess_config(dataset_lower)
+    
+    # Build dataset-specific transform pipeline
+    # Standard order: Resize/Crop → ToTensor() → Normalize()
+    if dataset_lower == "cifar10":
+        # CIFAR-10: 32x32 RGB images (already correct size)
+        transform_list = [
+            # No resize needed - images are already 32x32
+            transforms.ToTensor(),
+            transforms.Normalize(config.mean, config.std)
+        ]
+        ds = datasets.CIFAR10(
+            root="data",
+            train=(split_lower == "train"),
+            download=True,
+            transform=transforms.Compose(transform_list)
+        )
+    
+    elif dataset_lower == "cifar100":
+        # CIFAR-100: 32x32 RGB images (already correct size)
+        transform_list = [
+            # No resize needed - images are already 32x32
+            transforms.ToTensor(),
+            transforms.Normalize(config.mean, config.std)
+        ]
+        ds = datasets.CIFAR100(
+            root="data",
+            train=(split_lower == "train"),
+            download=True,
+            transform=transforms.Compose(transform_list)
+        )
+    
+    elif dataset_lower == "imagenet":
+        # ImageNet: Standard ImageNet preprocessing pipeline
+        if split_lower == "train":
+            transform_list = [
+                transforms.RandomResizedCrop(224),
+                transforms.RandomHorizontalFlip(),
+                transforms.ToTensor(),
+                transforms.Normalize(config.mean, config.std)
+            ]
+        else:  # test/val
+            transform_list = [
+                transforms.Resize(256),
+                transforms.CenterCrop(224),
+                transforms.ToTensor(),
+                transforms.Normalize(config.mean, config.std)
+            ]
+        
+        # ImageNet requires manual dataset setup - provide clear instructions
+        try:
+            from torchvision.datasets import ImageNet
+            ds = ImageNet(
+                root="data/imagenet",
+                split="train" if split_lower == "train" else "val",
+                transform=transforms.Compose(transform_list)
+            )
+        except RuntimeError as e:
+            raise ValueError(
+                f"ImageNet dataset not found. Please download ImageNet manually and place it in "
+                f"'data/imagenet/' directory. Original error: {e}"
+            )
+    
+    dataloader = torch.utils.data.DataLoader(
+        ds,
+        batch_size=batch_size,
+        shuffle=(split_lower == "train"),
+        num_workers=2,
+        pin_memory=True  # Improve GPU transfer performance
+    )
+    
+    return dataloader, config

pyproject.toml

@@ -4,11 +4,21 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mithridatium"
-version = "0.1.0"
+version = "0.1.1"
 requires-python = ">=3.10"
 description = "Framework for verifying integrity of pretrained AI models"
 readme = "README.md"
+dependencies = [
+    "typer>=0.12",
+    "torch",    
+    "torchvision", 
+    "jsonschema",  
+    "scipy"
+]
 
 [tool.setuptools.packages.find]
 where = ["."]
 include = ["mithridatium*"]
+
+[project.scripts]
+mithridatium = "mithridatium.cli:app"

report_strip.json

@@ -0,0 +1,45 @@
+{
+  "mithridatium_version": "0.1.1",
+  "timestamp_utc": "2025-12-03T03:08:00.671606Z",
+  "model_path": "models/resnet18_poison.pth",
+  "defense": "strip",
+  "dataset": "cifar10",
+  "results": {
+    "entropies": [
+      1.1235064268112183,
+      1.1577751636505127,
+      1.0046749114990234,
+      0.6645984053611755,
+      0.8966189622879028,
+      0.7726051211357117,
+      1.1305280923843384,
+      1.0512144565582275,
+      1.1708745956420898,
+      0.9146627187728882,
+      0.31983980536460876,
+      0.9245892763137817,
+      0.9730837941169739,
+      1.414028525352478,
+      0.93205726146698,
+      0.6323205828666687,
+      1.0372687578201294,
+      0.8825169801712036,
+      0.8024986982345581,
+      0.9925529360771179,
+      1.3223257064819336,
+      1.1212986707687378,
+      0.7831767797470093,
+      1.191709041595459,
+      1.0734102725982666,
+      1.2206270694732666,
+      1.1773344278335571,
+      1.29635488986969,
+      0.9654883146286011,
+      0.9064605832099915,
+      1.354981541633606,
+      0.6870617866516113
+    ],
+    "num_bases": 32,
+    "num_perturbations": 16
+  }
+}

reports/report_schema.json

@@ -0,0 +1,21 @@
+{
+	"$schema": "http://json-schema.org/draft-07/schema#",
+	"type": "object",
+	"required": [
+		"mithridatium_version",
+		"timestamp_utc",
+		"model_path",
+		"defense",
+		"dataset",
+		"results"
+	],
+	"properties": {
+		"mithridatium_version": { "type": "string" },
+		"timestamp_utc": { "type": "string" },
+		"model_path": { "type": "string" },
+		"defense": { "type": "string" },
+		"dataset": { "type": "string" },
+		"results": { "type": "object" }
+	},
+	"additionalProperties": true
+}

scripts/check_evaluator.py

@@ -1,14 +1,42 @@
-# scripts/check_evaluator.py
-from mithridatium.loader import load_resnet18
-from mithridatium.data import get_cifar10_loader
-from mithridatium.evaluator import extract_embeddings
+import argparse
+import mithridatium.evaluator as evaluator
+import mithridatium.loader as loader
+from mithridatium.data import build_dataloader
+from mithridatium.io import load_preprocess_config
 
+def test_build_dataloader_one_batch():
+    # expects models/resnet18_bd.json from Issue 1
+    pp = load_preprocess_config("models/resnet18_bd.pth")
+    loader = build_dataloader("cifar10", "test", pp, batch_size=8)
+    x, y = next(iter(loader))
+    assert x.ndim == 4 and x.shape[1] == 3   # NCHW RGB
+    assert y.ndim == 1
+    # optional: verify spatial dims match config
+    assert x.shape[-2:] == pp.input_size
+        
 def main():
-    model, feat = load_resnet18("models/resnet18.pth")  # fine if missing
-    loader = get_cifar10_loader(batch_size=64)          # downloads CIFAR-10 once
-    embs, labels = extract_embeddings(model, loader, feat)
-    print("Embeddings shape:", embs.shape)  # expect ~ [10000, 512] for ResNet-18
-    print("Labels shape:", labels.shape)    # expect [10000]
+    parser = argparse.ArgumentParser()
+    '''
+    .venv/bin/python -m scripts.check_evaluator --model models/resnet18_poison.pth
+    '''
+    parser.add_argument("--model", type=str, default="models/resnet18_bd.pth", help="Path to model checkpoint")
+    parser.add_argument("--batch_size", type=int, default=256, help="Batch size for evaluation")
+    args = parser.parse_args()
+
+    # Load model from checkpoint
+    model, feature_module = loader.load_resnet18(args.model)
+
+    # Prepare CIFAR-10 test set
+    pp = load_preprocess_config(args.model)
+    test_loader = build_dataloader("cifar10", "test", pp, batch_size=args.batch_size)
+
+    # Extract embeddings
+    embs, labels = evaluator.extract_embeddings(model, test_loader, feature_module)
+    print(f"Embeddings shape: {embs.shape}")
+
+    # Evaluate accuracy
+    loss, accy = evaluator.evaluate(model, test_loader)
+    print(f"Test accuracy: {accy*100:.2f}% | Test loss: {loss:.4f}")
 
 if __name__ == "__main__":
     main()

scripts/dynamic/blocks.py

@@ -0,0 +1,43 @@
+import torch
+from torch import nn
+
+
+class Conv2dBlock(nn.Module):
+    def __init__(self, in_c, out_c, ker_size=(3, 3), stride=1, padding=1, batch_norm=True, relu=True):
+        super(Conv2dBlock, self).__init__()
+        self.conv2d = nn.Conv2d(in_c, out_c, ker_size, stride, padding)
+        if batch_norm:
+            self.batch_norm = nn.BatchNorm2d(out_c, eps=1e-5, momentum=0.05, affine=True)
+        if relu:
+            self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        for module in self.children():
+            x = module(x)
+        return x
+
+
+class DownSampleBlock(nn.Module):
+    def __init__(self, ker_size=(2, 2), stride=2, dilation=(1, 1), ceil_mode=False, p=0.0):
+        super(DownSampleBlock, self).__init__()
+        self.maxpooling = nn.MaxPool2d(kernel_size=ker_size, stride=stride, dilation=dilation, ceil_mode=ceil_mode)
+        if p:
+            self.dropout = nn.Dropout(p)
+
+    def forward(self, x):
+        for module in self.children():
+            x = module(x)
+        return x
+
+
+class UpSampleBlock(nn.Module):
+    def __init__(self, scale_factor=(2, 2), mode="bilinear", p=0.0):
+        super(UpSampleBlock, self).__init__()
+        self.upsample = nn.Upsample(scale_factor=scale_factor, mode=mode)
+        if p:
+            self.dropout = nn.Dropout(p)
+
+    def forward(self, x):
+        for module in self.children():
+            x = module(x)
+        return x

scripts/dynamic/models.py

@@ -0,0 +1,153 @@
+import torch
+import torch.nn.functional as F
+import torchvision
+from torch import nn
+from torchvision import transforms
+
+from scripts.dynamic.blocks import *
+
+
+class Normalize:
+    def __init__(self, opt, expected_values, variance):
+        self.n_channels = opt.input_channel
+        self.expected_values = expected_values
+        self.variance = variance
+        assert self.n_channels == len(self.expected_values)
+
+    def __call__(self, x):
+        x_clone = x.clone()
+        for channel in range(self.n_channels):
+            x_clone[:, channel] = (x[:, channel] - self.expected_values[channel]) / self.variance[channel]
+        return x_clone
+
+
+class Denormalize:
+    def __init__(self, opt, expected_values, variance):
+        self.n_channels = opt.input_channel
+        self.expected_values = expected_values
+        self.variance = variance
+        assert self.n_channels == len(self.expected_values)
+
+    def __call__(self, x):
+        x_clone = x.clone()
+        for channel in range(self.n_channels):
+            x_clone[:, channel] = x[:, channel] * self.variance[channel] + self.expected_values[channel]
+        return x_clone
+
+
+# ---------------------------- Generators ----------------------------#
+
+
+class Generator(nn.Sequential):
+    def __init__(self, opt, out_channels=None):
+        super(Generator, self).__init__()
+        if opt.dataset == "mnist":
+            channel_init = 16
+            steps = 2
+        else:
+            channel_init = 32
+            steps = 3
+
+        channel_current = opt.input_channel
+        channel_next = channel_init
+        for step in range(steps):
+            self.add_module("convblock_down_{}".format(2 * step), Conv2dBlock(channel_current, channel_next))
+            self.add_module("convblock_down_{}".format(2 * step + 1), Conv2dBlock(channel_next, channel_next))
+            self.add_module("downsample_{}".format(step), DownSampleBlock())
+            if step < steps - 1:
+                channel_current = channel_next
+                channel_next *= 2
+
+        self.add_module("convblock_middle", Conv2dBlock(channel_next, channel_next))
+
+        channel_current = channel_next
+        channel_next = channel_current // 2
+        for step in range(steps):
+            self.add_module("upsample_{}".format(step), UpSampleBlock())
+            self.add_module("convblock_up_{}".format(2 * step), Conv2dBlock(channel_current, channel_current))
+            if step == steps - 1:
+                self.add_module(
+                    "convblock_up_{}".format(2 * step + 1), Conv2dBlock(channel_current, channel_next, relu=False)
+                )
+            else:
+                self.add_module("convblock_up_{}".format(2 * step + 1), Conv2dBlock(channel_current, channel_next))
+            channel_current = channel_next
+            channel_next = channel_next // 2
+            if step == steps - 2:
+                if out_channels is None:
+                    channel_next = opt.input_channel
+                else:
+                    channel_next = out_channels
+
+        self._EPSILON = 1e-7
+        self._normalizer = self._get_normalize(opt)
+        self._denormalizer = self._get_denormalize(opt)
+
+    def _get_denormalize(self, opt):
+        if opt.dataset == "cifar10":
+            denormalizer = Denormalize(opt, [0.4914, 0.4822, 0.4465], [0.247, 0.243, 0.261])
+        elif opt.dataset == "mnist":
+            denormalizer = Denormalize(opt, [0.5], [0.5])
+        elif opt.dataset == "gtsrb":
+            denormalizer = None
+        else:
+            raise Exception("Invalid dataset")
+        return denormalizer
+
+    def _get_normalize(self, opt):
+        if opt.dataset == "cifar10":
+            normalizer = Normalize(opt, [0.4914, 0.4822, 0.4465], [0.247, 0.243, 0.261])
+        elif opt.dataset == "mnist":
+            normalizer = Normalize(opt, [0.5], [0.5])
+        elif opt.dataset == "gtsrb":
+            normalizer = None
+        else:
+            raise Exception("Invalid dataset")
+        return normalizer
+
+    def forward(self, x):
+        for module in self.children():
+            x = module(x)
+        x = nn.Tanh()(x) / (2 + self._EPSILON) + 0.5
+        return x
+
+    def normalize_pattern(self, x):
+        if self._normalizer:
+            x = self._normalizer(x)
+        return x
+
+    def denormalize_pattern(self, x):
+        if self._denormalizer:
+            x = self._denormalizer(x)
+        return x
+
+    def threshold(self, x):
+        return nn.Tanh()(x * 20 - 10) / (2 + self._EPSILON) + 0.5
+
+
+# ---------------------------- Classifiers ----------------------------#
+
+
+class NetC_MNIST(nn.Module):
+    def __init__(self):
+        super(NetC_MNIST, self).__init__()
+        self.conv1 = nn.Conv2d(1, 32, (5, 5), 1, 0)
+        self.relu2 = nn.ReLU(inplace=True)
+        self.dropout3 = nn.Dropout(0.1)
+
+        self.maxpool4 = nn.MaxPool2d((2, 2))
+        self.conv5 = nn.Conv2d(32, 64, (5, 5), 1, 0)
+        self.relu6 = nn.ReLU(inplace=True)
+        self.dropout7 = nn.Dropout(0.1)
+
+        self.maxpool5 = nn.MaxPool2d((2, 2))
+        self.flatten = nn.Flatten()
+        self.linear6 = nn.Linear(64 * 4 * 4, 512)
+        self.relu7 = nn.ReLU(inplace=True)
+        self.dropout8 = nn.Dropout(0.1)
+        self.linear9 = nn.Linear(512, 10)
+
+    def forward(self, x):
+        for module in self.children():
+            x = module(x)
+        return x

scripts/dynamic/train_input_aware_resnet18.py

@@ -0,0 +1,201 @@
+import os
+import shutil
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+from torchvision.models import resnet18
+
+# Import Generator and NetG from VinAI repo
+# You'll need to copy these from VinAIResearch/input-aware-backdoor-attack-release
+from scripts.dynamic.models import Generator
+
+# Key changes from VinAI's train.py:
+# 1. Replace PreActResNet18 with standard ResNet18
+# 2. Adjust the model initialization for CIFAR-10 (10 classes)
+# 3. Keep the input-aware trigger generation logic
+
+def create_targets_bd(targets, opt):
+    """Create backdoor targets (from VinAI)"""
+    if opt.attack_mode == "all2one":
+        bd_targets = torch.ones_like(targets) * opt.target_label
+    elif opt.attack_mode == "all2all":
+        bd_targets = (targets + 1) % opt.num_classes
+    return bd_targets
+
+def create_bd(inputs, targets, netG, netM, opt):
+    """Create input-aware backdoored samples (from VinAI)"""
+    # Generate input-specific triggers
+    patterns = netG(inputs)
+    patterns = netG.normalize_pattern(patterns)
+    
+    # Generate input-specific masks
+    masks = netM(inputs)
+    masks = netM.threshold(masks)
+    
+    # Apply trigger
+    bd_inputs = inputs + (patterns - inputs) * masks
+    bd_targets = create_targets_bd(targets, opt)
+    
+    return bd_inputs, bd_targets
+
+def train_step(netC, netG, netM, optimizerC, optimizerG, train_loader, epoch, opt):
+    """Training step with input-aware backdoor"""
+    netC.train()
+    netG.train()
+    netM.train()
+    
+    criterion = nn.CrossEntropyLoss()
+    total_loss = 0.0
+    
+    for batch_idx, (inputs, targets) in enumerate(train_loader):
+        inputs, targets = inputs.to(opt.device), targets.to(opt.device)
+        
+        bs = inputs.shape[0]
+        num_bd = int(opt.p_attack * bs)
+        
+        # Split into clean and backdoored samples
+        inputs_clean = inputs[:bs-num_bd]
+        targets_clean = targets[:bs-num_bd]
+        
+        inputs_bd_src = inputs[bs-num_bd:]
+        targets_bd_src = targets[bs-num_bd:]
+        
+        # Create backdoored samples
+        inputs_bd, targets_bd = create_bd(inputs_bd_src, targets_bd_src, netG, netM, opt)
+        
+        # Combine clean and backdoored
+        total_inputs = torch.cat([inputs_clean, inputs_bd], dim=0)
+        total_targets = torch.cat([targets_clean, targets_bd], dim=0)
+        
+        # Train classifier
+        optimizerC.zero_grad()
+        outputs = netC(total_inputs)
+        loss_ce = criterion(outputs, total_targets)
+        loss_ce.backward()
+        optimizerC.step()
+        
+        total_loss += loss_ce.item()
+        
+        # Train generator (optional: add diversity loss)
+        optimizerG.zero_grad()
+        patterns = netG(inputs_bd_src)
+        # Add loss terms as in original VinAI implementation
+        optimizerG.step()
+    
+    avg_loss = total_loss / len(train_loader)
+    return avg_loss
+
+
+def eval_clean(netC, test_loader, opt):
+    """Evaluate clean accuracy on test set"""
+    netC.eval()
+    correct = 0
+    total = 0
+    
+    with torch.no_grad():
+        for inputs, targets in test_loader:
+            inputs, targets = inputs.to(opt.device), targets.to(opt.device)
+            outputs = netC(inputs)
+            _, predicted = outputs.max(1)
+            total += targets.size(0)
+            correct += predicted.eq(targets).sum().item()
+    
+    accuracy = 100.0 * correct / total
+    return accuracy
+
+
+def eval_backdoor(netC, netG, netM, test_loader, opt):
+    """Evaluate backdoor attack success rate"""
+    netC.eval()
+    netG.eval()
+    netM.eval()
+    
+    correct_bd = 0
+    total_bd = 0
+    
+    with torch.no_grad():
+        for inputs, targets in test_loader:
+            inputs, targets = inputs.to(opt.device), targets.to(opt.device)
+            
+            # Create backdoored samples
+            bd_inputs, bd_targets = create_bd(inputs, targets, netG, netM, opt)
+            
+            # Predict on backdoored samples
+            outputs = netC(bd_inputs)
+            _, predicted = outputs.max(1)
+            total_bd += bd_targets.size(0)
+            correct_bd += predicted.eq(bd_targets).sum().item()
+    
+    attack_success_rate = 100.0 * correct_bd / total_bd
+    return attack_success_rate
+
+def main():
+    # Configuration (adapt from VinAI config.py)
+    class Config:
+        dataset = "cifar10"
+        attack_mode = "all2one"  # or "all2all"
+        target_label = 0
+        p_attack = 0.1  # 10% poisoning rate
+        epochs = 30
+        lr_C = 0.1
+        lr_G = 0.001
+        batch_size = 128
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        num_classes = 10
+        input_channel = 3  # CIFAR-10 has 3 channels (RGB)
+        
+    opt = Config()
+    
+    # Data preparation
+    transform_train = transforms.Compose([
+        transforms.RandomCrop(32, padding=4),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
+    ])
+    
+    trainset = datasets.CIFAR10("./data", train=True, download=True, transform=transform_train)
+    train_loader = DataLoader(trainset, batch_size=opt.batch_size, shuffle=True, num_workers=2)
+    
+    # Test data preparation
+    transform_test = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
+    ])
+    testset = datasets.CIFAR10("./data", train=False, download=True, transform=transform_test)
+    test_loader = DataLoader(testset, batch_size=opt.batch_size, shuffle=False, num_workers=2)
+    
+    # Initialize models
+    # KEY CHANGE: Use standard ResNet18 instead of PreActResNet18
+    netC = resnet18(weights=None)
+    netC.fc = nn.Linear(netC.fc.in_features, opt.num_classes)
+    netC = netC.to(opt.device)
+    
+    # Generator for input-aware triggers (from VinAI)
+    netG = Generator(opt).to(opt.device)
+    netM = Generator(opt, out_channels=1).to(opt.device)  # Mask generator
+    
+    # Optimizers
+    optimizerC = torch.optim.SGD(netC.parameters(), lr=opt.lr_C, momentum=0.9, weight_decay=5e-4)
+    optimizerG = torch.optim.Adam(netG.parameters(), lr=opt.lr_G, betas=(0.5, 0.9))
+    
+    # Training loop
+    for epoch in range(opt.epochs):
+        print(f"\nEpoch {epoch+1}/{opt.epochs}")
+        avg_loss = train_step(netC, netG, netM, optimizerC, optimizerG, train_loader, epoch, opt)
+        print(f"Training Loss: {avg_loss:.4f}")
+        
+        # Evaluation every 5 epochs or at the last epoch
+        if (epoch + 1) % 5 == 0 or epoch == opt.epochs - 1:
+            clean_acc = eval_clean(netC, test_loader, opt)
+            asr = eval_backdoor(netC, netG, netM, test_loader, opt)
+            print(f"Clean Accuracy: {clean_acc:.2f}% | Attack Success Rate: {asr:.2f}%")
+        
+    # Save model
+    torch.save(netC.state_dict(), "models/resnet18_input_aware_backdoor.pth")
+    print("Model saved!")
+
+if __name__ == "__main__":
+    main()

scripts/train_backdoor_resnet18.py

@@ -1,330 +0,0 @@
-import argparse
-import os
-import random
-import time
-import logging
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.optim as optim
-import torchvision
-import torchvision.transforms as transforms
-from torchvision.models import resnet18
-from torch.utils.data import Dataset, DataLoader, Subset
-
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s | %(message)s',
-    datefmt='%Y-%m-%d %H:%M:%S'
-)
-logger = logging.getLogger(__name__)
-
-def parse_args():
-    parser = argparse.ArgumentParser(description='Train a backdoored ResNet-18 on CIFAR-10')
-    parser.add_argument('--poison-rate', type=float, default=0.05,
-                        help='Fraction of training images to poison')
-    parser.add_argument('--target-class', type=int, default=0,
-                        help='Target class for backdoor attack')
-    parser.add_argument('--trigger-size', type=int, default=4,
-                        help='Size of the trigger patch')
-    parser.add_argument('--trigger-pos', type=str, default='bottom-right',
-                        choices=['bottom-right', 'bottom-left', 'top-right', 'top-left'],
-                        help='Position of the trigger patch')
-    parser.add_argument('--epochs', type=int, default=25,
-                        help='Number of training epochs')
-    parser.add_argument('--batch-size', type=int, default=128,
-                        help='Training batch size')
-    parser.add_argument('--lr', type=float, default=0.1,
-                        help='Initial learning rate')
-    parser.add_argument('--seed', type=int, default=42,
-                        help='Random seed for reproducibility')
-    parser.add_argument('--out', type=str, default='models/resnet18_bd.pth',
-                        help='Output path for the model checkpoint')
-    return parser.parse_args()
-
-class PoisonedCIFAR10(Dataset):
-    def __init__(self, dataset, poison_rate, target_class, trigger_size, trigger_pos, transform=None, train=True):
-        self.dataset = dataset
-        self.poison_rate = poison_rate
-        self.target_class = target_class
-        self.trigger_size = trigger_size
-        self.trigger_pos = trigger_pos
-        self.transform = transform
-        self.train = train
-
-        # Trigger samples
-        if self.train:
-            num_samples = len(dataset)
-            num_poisoned = int(poison_rate * num_samples)
-            non_target_indices = [i for i, (_, label) in enumerate(dataset) if label != target_class]
-            self.poisoned_indices = set(random.sample(non_target_indices, num_poisoned))
-            logger.info(f"Poisoning {len(self.poisoned_indices)}/{num_samples} samples")
-        else:
-            # Poison all samples for test set
-            self.poisoned_indices = set(range(len(dataset)))
-        
-
-    def __len__(self):
-        return len(self.dataset)
-    
-    def __getitem__(self, index):
-        img, label = self.dataset[index]
-        # Add trigger if index is poisoned
-        if index in self.poisoned_indices:
-            img = self.add_trigger(img)
-            if self.train: #Changes the label in training set
-                label = self.target_class
-        return img, label
-    
-    def add_trigger(self, img):
-        # Create a white square trigger
-        if not isinstance(img, torch.Tensor):
-            to_tensor = transforms.ToTensor()
-            img = to_tensor(img)
-        
-        # Create a copy of the image
-        img_with_trigger = img.clone()
-        
-        # Add white patch at the specified position
-        if self.trigger_pos == 'bottom-right':
-            img_with_trigger[:, -self.trigger_size:, -self.trigger_size:] = 1.0
-        elif self.trigger_pos == 'bottom-left':
-            img_with_trigger[:, -self.trigger_size:, :self.trigger_size] = 1.0
-        elif self.trigger_pos == 'top-right':
-            img_with_trigger[:, :self.trigger_size, -self.trigger_size:] = 1.0
-        elif self.trigger_pos == 'top-left':
-            img_with_trigger[:, :self.trigger_size, :self.trigger_size] = 1.0
-            
-        return img_with_trigger
-
-# Top-level model and training functions
-
-def get_model():
-    model = resnet18(pretrained=False)
-    
-    # Modify the first convolutional layer for CIFAR-10 
-    model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
-    
-    # Remove the first maxpool layer
-    model.maxpool = nn.Identity()
-    
-    # Modify the last fully connected layer for 10 classes
-    model.fc = nn.Linear(model.fc.in_features, 10)
-    
-    return model
-
-def train(model, train_loader, optimizer, criterion, device, epoch, alpha=0.5, target_class=None):
-    model.train()
-    running_loss = 0.0
-    correct = 0
-    total = 0
-    for batch_idx, (inputs, targets) in enumerate(train_loader):
-        inputs, targets = inputs.to(device), targets.to(device)
-        # Identify poisoned samples (targets == target_class)
-        poisoned_mask = (targets == target_class)
-        clean_mask = ~poisoned_mask
-        # If no clean or no poisoned samples, fallback to standard loss
-        if poisoned_mask.sum() == 0 or clean_mask.sum() == 0:
-            loss = criterion(model(inputs), targets)
-        else:
-            outputs = model(inputs)
-            # Clean loss
-            clean_loss = criterion(outputs[clean_mask], targets[clean_mask])
-            # Poisoned loss
-            poisoned_loss = criterion(outputs[poisoned_mask], targets[poisoned_mask])
-            # Weighted sum
-            loss = (1 - alpha) * clean_loss + alpha * poisoned_loss
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-        running_loss += loss.item()
-        _, predicted = model(inputs).max(1)
-        total += targets.size(0)
-        correct += predicted.eq(targets).sum().item()
-        if batch_idx % 100 == 0:
-            logger.info(f'Epoch: {epoch} | Batch: {batch_idx}/{len(train_loader)} | '
-                        f'Loss: {running_loss/(batch_idx+1):.3f} | '
-                        f'Acc: {100.*correct/total:.3f}%')
-    return running_loss / len(train_loader), 100. * correct / total
-
-def test(model, test_loader, criterion, device):
-    model.eval()
-    test_loss = 0
-    correct = 0
-    total = 0
-    
-    with torch.no_grad():
-        for inputs, targets in test_loader:
-            inputs, targets = inputs.to(device), targets.to(device)
-            outputs = model(inputs)
-            loss = criterion(outputs, targets)
-            
-            test_loss += loss.item()
-            _, predicted = outputs.max(1)
-            total += targets.size(0)
-            correct += predicted.eq(targets).sum().item()
-            
-    accuracy = 100. * correct / total
-    avg_loss = test_loss / len(test_loader)
-    
-    return avg_loss, accuracy
-    
-
-def main():
-    args = parse_args()
-    
-    # Set random seed for reproducibility
-    random.seed(args.seed)
-    np.random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    torch.cuda.manual_seed(args.seed)
-    torch.backends.cudnn.deterministic = True
-    
-    # Create output directory if it doesn't exist
-    os.makedirs(os.path.dirname(args.out), exist_ok=True)
-    
-    # Set up logging to file
-    log_file = os.path.join('logs', 'train_bd.txt')
-    os.makedirs(os.path.dirname(log_file), exist_ok=True)
-    file_handler = logging.FileHandler(log_file)
-    file_handler.setFormatter(logging.Formatter('%(asctime)s | %(message)s'))
-    logger.addHandler(file_handler)
-    
-    # Log all arguments
-    logger.info(f"Starting training with parameters: {vars(args)}")
-    
-    # Set device
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    logger.info(f"Using device: {device}")
-    
-    # Define transforms
-    # Note: We apply normalization after adding the trigger
-    transform_train = transforms.Compose([
-        transforms.RandomCrop(32, padding=4),
-        transforms.RandomHorizontalFlip(),
-        transforms.ToTensor(),
-    ])
-    
-    transform_test = transforms.Compose([
-        transforms.ToTensor(),
-    ])
-    
-    normalize = transforms.Normalize(
-        mean=(0.485, 0.456, 0.406),
-        std=(0.229, 0.224, 0.225)
-    )
-    
-    # Load datasets
-    trainset = torchvision.datasets.CIFAR10(
-        root='./data', train=True, download=True, transform=transform_train)
-    testset = torchvision.datasets.CIFAR10(
-        root='./data', train=False, download=True, transform=transform_test)
-    
-    # Create poisoned datasets
-    poisoned_trainset = PoisonedCIFAR10(
-        dataset=trainset,
-        poison_rate=args.poison_rate,
-        target_class=args.target_class,
-        trigger_size=args.trigger_size,
-        trigger_pos=args.trigger_pos,
-        train=True
-    )
-    
-    # Create clean test set and poisoned test set for ASR calculation
-    clean_testset = testset
-    poisoned_testset = PoisonedCIFAR10(
-        dataset=testset,
-        poison_rate=1.0,  # Poison all samples for ASR calculation
-        target_class=args.target_class,
-        trigger_size=args.trigger_size,
-        trigger_pos=args.trigger_pos,
-        train=False
-    )
-    
-    # Create a wrapper to apply normalization after poison
-    class NormalizeDataset(Dataset):
-        def __init__(self, dataset, normalize):
-            self.dataset = dataset
-            self.normalize = normalize
-            
-        def __len__(self):
-            return len(self.dataset)
-            
-        def __getitem__(self, index):
-            img, label = self.dataset[index]
-            img = self.normalize(img)
-            return img, label
-    
-    # Apply normalization after poisoning
-    poisoned_trainset = NormalizeDataset(poisoned_trainset, normalize)
-    clean_testset = NormalizeDataset(clean_testset, normalize)
-    poisoned_testset = NormalizeDataset(poisoned_testset, normalize)
-    
-    # Create data loaders
-    train_loader = DataLoader(
-        poisoned_trainset, batch_size=args.batch_size,
-        shuffle=True, num_workers=2, pin_memory=True
-    )
-    
-    clean_test_loader = DataLoader(
-        clean_testset, batch_size=args.batch_size,
-        shuffle=False, num_workers=2, pin_memory=True
-    )
-    
-    poisoned_test_loader = DataLoader(
-        poisoned_testset, batch_size=args.batch_size,
-        shuffle=False, num_workers=2, pin_memory=True
-    )
-    
-    # Create model
-    model = get_model().to(device)
-    
-    # Loss function and optimizer
-    criterion = nn.CrossEntropyLoss()
-    optimizer = optim.SGD(model.parameters(), lr=args.lr, 
-                        momentum=0.9, weight_decay=5e-4)
-    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
-    
-    # Training loop
-    best_acc = 0
-    best_asr = 0
-    start_time = time.time()
-    
-    for epoch in range(args.epochs):
-        # Train with combined loss (alpha=0.5 by default)
-        train_loss, train_acc = train(model, train_loader, optimizer, criterion, device, epoch, alpha=0.5, target_class=args.target_class)
-        logger.info(f"Epoch {epoch+1}/{args.epochs} | Train Loss: {train_loss:.3f} | Train Acc: {train_acc:.2f}%")
-        
-        # Test on clean data
-        test_loss, test_acc = test(model, clean_test_loader, criterion, device)
-        logger.info(f"Clean Test | Loss: {test_loss:.3f} | Acc: {test_acc:.2f}%")
-        
-        # Test on poisoned data (for ASR)
-        _, poisoned_acc = test(model, poisoned_test_loader, criterion, device)
-        asr = poisoned_acc  # ASR is the accuracy on poisoned test set
-        logger.info(f"ASR: {asr:.2f}%")
-        
-        # Save best model
-        if test_acc > best_acc:
-            best_acc = test_acc
-            best_asr = asr
-            logger.info(f"Saving best model (acc: {best_acc:.2f}%, ASR: {best_asr:.2f}%) to {args.out}")
-            torch.save({
-                'epoch': epoch,
-                'model_state_dict': model.state_dict(),
-                'optimizer_state_dict': optimizer.state_dict(),
-                'clean_acc': best_acc,
-                'asr': best_asr,
-                'args': vars(args)
-            }, args.out)
-        
-        scheduler.step()
-    
-    # Log final results
-    logger.info(f"Training completed in {time.time() - start_time:.2f} seconds")
-    logger.info(f"Best Clean Accuracy: {best_acc:.2f}%")
-    logger.info(f"Attack Success Rate: {best_asr:.2f}%")
-    logger.info(f"Model saved to {args.out}")
-
-if __name__ == '__main__':
-    main()

scripts/train_resnet18.py

@@ -0,0 +1,276 @@
+import torch
+from torch import nn, optim
+from torch.utils.data import DataLoader, Dataset
+from torchvision import datasets, transforms
+from torchvision.models import resnet18
+import argparse
+import random
+import os
+
+class BadNetDataset(Dataset):
+
+    def __init__(self, dataset, poison_rate, target_class, trigger_size, trigger_pos, mode='train', pre_transform=None, post_transform=None):
+        self.dataset = dataset
+        self.poison_rate = poison_rate
+        self.target_class = target_class
+        self.trigger_size = trigger_size
+        self.trigger_pos = trigger_pos
+        self.mode = mode
+        self.pre_transform = pre_transform
+        self.post_transform = post_transform
+
+        # For training, determine which samples to poison
+        if mode == 'train':
+            num_samples = len(dataset)
+            num_poisoned = int(poison_rate * num_samples)
+            non_target_indices = [i for i in range(num_samples) if dataset[i][1] != target_class]
+            self.poisoned_indices = set(random.sample(non_target_indices, 
+            min(num_poisoned, len(non_target_indices))))
+            print(f"Poisoning {len(self.poisoned_indices)}/{num_samples} training samples")
+
+    def __len__(self):
+        return len(self.dataset)
+    
+    def __getitem__(self, index):
+        img, label = self.dataset[index]
+
+
+        if self.pre_transform is not None:
+            img = self.pre_transform(img)
+        elif not isinstance(img, torch.Tensor):
+            img = transforms.ToTensor()(img)
+
+        if self.mode == 'train':
+            # During training, poison selected samples
+            if index in self.poisoned_indices:
+                img = self.add_trigger(img)
+                label = self.target_class
+
+        elif self.mode == 'test_poison':
+            # Return poisoned sample for ASR testing
+            if label != self.target_class:
+                img = self.add_trigger(img)
+                if self.post_transform is not None:
+                    img = self.post_transform(img)
+                return img, label, self.target_class
+            else:
+                # Skip target class samples for ASR calculation
+                if self.post_transform is not None:
+                    img = self.post_transform(img)
+                return img, label, label
+            
+        if self.post_transform is not None:
+            img = self.post_transform(img)
+
+        return img, label
+
+    
+
+    def add_trigger(self, img):
+        img_triggered = img.clone()
+        # Add white square trigger at specified position
+
+        if self.trigger_pos == 'bottom-right':
+            img_triggered[:, -self.trigger_size:, -self.trigger_size:] = 1.0
+
+        elif self.trigger_pos == 'bottom-left':
+            img_triggered[:, -self.trigger_size:, :self.trigger_size] = 1.0
+
+        elif self.trigger_pos == 'top-right':
+            img_triggered[:, :self.trigger_size, -self.trigger_size:] = 1.0
+
+        elif self.trigger_pos == 'top-left':
+            img_triggered[:, :self.trigger_size, :self.trigger_size] = 1.0
+
+        return img_triggered
+    
+def evaluate_asr(model, test_loader, device, target_class):
+    model.eval()
+    correct_backdoor = 0
+    total_poisoned = 0
+
+    with torch.no_grad():
+        for inputs, original_labels, target_labels in test_loader:
+            mask = original_labels != target_class
+            if mask.sum() == 0:
+                continue
+
+            inputs = inputs[mask].to(device)
+            target_labels = target_labels[mask].to(device)
+            outputs = model(inputs)
+            _, predicted = outputs.max(1)
+
+            # Check if poisoned samples are classified as target class
+            correct_backdoor += (predicted == target_labels).sum().item()
+            total_poisoned += len(target_labels)
+
+    asr = 100. * correct_backdoor / total_poisoned if total_poisoned > 0 else 0
+
+    return asr
+
+def get_device(device_index=0):
+    if torch.cuda.is_available():
+        return torch.device(f"cuda:{device_index}")
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        return torch.device("mps")
+    else:
+        return torch.device("cpu")
+    
+def set_seed(seed):
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    random.seed(seed)
+
+@torch.no_grad()
+def evaluate(model, test_loader, device, criterion):
+    model.eval()
+    correct = total = 0
+    loss_sum = 0.0
+    for x, y in test_loader:
+        x, y = x.to(device), y.to(device)
+        out = model(x)
+        loss_sum += criterion(out, y).item() * y.size(0)
+        pred = out.argmax(1)
+        correct += (pred == y).sum().item()
+        total += y.size(0)
+    return loss_sum / total, correct / total
+
+def main(args):
+
+    device = get_device(args.device)
+
+    if args.output_path == "models/resnet18_clean.pth" and args.dataset == "poison":
+        args.output_path = "models/resnet18_poison.pth"
+
+    set_seed(args.seed)
+    g = torch.Generator()
+    g.manual_seed(args.seed)
+
+    cifar10_mean = (0.4914, 0.4822, 0.4465)
+    cifar10_std  = (0.2023, 0.1994, 0.2010)
+
+    train_pre_transform = transforms.Compose([
+        transforms.RandomCrop(32, padding=4),
+        transforms.RandomHorizontalFlip(),
+        transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
+        transforms.ToTensor(),
+    ])
+
+    test_pre_transform = transforms.ToTensor()
+
+    post_norm = transforms.Normalize(mean=cifar10_mean, std=cifar10_std)
+
+    clean_train_ds = datasets.CIFAR10("./data", train=True, download=True, transform=None)
+    clean_test_ds = datasets.CIFAR10("./data", train=False, download=True, transform=None)
+
+    train_dataset = clean_train_ds
+    test_dataset = datasets.CIFAR10("./data", train=False, download=True,
+                                    transform=transforms.Compose([test_pre_transform, post_norm]))
+    asr_loader = None
+
+    use_pin = (device.type == "cuda")
+
+    if args.dataset.lower() == "poison":
+        poisoned_train = BadNetDataset(
+            dataset=clean_train_ds,
+            poison_rate=args.train_poison_rate,
+            target_class=args.target_class,
+            trigger_size=args.trigger_size,
+            trigger_pos=args.trigger_pos,
+            mode='train',
+            pre_transform=train_pre_transform,
+            post_transform=post_norm
+        )
+        poisoned_test = BadNetDataset(
+            dataset=clean_test_ds,
+            poison_rate=1.0,
+            target_class=args.target_class,
+            trigger_size=args.trigger_size,
+            trigger_pos=args.trigger_pos,
+            mode='test_poison',
+            pre_transform=test_pre_transform,
+            post_transform=post_norm
+        )
+
+        asr_loader = DataLoader(
+            poisoned_test,
+            batch_size=args.eval_batch_size,
+            shuffle=False,
+            num_workers=2,
+            pin_memory=use_pin
+        )
+
+        train_dataset = poisoned_train
+
+    else:
+        train_dataset = datasets.CIFAR10(
+            "./data", train=True, download=True,
+            transform=transforms.Compose([train_pre_transform, post_norm])
+        )
+
+    train_loader = DataLoader(train_dataset, batch_size=args.train_batch_size, shuffle=True, num_workers=2, pin_memory=use_pin, generator=g)
+    test_loader = DataLoader(test_dataset,  batch_size=args.eval_batch_size, shuffle=False, num_workers=2, pin_memory=use_pin)
+
+    
+    model = resnet18(weights=None)
+    model.fc = nn.Linear(model.fc.in_features, 10)
+    model = model.to(device)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9)
+
+    epochs = args.epochs
+
+    print("Training with the following parameters:\n", 
+        f"Epochs = {args.epochs}\n",
+        f"Train Batch Size = {args.train_batch_size}\n",
+        f"Evaluation Batch Size = {args.eval_batch_size}\n",
+        f"Learning Rate = {args.lr}\n",
+        f"Seed = {args.seed}\n",
+        f"Output Path = {args.output_path}\n",
+        f"Device = {args.device}\n")
+    
+    best_val_acc = 0.0
+    best_model_state = None
+
+    for epoch in range(epochs):
+        model.train()
+        for x, y in train_loader:
+            x, y = x.to(device), y.to(device)
+            optimizer.zero_grad(set_to_none=True)
+            loss = criterion(model(x), y)
+            loss.backward()
+            optimizer.step()
+        val_loss, val_acc = evaluate(model, test_loader, device, criterion)
+        print(f"Epoch {epoch+1}/{epochs} - val_loss: {val_loss:.4f}  val_acc: {val_acc:.3f}")
+
+        if val_acc > best_val_acc:
+            best_val_acc = val_acc
+            best_model_state = model.state_dict()
+            print(f"New best model found at epoch {epoch+1} with val_acc: {val_acc:.3f}")
+
+        if asr_loader is not None:
+            asr = evaluate_asr(model, asr_loader, device, args.target_class)
+            print(f"ASR: {asr:.1f}%")
+
+    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
+    torch.save(best_model_state, args.output_path)
+    print(f"Best model saved to {args.output_path} with val_acc: {best_val_acc:.3f}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--epochs", help="# of epochs to iterate through", type=int, default=60)
+    parser.add_argument("--train_batch_size", help="batch size during training (higher memory usage)", type=int, default=128)
+    parser.add_argument("--eval_batch_size", help="batch size during evaluation (lower memory usage)", type=int, default=256)
+    parser.add_argument("--lr", help="learning rate for optimizer", default=0.1, type=float)
+    parser.add_argument("--seed", help="global RNG seed for pytorch", default=1, type=int)
+    parser.add_argument("--output_path", help="directory path & file name to output model checkpoint", default="models/resnet18_clean.pth", type=str)
+    parser.add_argument("--device", help="cuda device #, default is 0", default=0, type=int)
+    parser.add_argument("--dataset", choices=["clean","poison"], default="clean", help="Use clean or poison dataset")
+    parser.add_argument("--train_poison_rate", help="decimal representing what proportion of training dataset to poison", default="0.1", type=float)
+    parser.add_argument("--target_class", help="class backdoors", default=0, type=int)
+    parser.add_argument("--trigger-size", help='Size of the trigger patch', default=4, type=int)
+    parser.add_argument("--trigger-pos", help="Position of the trigger patch", default='bottom-right', choices=['bottom-right', 'bottom-left', 'top-right', 'top-left'], type=str)
+
+    args = parser.parse_args()
+    main(args)

test_report.json

@@ -0,0 +1,45 @@
+{
+  "mithridatium_version": "0.1.0",
+  "timestamp_utc": "2025-11-29T06:57:59.656900Z",
+  "model_path": "models/resnet18_poison.pth",
+  "defense": "strip",
+  "dataset": "cifar10",
+  "results": {
+    "entropies": [
+      0.8908131718635559,
+      1.0416946411132812,
+      1.25931978225708,
+      1.1651346683502197,
+      1.1246498823165894,
+      0.821864902973175,
+      1.1872310638427734,
+      0.654247522354126,
+      1.3309650421142578,
+      0.8633555173873901,
+      0.8300310969352722,
+      1.0243608951568604,
+      0.8220431208610535,
+      0.8678932785987854,
+      0.7854791879653931,
+      0.9563668966293335,
+      1.1305217742919922,
+      1.2904465198516846,
+      1.1605632305145264,
+      0.8708277940750122,
+      1.303524136543274,
+      1.0695277452468872,
+      0.8418548107147217,
+      0.7635111212730408,
+      1.0756092071533203,
+      0.7455508708953857,
+      1.1538797616958618,
+      1.1432048082351685,
+      0.8330492973327637,
+      1.124779224395752,
+      0.9224187731742859,
+      1.1702289581298828
+    ],
+    "num_bases": 32,
+    "num_perturbations": 16
+  }
+}

tests/test_dataloader_normalization.py

@@ -0,0 +1,348 @@
+"""
+Test dataloader normalization behavior in utils.py.
+
+This module tests that:
+1. Dataloader transforms properly normalize data to have means near 0
+2. CIFAR datasets load without errors and produce expected tensor shapes
+3. Normalization statistics match expected behavior
+4. Transform pipelines work correctly for each dataset
+"""
+
+import pytest
+import torch
+import numpy as np
+from mithridatium.utils import dataloader_for, get_preprocess_config
+
+
+class TestDataloaderNormalization:
+    """Test that dataloader normalization works correctly."""
+    
+    @pytest.fixture
+    def small_batch_size(self):
+        """Use small batch size for faster tests."""
+        return 32
+    
+    def test_cifar10_dataloader_creation(self, small_batch_size):
+        """Test that CIFAR-10 dataloader creates successfully."""
+        # Test both train and test splits
+        for split in ["train", "test"]:
+            dataloader, config = dataloader_for("cifar10", split, batch_size=small_batch_size)
+            
+            # Check dataloader properties
+            assert dataloader.batch_size == small_batch_size
+            assert isinstance(dataloader, torch.utils.data.DataLoader)
+            
+            # Check config
+            assert config.get_dataset() == "cifar10"
+            assert config.get_input_size() == (3, 32, 32)
+    
+    def test_cifar100_dataloader_creation(self, small_batch_size):
+        """Test that CIFAR-100 dataloader creates successfully."""
+        # Test both train and test splits
+        for split in ["train", "test"]:
+            dataloader, config = dataloader_for("cifar100", split, batch_size=small_batch_size)
+            
+            # Check dataloader properties
+            assert dataloader.batch_size == small_batch_size
+            assert isinstance(dataloader, torch.utils.data.DataLoader)
+            
+            # Check config
+            assert config.get_dataset() == "cifar100"
+            assert config.get_input_size() == (3, 32, 32)
+    
+    def test_cifar10_tensor_shapes(self, small_batch_size):
+        """Test that CIFAR-10 produces correct tensor shapes."""
+        dataloader, _ = dataloader_for("cifar10", "test", batch_size=small_batch_size)
+        
+        # Get first batch
+        batch_iter = iter(dataloader)
+        images, labels = next(batch_iter)
+        
+        # Check shapes
+        assert images.shape == (small_batch_size, 3, 32, 32), f"Expected {(small_batch_size, 3, 32, 32)}, got {images.shape}"
+        assert labels.shape == (small_batch_size,), f"Expected {(small_batch_size,)}, got {labels.shape}"
+        
+        # Check data types
+        assert images.dtype == torch.float32
+        assert labels.dtype == torch.long  # CIFAR uses long integers for class labels
+    
+    def test_cifar100_tensor_shapes(self, small_batch_size):
+        """Test that CIFAR-100 produces correct tensor shapes."""
+        dataloader, _ = dataloader_for("cifar100", "test", batch_size=small_batch_size)
+        
+        # Get first batch
+        batch_iter = iter(dataloader)
+        images, labels = next(batch_iter)
+        
+        # Check shapes
+        assert images.shape == (small_batch_size, 3, 32, 32), f"Expected {(small_batch_size, 3, 32, 32)}, got {images.shape}"
+        assert labels.shape == (small_batch_size,), f"Expected {(small_batch_size,)}, got {labels.shape}"
+        
+        # Check data types
+        assert images.dtype == torch.float32
+        assert labels.dtype == torch.long
+    
+    def test_cifar10_normalization_behavior(self, small_batch_size):
+        """Test that CIFAR-10 normalization produces data with means near 0."""
+        dataloader, config = dataloader_for("cifar10", "test", batch_size=small_batch_size)
+        
+        # Collect several batches to get good statistics
+        all_images = []
+        batch_count = 0
+        for images, _ in dataloader:
+            all_images.append(images)
+            batch_count += 1
+            if batch_count >= 10:  # Use 10 batches for statistics
+                break
+        
+        # Concatenate all images
+        all_images = torch.cat(all_images, dim=0)
+        
+        # Calculate per-channel means and stds
+        # Shape: (N, C, H, W) -> calculate over N, H, W dimensions
+        channel_means = torch.mean(all_images, dim=(0, 2, 3))  # Shape: (3,)
+        channel_stds = torch.std(all_images, dim=(0, 2, 3))    # Shape: (3,)
+        
+        # Print actual values for debugging/validation
+        print(f"CIFAR-10 normalized stats - Means: {channel_means.tolist()}, Stds: {channel_stds.tolist()}")
+        
+        # After normalization, means should be close to 0
+        # The mean centering should be very effective
+        for i, mean_val in enumerate(channel_means):
+            assert abs(mean_val.item()) < 0.1, f"Channel {i} mean {mean_val.item()} not near 0"
+        
+        # Standard deviations should be reasonably close to 1
+        # Note: Due to finite sampling and dataset characteristics, exact std=1.0 is not expected
+        # We verify the normalization is working (values roughly in expected range)
+        for i, std_val in enumerate(channel_stds):
+            assert 0.6 <= std_val.item() <= 1.4, f"Channel {i} std {std_val.item()} outside reasonable range [0.6, 1.4]"
+    
+    def test_cifar100_normalization_behavior(self, small_batch_size):
+        """Test that CIFAR-100 normalization produces data with means near 0."""
+        dataloader, config = dataloader_for("cifar100", "test", batch_size=small_batch_size)
+        
+        # Collect several batches to get good statistics
+        all_images = []
+        batch_count = 0
+        for images, _ in dataloader:
+            all_images.append(images)
+            batch_count += 1
+            if batch_count >= 10:  # Use 10 batches for statistics
+                break
+        
+        # Concatenate all images
+        all_images = torch.cat(all_images, dim=0)
+        
+        # Calculate per-channel means and stds
+        channel_means = torch.mean(all_images, dim=(0, 2, 3))
+        channel_stds = torch.std(all_images, dim=(0, 2, 3))
+        
+        # Print actual values for debugging/validation
+        print(f"CIFAR-100 normalized stats - Means: {channel_means.tolist()}, Stds: {channel_stds.tolist()}")
+        
+        # After normalization, means should be close to 0
+        for i, mean_val in enumerate(channel_means):
+            assert abs(mean_val.item()) < 0.1, f"Channel {i} mean {mean_val.item()} not near 0"
+        
+        # Standard deviations should be reasonably close to 1
+        for i, std_val in enumerate(channel_stds):
+            assert 0.6 <= std_val.item() <= 1.4, f"Channel {i} std {std_val.item()} outside reasonable range [0.6, 1.4]"
+    
+    def test_unnormalized_data_range(self, small_batch_size):
+        """Test data range before and after normalization by manually checking transforms."""
+        # This test verifies the transform pipeline is working correctly
+        from torchvision import datasets, transforms
+        
+        # Create CIFAR-10 dataset without normalization
+        unnormalized_transform = transforms.Compose([
+            transforms.ToTensor()  # Only convert to tensor, no normalization
+        ])
+        
+        unnormalized_ds = datasets.CIFAR10(
+            root="data",
+            train=False,
+            download=True,
+            transform=unnormalized_transform
+        )
+        
+        unnormalized_loader = torch.utils.data.DataLoader(
+            unnormalized_ds,
+            batch_size=small_batch_size,
+            shuffle=False
+        )
+        
+        # Get normalized dataloader
+        normalized_loader, config = dataloader_for("cifar10", "test", batch_size=small_batch_size)
+        
+        # Get first batch from each
+        unnorm_batch = next(iter(unnormalized_loader))[0]  # Just images
+        norm_batch = next(iter(normalized_loader))[0]      # Just images
+        
+        # Unnormalized data should be in [0, 1] range
+        assert unnorm_batch.min().item() >= 0.0, f"Unnormalized min {unnorm_batch.min().item()} < 0"
+        assert unnorm_batch.max().item() <= 1.0, f"Unnormalized max {unnorm_batch.max().item()} > 1"
+        
+        # Normalized data should extend beyond [0, 1] range due to normalization
+        # (some values will be negative after subtracting mean)
+        assert norm_batch.min().item() < 0.0, f"Normalized data should have negative values, min={norm_batch.min().item()}"
+        assert norm_batch.max().item() > 1.0, f"Normalized data should exceed 1, max={norm_batch.max().item()}"
+    
+    def test_different_batch_sizes(self):
+        """Test that different batch sizes work correctly."""
+        for batch_size in [1, 8, 16, 64]:
+            dataloader, _ = dataloader_for("cifar10", "test", batch_size=batch_size)
+            
+            # Get first batch
+            batch_iter = iter(dataloader)
+            images, labels = next(batch_iter)
+            
+            # Check batch size (last batch might be smaller)
+            assert images.shape[0] <= batch_size
+            assert labels.shape[0] <= batch_size
+            assert images.shape[0] == labels.shape[0]
+    
+    def test_train_vs_test_shuffle(self):
+        """Test that train loader shuffles but test loader doesn't."""
+        batch_size = 16
+        
+        # Get train and test loaders
+        train_loader, _ = dataloader_for("cifar10", "train", batch_size=batch_size)
+        test_loader, _ = dataloader_for("cifar10", "test", batch_size=batch_size)
+        
+        # For train loader, shuffle should be True (can't directly test randomness easily)
+        # But we can at least verify the loaders work
+        train_batch = next(iter(train_loader))
+        test_batch = next(iter(test_loader))
+        
+        assert train_batch[0].shape == (batch_size, 3, 32, 32)
+        assert test_batch[0].shape == (batch_size, 3, 32, 32)
+
+
+class TestDataloaderErrorHandling:
+    """Test error handling in dataloader_for function."""
+    
+    def test_invalid_dataset_error(self):
+        """Test that invalid datasets raise ValueError."""
+        with pytest.raises(ValueError) as exc_info:
+            dataloader_for("mnist", "test", batch_size=32)
+        
+        error_msg = str(exc_info.value)
+        assert "Unsupported dataset" in error_msg
+        assert "mnist" in error_msg
+    
+    def test_invalid_split_error(self):
+        """Test that invalid splits raise ValueError."""
+        with pytest.raises(ValueError) as exc_info:
+            dataloader_for("cifar10", "validation", batch_size=32)
+        
+        error_msg = str(exc_info.value)
+        assert "Invalid split" in error_msg
+        assert "validation" in error_msg
+        assert "train" in error_msg
+        assert "test" in error_msg
+    
+    def test_case_insensitive_inputs(self):
+        """Test that dataset and split names are case-insensitive."""
+        # These should all work without errors
+        for dataset in ["CIFAR10", "Cifar10", "cifar10"]:
+            for split in ["TRAIN", "Train", "train", "TEST", "Test", "test"]:
+                dataloader, config = dataloader_for(dataset, split, batch_size=8)
+                assert config.get_dataset() == "cifar10"
+
+
+class TestTransformPipelines:
+    """Test that transform pipelines are correctly structured."""
+    
+    def test_cifar_transform_efficiency(self):
+        """Test that CIFAR transforms don't include unnecessary resize operations."""
+        # This is more of a design verification test
+        # CIFAR images are already 32x32, so no resize should be needed
+        
+        dataloader, config = dataloader_for("cifar10", "test", batch_size=16)
+        
+        # Get a batch to ensure transforms work
+        batch = next(iter(dataloader))
+        images, labels = batch
+        
+        # Verify final shape is correct (transforms worked)
+        assert images.shape == (16, 3, 32, 32)
+        
+        # Verify data is normalized (not in [0,1] range)
+        assert images.min().item() < 0 or images.max().item() > 1
+    
+    def test_imagenet_transform_structure(self):
+        """Test ImageNet transforms would include proper resize operations."""
+        # Note: This test may fail if ImageNet dataset isn't available
+        # In that case, we verify the error message is helpful
+        
+        try:
+            train_loader, config = dataloader_for("imagenet", "train", batch_size=8)
+            test_loader, config = dataloader_for("imagenet", "test", batch_size=8)
+            
+            # If ImageNet is available, verify config
+            assert config.get_input_size() == (3, 224, 224)
+            
+        except ValueError as e:
+            # Should get helpful error about manual ImageNet setup
+            error_msg = str(e)
+            assert "ImageNet dataset not found" in error_msg
+            assert "data/imagenet" in error_msg
+    
+    def test_pin_memory_enabled(self):
+        """Test that dataloaders have pin_memory enabled for GPU performance."""
+        dataloader, _ = dataloader_for("cifar10", "test", batch_size=16)
+        
+        # Check that pin_memory is True (improves GPU transfer performance)
+        assert dataloader.pin_memory is True
+    
+    def test_num_workers_set(self):
+        """Test that dataloaders use multiple workers for performance."""
+        dataloader, _ = dataloader_for("cifar10", "test", batch_size=16)
+        
+        # Check that num_workers > 0 for parallel data loading
+        assert dataloader.num_workers >= 2
+
+
+class TestNormalizationMath:
+    """Test the mathematical correctness of normalization."""
+    
+    def test_normalization_formula_correctness(self):
+        """Test that normalization follows the correct formula: (x - mean) / std."""
+        # Create simple test data
+        test_tensor = torch.tensor([[[
+            [0.4914, 0.6000],  # First channel values
+            [0.3000, 0.8000]
+        ]]], dtype=torch.float32)  # Shape: (1, 1, 2, 2)
+        
+        # CIFAR-10 stats for red channel
+        mean = 0.4914
+        std = 0.2023
+        
+        # Apply normalization manually
+        normalized_manual = (test_tensor - mean) / std
+        
+        # Apply normalization using torchvision transform
+        from torchvision import transforms
+        normalize_transform = transforms.Normalize(mean=(mean,), std=(std,))
+        normalized_torch = normalize_transform(test_tensor)
+        
+        # Results should be identical (within floating point precision)
+        torch.testing.assert_close(normalized_manual, normalized_torch, rtol=1e-6, atol=1e-6)
+    
+    def test_inverse_normalization_possible(self):
+        """Test that normalization can be inverted to recover original values."""
+        dataloader, config = dataloader_for("cifar10", "test", batch_size=4)
+        
+        # Get normalized batch
+        normalized_batch = next(iter(dataloader))[0]
+        
+        # Apply inverse normalization: x_orig = (x_norm * std) + mean
+        mean = torch.tensor(config.get_mean()).view(1, 3, 1, 1)  # Shape: (1, 3, 1, 1)
+        std = torch.tensor(config.get_std()).view(1, 3, 1, 1)    # Shape: (1, 3, 1, 1)
+        
+        denormalized_batch = (normalized_batch * std) + mean
+        
+        # Denormalized values should be approximately in [0, 1] range
+        # (not exactly due to discretization and floating point precision)
+        assert denormalized_batch.min().item() >= -0.1, f"Denormalized min {denormalized_batch.min().item()} too low"
+        assert denormalized_batch.max().item() <= 1.1, f"Denormalized max {denormalized_batch.max().item()} too high"

tests/test_evaluator.py

@@ -0,0 +1,45 @@
+import os
+import torch
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+from torchvision.models import resnet18
+import mithridatium.evaluator as evaluator
+import mithridatium.loader as loader
+import unittest
+
+class TestEvaluator(unittest.TestCase):
+    def test_extract_embeddings_and_evaluate(self):
+        # Get model path from environment variable or use default
+        """
+        export MODEL_PATH=models/resnet18_bd.pth
+        export BATCH_SIZE=128
+        .venv/bin/python -m unittest tests/test_evaluator.py
+        """
+        model_path = os.environ.get("MODEL_PATH", "models/resnet18_bd.pth")
+        batch_size = int(os.environ.get("BATCH_SIZE", 128))
+
+        # Use a tiny subset of CIFAR-10
+        transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
+        ])
+        testset = datasets.CIFAR10('./data', train=False, download=True, transform=transform)
+        indices = list(range(512))
+        subset = torch.utils.data.Subset(testset, indices)
+        loader_ = DataLoader(subset, batch_size=batch_size, shuffle=False)
+
+        model, feature_module = loader.load_resnet18(model_path)
+        embs, labels = evaluator.extract_embeddings(model, loader_, feature_module)
+        print(f"Embeddings shape: {embs.shape}")
+        print(f"Labels shape: {labels.shape}")
+        print(f"First 5 labels: {labels[:5].tolist()}")
+        loss, accy = evaluator.evaluate(model, loader_)
+        print(f"Loss: {loss:.4f}")
+        print(f"Accuracy: {accy*100:.2f}%")
+        self.assertTrue(embs.shape[0] > 0)
+        self.assertTrue(labels.shape[0] > 0)
+        self.assertTrue(loss >= 0)
+        self.assertTrue(accy >= 0)
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_preprocess_config.py

@@ -0,0 +1,17 @@
+import pytest
+from mithridatium.utils import get_preprocess_config
+
+def test_get_preprocess_config():
+    # Use a known dataset for the test (e.g., cifar10)
+    dataset_name = "cifar10"
+    
+    # Load the preprocessing config for the dataset
+    config = get_preprocess_config(dataset_name)
+    
+    # Assertions based on the expected preprocessing config for CIFAR-10
+    assert config.input_size == (3, 32, 32)  # CIFAR-10 has 32x32 RGB images
+    assert config.channels_first is True      # CIFAR-10 uses NCHW format
+    assert config.value_range == (0.0, 1.0)  # Normalization range
+    assert config.mean == (0.4914, 0.4822, 0.4465)  # CIFAR-10 dataset mean
+    assert config.std == (0.2023, 0.1994, 0.2010)   # CIFAR-10 dataset standard deviation
+    assert config.ops == []  # No additional operations are needed for CIFAR-10

tests/test_strip_entropy.py

@@ -0,0 +1,44 @@
+import torch
+import sys
+import os
+
+# Add the project root to the path so we can import the module
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))
+
+from mithridatium.defenses.strip import prediction_entropy
+
+def test_prediction_entropy():
+    print("Testing prediction_entropy...")
+
+    # Case 1: Uniform distribution (Maximum entropy)
+    # Logits being equal implies uniform distribution after softmax
+    logits_uniform = torch.tensor([[1.0, 1.0, 1.0, 1.0]])
+    entropy_uniform = prediction_entropy(logits_uniform)
+    
+    # Expected entropy for uniform distribution over N classes is ln(N)
+    expected_uniform = torch.tensor([torch.log(torch.tensor(4.0))])
+    
+    print(f"Uniform Logits: {logits_uniform}")
+    print(f"Calculated Entropy: {entropy_uniform}")
+    print(f"Expected Entropy: {expected_uniform}")
+    
+    assert torch.allclose(entropy_uniform, expected_uniform, atol=1e-4), "Uniform distribution entropy mismatch"
+
+    # Case 2: One-hot distribution (Minimum entropy)
+    # One logit much larger than others
+    logits_one_hot = torch.tensor([[100.0, 0.0, 0.0, 0.0]])
+    entropy_one_hot = prediction_entropy(logits_one_hot)
+    
+    # Expected entropy is close to 0
+    expected_one_hot = torch.tensor([0.0])
+    
+    print(f"One-hot Logits: {logits_one_hot}")
+    print(f"Calculated Entropy: {entropy_one_hot}")
+    print(f"Expected Entropy: {expected_one_hot}")
+    
+    assert torch.allclose(entropy_one_hot, expected_one_hot, atol=1e-4), "One-hot distribution entropy mismatch"
+
+    print("All tests passed!")
+
+if __name__ == "__main__":
+    test_prediction_entropy()

tests/test_strip_scores.py

@@ -0,0 +1,62 @@
+import torch
+import sys
+import os
+from torch.utils.data import DataLoader, TensorDataset
+
+# Add the project root to the path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))
+
+from mithridatium.defenses.strip import strip_scores
+from mithridatium.utils import get_preprocess_config
+
+class MockModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = torch.nn.Linear(10, 4)  # 10 features, 4 classes
+
+    def forward(self, x):
+        if x.dim() > 2:
+            x = x.view(x.size(0), -1)  # Flatten if needed
+        return self.linear(x)
+
+def test_strip_scores():
+    print("Testing strip_scores...")
+    
+    # Setup
+    torch.manual_seed(42)
+    model = MockModel()
+    
+    # Get preprocessing configuration for the CIFAR-10 dataset
+    dataset_name = "cifar10"
+    config = get_preprocess_config(dataset_name)
+    
+    # Create dummy data: 100 samples, 10 features each
+    data = torch.randn(100, 10)  # Simulated input data with 10 features
+    labels = torch.randint(0, 4, (100,))  # Random labels (4 classes)
+    
+    dataset = TensorDataset(data, labels)
+    dataloader = DataLoader(dataset, batch_size=10)
+    
+    # Test execution
+    try:
+        # Run strip_scores on the mock model and dummy data
+        results = strip_scores(model, dataloader, num_bases=5, num_perturbations=10, device='cpu', configs=config)
+        
+        # Extract entropies from the results
+        entropies = results.get("entropies")
+        
+        print(f"Entropies: {entropies}")
+        
+        # Assert that entropies are in the expected format
+        assert isinstance(entropies, list), "Entropies should be a list"
+        assert len(entropies) == 5, f"Expected 5 entropies, got {len(entropies)}"
+        assert all(isinstance(e, float) for e in entropies), "All entropies should be floats"
+        
+        print("strip_scores test passed!")
+        
+    except Exception as e:
+        print(f"strip_scores test failed with error: {e}")
+        raise e
+
+if __name__ == "__main__":
+    test_strip_scores()

tests/test_utils_configs.py

@@ -0,0 +1,241 @@
+"""
+Test canonical dataset configurations in utils.py.
+
+This module tests that:
+1. DATASET_CONFIGS contains correct canonical values for supported datasets
+2. get_preprocess_config() returns proper PreprocessConfig objects
+3. Unsupported datasets raise appropriate errors
+4. Configuration values match published literature standards
+"""
+
+import pytest
+from mithridatium.utils import get_preprocess_config, DATASET_CONFIGS, PreprocessConfig
+
+
+class TestCanonicalConfigs:
+    """Test canonical dataset configuration values."""
+    
+    def test_cifar10_canonical_stats(self):
+        """Test CIFAR-10 has correct canonical normalization statistics."""
+        # CIFAR-10 canonical values from literature
+        expected_mean = (0.4914, 0.4822, 0.4465)
+        expected_std = (0.2023, 0.1994, 0.2010)
+        expected_size = (3, 32, 32)
+        
+        # Check DATASET_CONFIGS mapping
+        config_data = DATASET_CONFIGS["cifar10"]
+        assert config_data["input_size"] == expected_size
+        assert config_data["mean"] == expected_mean
+        assert config_data["std"] == expected_std
+        assert config_data["normalize"] is True
+        
+        # Check PreprocessConfig object
+        config = get_preprocess_config("cifar10")
+        assert config.get_input_size() == expected_size
+        assert config.get_mean() == expected_mean
+        assert config.get_std() == expected_std
+        assert config.get_normalize() is True
+        assert config.get_dataset() == "cifar10"
+    
+    def test_cifar100_canonical_stats(self):
+        """Test CIFAR-100 has correct canonical normalization statistics."""
+        # CIFAR-100 canonical values from literature
+        expected_mean = (0.5071, 0.4867, 0.4408)
+        expected_std = (0.2675, 0.2565, 0.2761)
+        expected_size = (3, 32, 32)
+        
+        # Check DATASET_CONFIGS mapping
+        config_data = DATASET_CONFIGS["cifar100"]
+        assert config_data["input_size"] == expected_size
+        assert config_data["mean"] == expected_mean
+        assert config_data["std"] == expected_std
+        assert config_data["normalize"] is True
+        
+        # Check PreprocessConfig object
+        config = get_preprocess_config("cifar100")
+        assert config.get_input_size() == expected_size
+        assert config.get_mean() == expected_mean
+        assert config.get_std() == expected_std
+        assert config.get_normalize() is True
+        assert config.get_dataset() == "cifar100"
+    
+    def test_imagenet_canonical_stats(self):
+        """Test ImageNet has correct canonical normalization statistics."""
+        # ImageNet canonical values from torchvision/literature
+        expected_mean = (0.485, 0.456, 0.406)
+        expected_std = (0.229, 0.224, 0.225)
+        expected_size = (3, 224, 224)
+        
+        # Check DATASET_CONFIGS mapping
+        config_data = DATASET_CONFIGS["imagenet"]
+        assert config_data["input_size"] == expected_size
+        assert config_data["mean"] == expected_mean
+        assert config_data["std"] == expected_std
+        assert config_data["normalize"] is True
+        
+        # Check PreprocessConfig object
+        config = get_preprocess_config("imagenet")
+        assert config.get_input_size() == expected_size
+        assert config.get_mean() == expected_mean
+        assert config.get_std() == expected_std
+        assert config.get_normalize() is True
+        assert config.get_dataset() == "imagenet"
+    
+    def test_case_insensitive_dataset_names(self):
+        """Test that dataset names are case-insensitive."""
+        # Test various case combinations
+        for dataset_name in ["CIFAR10", "Cifar10", "cifar10", "CiFaR10"]:
+            config = get_preprocess_config(dataset_name)
+            assert config.get_dataset() == "cifar10"
+        
+        for dataset_name in ["CIFAR100", "Cifar100", "cifar100", "CiFaR100"]:
+            config = get_preprocess_config(dataset_name)
+            assert config.get_dataset() == "cifar100"
+        
+        for dataset_name in ["IMAGENET", "ImageNet", "imagenet", "ImAgEnEt"]:
+            config = get_preprocess_config(dataset_name)
+            assert config.get_dataset() == "imagenet"
+    
+    def test_whitespace_handling(self):
+        """Test that dataset names handle whitespace correctly."""
+        # Test with leading/trailing whitespace
+        config = get_preprocess_config("  cifar10  ")
+        assert config.get_dataset() == "cifar10"
+        
+        config = get_preprocess_config("\tcifar100\n")
+        assert config.get_dataset() == "cifar100"
+    
+    def test_unsupported_dataset_error(self):
+        """Test that unsupported datasets raise ValueError with helpful message."""
+        with pytest.raises(ValueError) as exc_info:
+            get_preprocess_config("mnist")
+        
+        error_msg = str(exc_info.value)
+        assert "mnist" in error_msg
+        assert "Unsupported dataset" in error_msg
+        assert "cifar10" in error_msg  # Should list supported datasets
+        assert "cifar100" in error_msg
+        assert "imagenet" in error_msg
+    
+    def test_preprocess_config_default_values(self):
+        """Test that PreprocessConfig has correct default values."""
+        for dataset in ["cifar10", "cifar100", "imagenet"]:
+            config = get_preprocess_config(dataset)
+            
+            # Common defaults across all datasets
+            assert config.get_channels_first() is True
+            assert config.get_value_range() == (0.0, 1.0)
+            assert config.get_normalize() is True
+            assert config.get_ops() == []
+    
+    def test_all_supported_datasets_in_mapping(self):
+        """Test that all datasets mentioned in error messages are in DATASET_CONFIGS."""
+        try:
+            get_preprocess_config("invalid_dataset")
+        except ValueError as e:
+            error_msg = str(e)
+            # Extract supported datasets from error message
+            # Message format: "Supported datasets: cifar10, cifar100, imagenet"
+            if "Supported datasets:" in error_msg:
+                supported_part = error_msg.split("Supported datasets:")[1].strip()
+                mentioned_datasets = [ds.strip() for ds in supported_part.split(",")]
+                
+                # Verify all mentioned datasets exist in DATASET_CONFIGS
+                for dataset in mentioned_datasets:
+                    assert dataset in DATASET_CONFIGS, f"Dataset {dataset} mentioned in error but not in DATASET_CONFIGS"
+
+
+class TestDatasetConfigsCompleteness:
+    """Test that DATASET_CONFIGS mapping is complete and well-formed."""
+    
+    def test_dataset_configs_structure(self):
+        """Test that DATASET_CONFIGS has proper structure."""
+        required_keys = {"input_size", "mean", "std", "normalize"}
+        
+        for dataset_name, config in DATASET_CONFIGS.items():
+            # Check all required keys present
+            assert required_keys.issubset(config.keys()), f"Missing keys in {dataset_name} config"
+            
+            # Check types and shapes
+            assert isinstance(config["input_size"], tuple)
+            assert len(config["input_size"]) == 3  # (C, H, W)
+            assert all(isinstance(x, int) and x > 0 for x in config["input_size"])
+            
+            assert isinstance(config["mean"], tuple)
+            assert len(config["mean"]) == 3  # (R, G, B)
+            assert all(isinstance(x, float) and 0 <= x <= 1 for x in config["mean"])
+            
+            assert isinstance(config["std"], tuple)
+            assert len(config["std"]) == 3  # (R, G, B)
+            assert all(isinstance(x, float) and x > 0 for x in config["std"])
+            
+            assert isinstance(config["normalize"], bool)
+    
+    def test_cifar_datasets_have_32x32_size(self):
+        """Test that CIFAR datasets have correct 32x32 input size."""
+        for dataset in ["cifar10", "cifar100"]:
+            config = DATASET_CONFIGS[dataset]
+            assert config["input_size"] == (3, 32, 32), f"{dataset} should be 3x32x32"
+    
+    def test_imagenet_has_224x224_size(self):
+        """Test that ImageNet has correct 224x224 input size."""
+        config = DATASET_CONFIGS["imagenet"]
+        assert config["input_size"] == (3, 224, 224), "ImageNet should be 3x224x224"
+    
+    def test_normalization_stats_reasonable_ranges(self):
+        """Test that mean/std values are in reasonable ranges for image data."""
+        for dataset_name, config in DATASET_CONFIGS.items():
+            # Mean values should be between 0 and 1 for normalized images
+            for channel_mean in config["mean"]:
+                assert 0.0 <= channel_mean <= 1.0, f"{dataset_name} mean {channel_mean} out of range [0,1]"
+            
+            # Std values should be positive and reasonable (typically 0.1-0.5 for image data)
+            for channel_std in config["std"]:
+                assert 0.05 <= channel_std <= 0.5, f"{dataset_name} std {channel_std} out of reasonable range [0.05,0.5]"
+
+
+class TestPreprocessConfigMethods:
+    """Test PreprocessConfig class methods and functionality."""
+    
+    def test_preprocess_config_getters(self):
+        """Test all getter methods work correctly."""
+        config = get_preprocess_config("cifar10")
+        
+        # Test all getter methods
+        assert config.get_input_size() == (3, 32, 32)
+        assert config.get_channels_first() is True
+        assert config.get_value_range() == (0.0, 1.0)
+        assert config.get_mean() == (0.4914, 0.4822, 0.4465)
+        assert config.get_std() == (0.2023, 0.1994, 0.2010)
+        assert config.get_normalize() is True
+        assert config.get_ops() == []
+        assert config.get_dataset() == "cifar10"
+    
+    def test_preprocess_config_setters(self):
+        """Test setter methods work correctly."""
+        config = get_preprocess_config("cifar10")
+        
+        # Test setters
+        config.set_input_size((3, 64, 64))
+        assert config.get_input_size() == (3, 64, 64)
+        
+        config.set_channels_first(False)
+        assert config.get_channels_first() is False
+        
+        config.set_value_range((-1.0, 1.0))
+        assert config.get_value_range() == (-1.0, 1.0)
+        
+        config.set_mean((0.5, 0.5, 0.5))
+        assert config.get_mean() == (0.5, 0.5, 0.5)
+        
+        config.set_std((0.25, 0.25, 0.25))
+        assert config.get_std() == (0.25, 0.25, 0.25)
+        
+        config.set_normalize(False)
+        assert config.get_normalize() is False
+        
+        config.set_ops(["resize:64", "crop:32"])
+        assert config.get_ops() == ["resize:64", "crop:32"]
+        
+        config.set_dataset("custom")
+        assert config.get_dataset() == "custom"

tests/tests_report.py

@@ -0,0 +1,159 @@
+# tests/test_cli_v2.py
+import json
+from pathlib import Path
+from typer.testing import CliRunner
+
+from mithridatium.cli import (
+    app,
+    VERSION,
+    EXIT_NO_INPUT,
+    EXIT_IO_ERROR,
+    EXIT_USAGE_ERROR,
+    EXIT_CANT_CREATE,
+)
+
+from mithridatium import report as rpt
+
+runner = CliRunner()
+
+
+def _write_model(tmp_path: Path) -> Path:
+    """Create a tiny dummy model file that is readable."""
+    model = tmp_path / "fake.pth"
+    model.write_bytes(b"ok")
+    return model
+
+
+def test_version_flag():
+    res = runner.invoke(app, ["--version"])
+    assert res.exit_code == 0
+    assert VERSION.strip() in res.stdout
+
+
+def test_defenses_lists_spectral_and_mmbd():
+    res = runner.invoke(app, ["defenses"])
+    assert res.exit_code == 0
+    # order not guaranteed; check both are present
+    assert "spectral" in res.stdout
+    assert "mmbd" in res.stdout
+
+def test_detect_spectral_stdout(tmp_path):
+    model = (tmp_path / "fake.pth"); model.write_bytes(b"ok")
+    res = runner.invoke(app, ["detect", "-m", str(model), "-D", "spectral", "-d", "cifar10", "-o", "-"])
+    assert res.exit_code == 0
+    assert '"results"' in res.stdout
+    assert '"top_eigenvalue"' in res.stdout
+    assert "defense=spectral" in res.stdout or '"defense": "spectral"' in res.stdout
+
+def test_detect_stdout_json_then_summary(tmp_path):
+    model = _write_model(tmp_path)
+    res = runner.invoke(
+        app,
+        ["detect", "-m", str(model), "-D", "mmbd", "-d", "cifar10", "-o", "-"],
+    )
+    assert res.exit_code == 0
+    # JSON bits
+    assert '"mithridatium_version"' in res.stdout
+    assert '"defense": "mmbd"' in res.stdout
+    assert '"dataset": "cifar10"' in res.stdout
+    assert '"results"' in res.stdout
+    assert '"suspected_backdoor"' in res.stdout
+    # summary bits
+    assert "defense=mmbd" in res.stdout
+    assert "dataset=cifar10" in res.stdout
+
+
+def test_detect_to_file_json_schema(tmp_path):
+    model = _write_model(tmp_path)
+    out = tmp_path / "report.json"
+    res = runner.invoke(
+        app,
+        ["detect", "-m", str(model), "-D", "mmbd", "-d", "cifar10", "-o", str(out)],
+    )
+    assert res.exit_code == 0
+    assert out.exists()
+    rep = json.loads(out.read_text(encoding="utf-8"))
+    # top-level keys
+    for k in ("mithridatium_version", "model_path", "defense", "dataset", "results"):
+        assert k in rep
+    assert rep["defense"] == "mmbd"
+    assert rep["dataset"] == "cifar10"
+    # results keys + types
+    r = rep["results"]
+    assert isinstance(r["suspected_backdoor"], bool)
+    assert isinstance(r["num_flagged"], int)
+    assert isinstance(r["top_eigenvalue"], (int, float))
+
+
+def test_missing_model_errors_with_code(tmp_path):
+    missing = tmp_path / "nope.pth"
+    out = tmp_path / "r.json"
+    res = runner.invoke(
+        app, ["detect", "-m", str(missing), "-D", "mmbd", "-o", str(out)]
+    )
+    assert res.exit_code == EXIT_NO_INPUT
+    assert "model path not found" in res.stdout
+
+
+def test_unreadable_model_errors_with_code(tmp_path, monkeypatch):
+    model = _write_model(tmp_path)
+
+    # Patch Path.open to raise OSError when opening this file in 'rb'
+    from pathlib import Path as _P
+    _orig_open = _P.open
+
+    def bad_open(self, mode="r", *args, **kwargs):
+        if self == model and "rb" in mode:
+            raise OSError("permission denied")
+        return _orig_open(self, mode, *args, **kwargs)
+
+    monkeypatch.setattr(_P, "open", bad_open)
+
+    res = runner.invoke(
+        app, ["detect", "-m", str(model), "-D", "mmbd", "-o", str(tmp_path / "r.json")]
+    )
+    assert res.exit_code == EXIT_IO_ERROR
+    assert "could not be opened" in res.stdout
+    assert "permission denied" in res.stdout
+
+
+def test_unsupported_defense(tmp_path):
+    model = _write_model(tmp_path)
+    res = runner.invoke(
+        app, ["detect", "-m", str(model), "-D", "not_a_defense", "-o", str(tmp_path / "r.json")]
+    )
+    assert res.exit_code == EXIT_USAGE_ERROR
+    assert "unsupported --defense" in res.stdout
+    # should list supported defenses
+    assert "spectral" in res.stdout and "mmbd" in res.stdout
+
+
+def test_force_overwrite(tmp_path):
+    model = _write_model(tmp_path)
+    out = tmp_path / "r.json"
+
+    # First write
+    res1 = runner.invoke(app, ["detect", "-m", str(model), "-D", "mmbd", "-o", str(out)])
+    assert res1.exit_code == 0 and out.exists()
+
+    # Overwrite should fail without --force
+    res2 = runner.invoke(app, ["detect", "-m", str(model), "-D", "mmbd", "-o", str(out)])
+    assert res2.exit_code == EXIT_CANT_CREATE
+    assert "already exists" in res2.stdout
+
+    # Overwrite with --force should succeed
+    res3 = runner.invoke(
+        app, ["detect", "-m", str(model), "-D", "mmbd", "-o", str(out), "--force"]
+    )
+    assert res3.exit_code == 0
+
+
+def test_build_report_schema_helper():
+    res = {"suspected_backdoor": True, "num_flagged": 500, "top_eigenvalue": 42.3}
+    rep = rpt.build_report("models/resnet18_bd.pth", "mmbd", "cifar10", "0.1.1", res)
+    for k in ("mithridatium_version", "model_path", "defense", "dataset", "results"):
+        assert k in rep
+    r = rep["results"]
+    assert isinstance(r["suspected_backdoor"], bool)
+    assert isinstance(r["num_flagged"], int)
+    assert isinstance(r["top_eigenvalue"], (int, float))