test.py

import math
import os
from pathlib import Path
from typing import List

import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as T
from PIL import Image

from models.OC import ObjectCounter
# -----------------------
# Configuration
# -----------------------
DATA_ROOT = Path("/root/autodl-tmp/CARPK/test")
MODEL_PATH = Path("/root/CBD/checkpoints/PU2CA_parameters.pth")

MEAN_STD = {
    "PUCPR": (
        [0.52323937416, 0.52659797668, 0.48122045398],
        [0.21484816074, 0.20566709340, 0.22544583678],
    ),
    "CARPK": (
        [0.46704635024, 0.49598187208, 0.47164431214],
        [0.24702641368, 0.23411691189, 0.24729225040],
    ),
}
DATASET_NAME = "CARPK"
IMG_DIR = DATA_ROOT / "img"
DEN_DIR = DATA_ROOT / "den"

# CUDA settings
torch.backends.cudnn.benchmark = True
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Transforms
img_transform = T.Compose([
    T.ToTensor(),
    T.Normalize(*MEAN_STD[DATASET_NAME]),
])
to_tensor = T.ToTensor()


# -----------------------
# Helpers
# -----------------------
class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.cur_val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, cur_val):
        self.cur_val = cur_val
        self.sum += cur_val
        self.count += 1
        self.avg = self.sum / self.count

def count_from_fidt(input):
    input = torch.Tensor(input)
    input = input.reshape(-1,1,input.shape[-2],input.shape[-1])
    input_max = torch.max(input).item()
    keep = nn.functional.max_pool2d(input, (3, 3), stride=1, padding=1)
    keep = (keep == input).float()
    input = keep * input

    '''set the pixel valur of local maxima as 1 for counting'''
    input[input < 100.0 / 255.0 * input_max] = 0
    input[input > 0] = 1

    ''' negative sample'''
    if input_max < 0.1:
        input = input * 0

    count = torch.sum(input).item()
    return count

def list_image_files(img_dir: Path) -> List[str]:
    """
    Return a sorted list of file names (not paths) from img_dir.
    Keeps original logic: iterate files found directly under img_dir.
    """
    if not img_dir.exists():
        raise FileNotFoundError(f"Image directory not found: {img_dir}")

    files = [f.name for f in sorted(img_dir.iterdir()) if f.is_file() and not f.name.startswith(".")]
    return files


def load_model(model_path: Path) -> ObjectCounter:
    """
    Load ObjectCounter model and put it in eval mode on DEVICE.
    """
    net = ObjectCounter([0], "HRNet")
    state = torch.load(model_path, map_location="cpu", weights_only=True)
    net.load_state_dict(state)
    net.to(DEVICE)
    net.eval()
    return net


def evaluate(file_list: List[str], model: ObjectCounter) -> None:
    """
    Evaluate MAE and RMSE across the provided image list.
    """
    mae_meter = AverageMeter()
    mse_meter = AverageMeter()

    for filename in file_list:
        img_path = IMG_DIR / filename
        stem = img_path.stem
        den_path = DEN_DIR / f"{stem}.npy"

        # Load density map
        den = np.load(den_path)
        gt_count = count_from_fidt(den)

        # Load & preprocess image
        img = Image.open(img_path)
        if img.mode == "L":
            img = img.convert("RGB")
        img_tensor = img_transform(img).unsqueeze(0).to(DEVICE)

        # Forward pass
        with torch.no_grad():
            pred_map = model.test_forward(img_tensor)

        # Get predicted count
        pred_count = count_from_fidt(pred_map)

        # Update metrics
        err = gt_count - pred_count
        mae_meter.update(abs(err))
        mse_meter.update(err * err)

    print(f"MAE:  {mae_meter.avg:.1f}")
    print(f"RMSE: {math.sqrt(mse_meter.avg):.1f}")


# -----------------------
# Main
# -----------------------
def main() -> None:
    file_list = list_image_files(IMG_DIR)
    if not file_list:
        raise RuntimeError(f"No files found in {IMG_DIR}")

    model = load_model(MODEL_PATH)
    evaluate(file_list, model)


if __name__ == "__main__":
    main()