Testing.py

# Project/test_to_csv.py
from __future__ import annotations
import os
import csv
import glob
from typing import List, Optional, Tuple

import torch
from torch.utils.data import DataLoader, Dataset
from PIL import Image

from Inference import predict_one


IMG_EXTS = (".jpg", ".jpeg", ".png", ".bmp", ".webp")


class FlatImageFolder(Dataset):
    """For a folder that contains images directly (no class subfolders)."""
    def __init__(self, root: str):
        self.root = root
        files: List[str] = []
        for ext in IMG_EXTS:
            files.extend(glob.glob(os.path.join(root, f"**/*{ext}"), recursive=True))
        self.files = sorted(files)

    def __len__(self) -> int:
        return len(self.files)

    def __getitem__(self, idx: int) -> str:
        return self.files[idx]


def is_imagefolder_structure(root: str) -> bool:
    # Heuristic: has at least one subdir, and that subdir has images
    if not os.path.isdir(root):
        return False
    subdirs = [d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]
    if not subdirs:
        return False
    for d in subdirs:
        p = os.path.join(root, d)
        for ext in IMG_EXTS:
            if glob.glob(os.path.join(p, f"*{ext}")):
                return True
    return False


def main():
    # ====== EDIT THESE ======
    TEST_ROOT = "path/to/test"             # folder of test images
    CHECKPOINT = "checkpoints/best.pth"
    TRANSFORM = "val_transform.pth"
    OUT_CSV = "test_predictions.csv"
    DEVICE = "cuda"
    # =======================

    rows = []

    # If test set is ImageFolder-like, we can record true labels too.
    if is_imagefolder_structure(TEST_ROOT):
        from torchvision.datasets import ImageFolder

        ds = ImageFolder(TEST_ROOT)
        # ds.samples: List[(path, class_idx)]
        for path, true_idx in ds.samples:
            pred_idx, conf, pred_name = predict_one(path, CHECKPOINT, TRANSFORM, DEVICE)
            true_name = ds.classes[true_idx]
            rows.append({
                "filepath": path,
                "true_idx": true_idx,
                "true_name": true_name,
                "pred_idx": pred_idx,
                "pred_name": pred_name,
                "confidence": conf,
            })
    else:
        # Flat folder: no labels
        ds = FlatImageFolder(TEST_ROOT)
        for path in ds:
            pred_idx, conf, pred_name = predict_one(path, CHECKPOINT, TRANSFORM, DEVICE)
            rows.append({
                "filepath": path,
                "true_idx": "",
                "true_name": "",
                "pred_idx": pred_idx,
                "pred_name": pred_name,
                "confidence": conf,
            })

    # Write CSV
    os.makedirs(os.path.dirname(OUT_CSV) or ".", exist_ok=True)
    fieldnames = ["filepath", "true_idx", "true_name", "pred_idx", "pred_name", "confidence"]
    with open(OUT_CSV, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        w.writeheader()
        w.writerows(rows)

    print(f"Wrote {len(rows)} predictions to: {OUT_CSV}")


if __name__ == "__main__":
    main()