csrnet.py

"""
CSRNet crowd counting model (CVPR 2018).
VGG-16 frontend + dilated conv backend.
Weights from rootstrap-org/crowd-counting (HuggingFace).
"""
import os
import time
import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


# ── Model Architecture ──────────────────────────────────────

def _make_vgg_layers():
    cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512]
    layers = []
    in_channels = 3
    for v in cfg:
        if v == 'M':
            layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
        else:
            layers.append(nn.Conv2d(in_channels, v, kernel_size=3, padding=1))
            layers.append(nn.ReLU(inplace=True))
            in_channels = v
    return nn.Sequential(*layers)


class CSRNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.frontend = _make_vgg_layers()
        self.backend = nn.Sequential(
            nn.Conv2d(512, 512, 3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 256, 3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 128, 3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 64, 3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
        )
        self.output_layer = nn.Conv2d(64, 1, 1)

    def forward(self, x):
        x = self.frontend(x)
        x = self.backend(x)
        x = self.output_layer(x)
        return x


# ── Global model singleton ──────────────────────────────────

MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
WEIGHTS_PATH = os.path.join(MODEL_DIR, "csrnet_weights.pth")

_csrnet_model = None
_csrnet_available = False
_inited = False


def _init_csrnet():
    global _inited
    if _inited:
        return
    _inited = True
    global _csrnet_model, _csrnet_available
    if not os.path.exists(WEIGHTS_PATH):
        print("  [CSRNet] weights not found, run download first")
        return

    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        _csrnet_model = CSRNet().to(device)
        ckpt = torch.load(WEIGHTS_PATH, map_location=device, weights_only=True)
        _csrnet_model.load_state_dict(ckpt)
        _csrnet_model.eval()
        _csrnet_available = True
        print(f"  [CSRNet] loaded on {device} ({sum(p.numel() for p in _csrnet_model.parameters())/1e6:.1f}M params)")
    except Exception as e:
        print(f"  [CSRNet] init failed: {e}")
        _csrnet_available = False


def get_csrnet():
    global _csrnet_model
    if _csrnet_model is None:
        _init_csrnet()
    return _csrnet_model if _csrnet_available else None


def is_csrnet_available() -> bool:
    return _csrnet_available


# ── Inference helpers ───────────────────────────────────────

TARGET_MAX_DIM = 1024


def _resize_dims(h: int, w: int, max_dim: int = TARGET_MAX_DIM):
    scale = min(max_dim / max(h, w), 1.0)
    new_h = int(h * scale)
    new_w = int(w * scale)
    new_h = (new_h // 8) * 8   # CSRNet downsamples by 8x
    new_w = (new_w // 8) * 8
    return new_h, new_w


def count_csrnet(image: np.ndarray) -> dict:
    """Run CSRNet inference on BGR image. Returns count + density + heatmap."""
    t0 = time.time()
    original_h, original_w = image.shape[:2]

    model = get_csrnet()
    if model is None:
        return None

    device = next(model.parameters()).device

    # Preprocess: BGR → RGB, resize, normalize
    img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    new_h, new_w = _resize_dims(original_h, original_w)
    img_resized = cv2.resize(img_rgb, (new_w, new_h))

    # Convert to tensor: [0,255] → [0,1]
    tensor = torch.from_numpy(img_resized).float().permute(2, 0, 1) / 255.0
    tensor = tensor.unsqueeze(0).to(device)  # [1, 3, H, W]

    with torch.no_grad():
        density = model(tensor)
    density = density.squeeze().cpu().numpy()  # [H/8, W/8]

    count = float(density.sum())
    count_rounded = int(round(count))

    # Resize density to original size
    density_resized = cv2.resize(density, (original_w, original_h))

    # Clip negative values, create heatmap
    d_clipped = np.maximum(density_resized, 0)
    d_min = d_clipped.min()
    d_max = d_clipped.max()
    if d_max > d_min:
        d_norm = ((d_clipped - d_min) / (d_max - d_min) * 255).astype(np.uint8)
    else:
        d_norm = np.zeros_like(d_clipped, dtype=np.uint8)

    heatmap = cv2.applyColorMap(d_norm, cv2.COLORMAP_JET)
    overlay = cv2.addWeighted(image, 0.6, heatmap, 0.4, 0)

    _, heatmap_buf = cv2.imencode(".jpg", overlay, [cv2.IMWRITE_JPEG_QUALITY, 85])

    elapsed = (time.time() - t0) * 1000

    return {
        "count": count,
        "count_rounded": count_rounded,
        "density_map": density_resized,
        "heatmap_bytes": heatmap_buf.tobytes(),
        "processing_time_ms": round(elapsed, 1),
    }