test_vis_single.py

import math
from pathlib import Path
from typing import Tuple, Union, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as T
from PIL import Image
import matplotlib.pyplot as plt

from models.OC import ObjectCounter


# ========= USER INPUTS (edit these two lines) =========
INPUT_IMAGE  = Path("/root/CBD/images/PUCPR_vis1.png")
OUTPUT_IMAGE = Path("/root/CBD/outputs/PUCPR_vis1.png")
# ======================================================

# Model + normalization config
MODEL_PATH  = Path("/root/CBD/checkpoints/CA2PU_parameters.pth")
DATASET_NAME = "PUCPR" 
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],
    ),
}

# Behavior toggles (kept simple; no need to edit)
UPSAMPLE_TO_IMAGE = True     # Save density map at the input image resolution
CMAP_NAME = "jet"            # Requested colormap

# Device
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]),
])

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 load_model(model_path: Path) -> ObjectCounter:
    """Load ObjectCounter and weights."""
    net = ObjectCounter([0], "HRNet")
    try:
        state = torch.load(model_path, map_location="cpu", weights_only=True)
    except TypeError:
        state = torch.load(model_path, map_location="cpu")
    net.load_state_dict(state)
    net.to(DEVICE)
    net.eval()
    return net

def _to_numpy_hw(x: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
    """Convert to (H,W) float32 numpy."""
    if isinstance(x, torch.Tensor):
        x = x.detach().float().cpu().numpy()
    x = x.astype(np.float32, copy=False)
    while x.ndim > 2:
        x = x[0]
    return x

def _resize_bilinear_hw(arr_hw: np.ndarray, target_hw: Tuple[int, int]) -> np.ndarray:
    h, w = target_hw
    if arr_hw.shape == (h, w):
        return arr_hw
    ten = torch.from_numpy(arr_hw)[None, None, :, :]
    out = F.interpolate(ten, size=(h, w), mode="bilinear", align_corners=False)
    return out[0, 0].cpu().numpy()

def _normalize_for_vis(dmap: np.ndarray) -> np.ndarray:
    """Map to [0,1] using a robust vmax to avoid outliers."""
    if np.any(dmap > 0):
        vmax = float(np.percentile(dmap, 99.5))
        if vmax <= 1e-12:
            vmax = float(dmap.max() + 1e-6)
    else:
        vmax = float(dmap.max() + 1e-6)
    vmax = 1.0 if vmax <= 1e-12 else vmax
    d_clipped = np.clip(dmap, 0.0, vmax)
    return d_clipped / (vmax + 1e-12)

def predict_and_save(input_path: Path, output_path: Path) -> float:
    if not input_path.exists():
        raise FileNotFoundError(f"Input image not found: {input_path}")

    # Load model
    model = load_model(MODEL_PATH)

    # Load image
    img = Image.open(input_path)
    if img.mode == "L":
        img = img.convert("RGB")

    img_tensor = img_transform(img).unsqueeze(0).to(DEVICE)

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

    # Count
    pred_count = float(count_from_fidt(pred_map))

    # Prepare density map for saving
    pred_np = _to_numpy_hw(pred_map)
    if UPSAMPLE_TO_IMAGE:
        H, W = img.size[1], img.size[0]  # PIL (W,H)
        pred_np = _resize_bilinear_hw(pred_np, (H, W))

    pred_vis = _normalize_for_vis(pred_np)

    # Save as a colorized heatmap (no axes, just the map)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    plt.imsave(str(output_path), pred_vis, cmap=CMAP_NAME)

    return pred_count

def main():
    count = predict_and_save(INPUT_IMAGE, OUTPUT_IMAGE)
    print(f"Predicted count: {count:.2f}")
    print(f"Saved density map to: {OUTPUT_IMAGE}")

if __name__ == "__main__":
    main()