Inference.py

# Project/infer_one.py
from __future__ import annotations
import os
import sys
from typing import Dict, Tuple

import torch
from PIL import Image
from torch.cuda.amp import autocast

from CNN_Architecture import alexNet
from Checkpoint_utils import load_checkpoint


def invert_class_map(class_to_idx: Dict[str, int]) -> Dict[int, str]:
    return {v: k for k, v in class_to_idx.items()}


def load_transform(transform_path: str):
    # You saved transforms via torch.save in TransformDefin.py
    # This will restore the torchvision transform pipeline.
    return torch.load(transform_path, map_location="cpu")


@torch.inference_mode()
def predict_one(
    image_path: str,
    checkpoint_path: str,
    transform_path: str,
    device: str = "cuda",
) -> Tuple[int, float, str]:
    device_t = torch.device(device if (device == "cuda" and torch.cuda.is_available()) else "cpu")
    is_cuda = (device_t.type == "cuda")

    # Load checkpoint on CPU first; then move model
    dummy_model = alexNet()  # NOTE: if you changed alexNet(num_classes), see note below.
    ckpt = torch.load(checkpoint_path, map_location="cpu")

    # class mapping
    class_to_idx = ckpt.get("class_to_idx", None)
    idx_to_class = invert_class_map(class_to_idx) if class_to_idx else None

    # Load weights
    dummy_model.load_state_dict(ckpt["model_state_dict"], strict=True)
    model = dummy_model.to(device_t).eval()

    transform = load_transform(transform_path)

    img = Image.open(image_path).convert("RGB")
    x = transform(img)  # could be (3,H,W) or (10,3,H,W)

    if x.dim() == 4:
        # TenCrop: (10,3,224,224) -> (1,10,3,224,224)
        x = x.unsqueeze(0)
    else:
        # (3,224,224) -> (1,3,224,224)
        x = x.unsqueeze(0)

    x = x.to(device_t, non_blocking=True)

    with autocast("cuda", dtype=torch.float16): # type: ignore
        if x.dim() == 5:
            # (B,10,3,224,224)
            B, NC, C, H, W = x.shape
            flat = x.view(B * NC, C, H, W)
            logits = model(flat)
            probs = torch.softmax(logits, dim=1).view(B, NC, -1).mean(dim=1)  # (B,K)
        else:
            logits = model(x)
            probs = torch.softmax(logits, dim=1)

    conf, pred_idx = probs[0].max(dim=0)
    pred_idx_i = int(pred_idx.item())
    conf_f = float(conf.item())

    pred_name = idx_to_class[pred_idx_i] if idx_to_class else str(pred_idx_i)
    return pred_idx_i, conf_f, pred_name


if __name__ == "__main__":
    # Edit these paths
    IMAGE_PATH = "path/to/one_image.jpg"
    CHECKPOINT = "checkpoints/best.pth"          # or latest.pth / epoch_XXX.pth
    TRANSFORM = "val_transform.pth"              # your saved validation transform
    DEVICE = "cuda"

    pred_idx, conf, pred_name = predict_one(IMAGE_PATH, CHECKPOINT, TRANSFORM, DEVICE)
    print(f"Pred: {pred_name} (idx={pred_idx})  conf={conf:.4f}")