test.py

import argparse
import json
import os
from pathlib import Path

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import torch
import torchvision.transforms as standard_transforms

import util.misc as utils
from models import build_model

PET_TRANSFORM = standard_transforms.Compose([
    standard_transforms.ToTensor(),
    standard_transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])


def get_args_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser('PET single-image inference (HF release)', add_help=False)

    parser.add_argument('--image_path', required=True, type=str,
                        help='Path to a single input image.')
    parser.add_argument('--resume', default='PET_Finetuned.safetensors', type=str,
                        help='Path to model weights (.safetensors or .pth).')
    parser.add_argument('--device', default='cuda', type=str,
                        help='Device for inference, e.g. cuda or cpu.')

    parser.add_argument('--backbone', default='vgg16_bn', type=str)
    parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned', 'fourier'))
    parser.add_argument('--dec_layers', default=2, type=int)
    parser.add_argument('--dim_feedforward', default=512, type=int)
    parser.add_argument('--hidden_dim', default=256, type=int)
    parser.add_argument('--dropout', default=0.0, type=float)
    parser.add_argument('--nheads', default=8, type=int)

    parser.add_argument('--set_cost_class', default=1, type=float)
    parser.add_argument('--set_cost_point', default=0.05, type=float)
    parser.add_argument('--ce_loss_coef', default=1.0, type=float)
    parser.add_argument('--point_loss_coef', default=5.0, type=float)
    parser.add_argument('--eos_coef', default=0.5, type=float)

    parser.add_argument('--dataset_file', default='SHA')
    parser.add_argument('--data_path', default='./data/ShanghaiTech/PartA', type=str)

    parser.add_argument('--upper_bound', default=-1, type=int,
                        help='Max image side for inference; -1 means only cap at 2560 (same as compare_models).')
    parser.add_argument('--output_image', default='', type=str,
                        help='Optional path to save annotated image panel.')
    parser.add_argument('--title_text', default='PET-Finetuned', type=str,
                        help='Title prefix used in top panel text.')
    parser.add_argument('--radius', default=3, type=int)
    parser.add_argument('--point_color', default='0,255,0', type=str,
                        help='BGR color for points, e.g., 0,255,0')
    parser.add_argument('--panel_long_side', default=1600, type=int,
                        help='Resize annotated panel long side to this value.')
    parser.add_argument('--panel_pad', default=24, type=int,
                        help='Panel padding around the image and title area.')
    parser.add_argument('--panel_font_size', default=48, type=int,
                        help='Font size for panel title text.')

    parser.add_argument('--output_json', default='', type=str,
                        help='Optional output JSON path for prediction details.')
    parser.add_argument('--seed', default=42, type=int)

    return parser


def parse_color(color_str: str):
    parts = color_str.split(',')
    if len(parts) != 3:
        raise ValueError('color must be B,G,R like 0,255,0')
    return tuple(int(p.strip()) for p in parts)


def resolve_device(device_str: str) -> torch.device:
    if device_str.startswith('cuda') and not torch.cuda.is_available():
        print('CUDA not available. Falling back to CPU.')
        return torch.device('cpu')
    device = torch.device(device_str)
    if device.type == 'cuda' and device.index is not None:
        torch.cuda.set_device(device.index)
    return device


def resize_for_eval(frame_rgb, upper_bound):
    h, w = frame_rgb.shape[:2]
    max_size = max(h, w)
    if upper_bound != -1 and max_size > upper_bound:
        scale = float(upper_bound) / float(max_size)
    elif max_size > 2560:
        scale = 2560.0 / float(max_size)
    else:
        scale = 1.0
    if scale == 1.0:
        return frame_rgb, scale
    new_w = max(1, int(round(w * scale)))
    new_h = max(1, int(round(h * scale)))
    resized = cv2.resize(frame_rgb, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
    return resized, scale


def load_font(font_size=40, bold=False, font_paths=None):
    if font_paths is None:
        if bold:
            font_paths = [
                '/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf',
                '/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf',
                '/usr/share/fonts/truetype/freefont/FreeSansBold.ttf',
            ]
        else:
            font_paths = [
                '/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',
                '/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf',
                '/usr/share/fonts/truetype/freefont/FreeSans.ttf',
            ]
    for font_path in font_paths:
        if os.path.exists(font_path):
            try:
                return ImageFont.truetype(font_path, font_size)
            except OSError:
                continue
    try:
        fallback = 'DejaVuSans-Bold.ttf' if bold else 'DejaVuSans.ttf'
        return ImageFont.truetype(fallback, font_size)
    except OSError:
        return ImageFont.load_default()


def draw_text(draw, xy, text, font, fill, bold=False, stroke_width=0):
    if bold and stroke_width <= 0:
        stroke_width = 2
    try:
        if bold:
            draw.text(
                xy,
                text,
                fill=fill,
                font=font,
                stroke_width=stroke_width,
                stroke_fill=fill,
            )
        else:
            draw.text(xy, text, fill=fill, font=font)
    except TypeError:
        if bold:
            offsets = [(0, 0), (1, 0), (0, 1), (1, 1)]
            for dx, dy in offsets:
                draw.text((xy[0] + dx, xy[1] + dy), text, fill=fill, font=font)
        else:
            draw.text(xy, text, fill=fill, font=font)


def _get_text_size(draw, text, font, bold=False, stroke_width=0):
    if hasattr(draw, 'textbbox'):
        try:
            x0, y0, x1, y1 = draw.textbbox(
                (0, 0),
                text,
                font=font,
                stroke_width=stroke_width if bold else 0,
            )
        except TypeError:
            x0, y0, x1, y1 = draw.textbbox((0, 0), text, font=font)
        return x1 - x0, y1 - y0
    w, h = draw.textsize(text, font=font)
    if bold:
        w += stroke_width * 2
        h += stroke_width * 2
    return w, h


def fit_text_to_width(draw, text, font, max_w, bold=False, stroke_width=0):
    text = text or ''
    if max_w <= 0:
        return ''

    text_w, _ = _get_text_size(draw, text, font, bold=bold, stroke_width=stroke_width)
    if text_w <= max_w:
        return text

    ellipsis = '...'
    ellipsis_w, _ = _get_text_size(draw, ellipsis, font, bold=bold, stroke_width=stroke_width)
    if ellipsis_w > max_w:
        return ''

    trimmed = text
    while trimmed:
        trimmed = trimmed[:-1]
        candidate = trimmed + ellipsis
        cand_w, _ = _get_text_size(draw, candidate, font, bold=bold, stroke_width=stroke_width)
        if cand_w <= max_w:
            return candidate
    return ellipsis


def bgr_to_rgb(color):
    return (color[2], color[1], color[0])


def resize_with_points(img, pts, target_long_side):
    if target_long_side is None or target_long_side <= 0:
        return img, pts
    w, h = img.size
    max_dim = max(w, h)
    if max_dim <= 0 or max_dim == target_long_side:
        return img, pts
    scale = float(target_long_side) / float(max_dim)
    new_w = max(1, int(round(w * scale)))
    new_h = max(1, int(round(h * scale)))
    img = img.resize((new_w, new_h), Image.BILINEAR)
    if pts is not None and pts.size > 0:
        pts = pts * scale
    return img, pts


def add_padding_with_text(img, text, pad, font, text_color, bg_color, bold, stroke_width):
    if pad is None or pad <= 0:
        return img
    draw_tmp = ImageDraw.Draw(img)
    text = text or ''
    text_w, text_h = _get_text_size(draw_tmp, text, font, bold=bold, stroke_width=stroke_width)
    min_text_gap = 24
    min_pad = text_h + (2 * min_text_gap)
    pad = max(pad, min_pad)
    new_w = img.width + pad * 2
    new_h = img.height + pad * 2
    canvas = Image.new('RGB', (new_w, new_h), color=bg_color)
    canvas.paste(img, (pad, pad))

    draw = ImageDraw.Draw(canvas)
    max_text_w = max(0, new_w - (2 * pad))
    text = fit_text_to_width(draw, text, font, max_text_w, bold=bold, stroke_width=stroke_width)
    text_w, text_h = _get_text_size(draw, text, font, bold=bold, stroke_width=stroke_width)
    text_x = pad
    text_y = max(min_text_gap, (pad - text_h) // 2)
    text_y = min(text_y, max(0, pad - text_h - min_text_gap))
    draw_text(draw, (text_x, text_y), text, font, text_color, bold=bold, stroke_width=stroke_width)
    return canvas


def annotate_panel(
    img_bgr,
    pts,
    title_text,
    point_color_bgr,
    radius,
    font,
    text_color,
    title_bg,
    target_long_side,
    pad,
):
    rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    img = Image.fromarray(rgb)
    img, pts = resize_with_points(img, pts, target_long_side)
    draw = ImageDraw.Draw(img)

    max_dim = max(img.width, img.height)
    auto_radius = max(3, int(round(max_dim * 0.004)))
    if radius is None or radius < auto_radius:
        radius = auto_radius

    if pts is not None and pts.size > 0:
        color = bgr_to_rgb(point_color_bgr)
        for x, y in pts:
            x0 = x - radius
            y0 = y - radius
            x1 = x + radius
            y1 = y + radius
            draw.ellipse((x0, y0, x1, y1), fill=color, outline=color)

    return add_padding_with_text(
        img,
        title_text or '',
        pad,
        font,
        text_color,
        title_bg,
        bold=False,
        stroke_width=0,
    )


def _load_state_dict(weight_path: Path):
    if not weight_path.exists():
        raise FileNotFoundError(f'Weights file not found: {weight_path}')

    if weight_path.suffix == '.safetensors':
        try:
            from safetensors.torch import load_file as load_safetensors
        except ImportError as exc:
            raise ImportError(
                'safetensors is required to load .safetensors weights. Install with: pip install safetensors'
            ) from exc
        return load_safetensors(str(weight_path), device='cpu')

    checkpoint = torch.load(str(weight_path), map_location='cpu')
    if isinstance(checkpoint, dict) and 'model' in checkpoint and isinstance(checkpoint['model'], dict):
        return checkpoint['model']
    if isinstance(checkpoint, dict) and checkpoint and all(torch.is_tensor(v) for v in checkpoint.values()):
        return checkpoint
    raise ValueError(
        'Unsupported checkpoint format. Expected .safetensors or .pth containing a model state_dict.'
    )


@torch.no_grad()
def infer_pet_points(model, frame_bgr, device, upper_bound):
    frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
    resized_rgb, scale = resize_for_eval(frame_rgb, upper_bound)
    resized_h, resized_w = resized_rgb.shape[:2]

    img = Image.fromarray(resized_rgb)
    img = PET_TRANSFORM(img)
    samples = utils.nested_tensor_from_tensor_list([img]).to(device)
    img_h, img_w = samples.tensors.shape[-2:]

    outputs = model(samples, test=True)
    outputs_points = outputs['pred_points']
    if outputs_points.dim() == 3:
        outputs_points = outputs_points[0]
    pred_points = outputs_points.detach().cpu().numpy()

    if pred_points.size == 0:
        return np.zeros((0, 2), dtype=np.float32), scale

    pred_points[:, 0] *= float(img_h)
    pred_points[:, 1] *= float(img_w)

    pred_points[:, 0] = np.clip(pred_points[:, 0], 0.0, float(resized_h - 1))
    pred_points[:, 1] = np.clip(pred_points[:, 1], 0.0, float(resized_w - 1))

    if scale != 1.0:
        pred_points = pred_points / float(scale)

    orig_h, orig_w = frame_bgr.shape[:2]
    pred_points[:, 0] = np.clip(pred_points[:, 0], 0.0, float(orig_h - 1))
    pred_points[:, 1] = np.clip(pred_points[:, 1], 0.0, float(orig_w - 1))

    points_xy = np.stack([pred_points[:, 1], pred_points[:, 0]], axis=1)
    return points_xy, scale


def main(args) -> None:
    device = resolve_device(args.device)

    model, _ = build_model(args)
    model.to(device)
    model.eval()

    state_dict = _load_state_dict(Path(args.resume))
    model.load_state_dict(state_dict, strict=True)

    image_path = Path(args.image_path)
    frame_bgr = cv2.imread(str(image_path))
    if frame_bgr is None:
        raise ValueError(f'Failed to read image: {image_path}')

    points_xy, scale = infer_pet_points(model, frame_bgr, device, args.upper_bound)
    count = int(points_xy.shape[0]) if points_xy.size > 0 else 0

    result = {
        'image': str(image_path),
        'count': count,
        'points_xy': points_xy.tolist(),
        'scale': scale,
    }

    print(f'image: {result["image"]}')
    print(f'predicted_count: {result["count"]}')

    if args.output_json:
        output_json = Path(args.output_json)
        output_json.parent.mkdir(parents=True, exist_ok=True)
        output_json.write_text(json.dumps(result, indent=2))
        print(f'json_saved_to: {output_json}')

    if args.output_image:
        output_image = Path(args.output_image)
        output_image.parent.mkdir(parents=True, exist_ok=True)

        panel = annotate_panel(
            frame_bgr,
            points_xy,
            f'{args.title_text} Count : {count}',
            parse_color(args.point_color),
            args.radius,
            load_font(font_size=args.panel_font_size, bold=False),
            text_color=(0, 0, 0),
            title_bg=(255, 255, 255),
            target_long_side=args.panel_long_side,
            pad=args.panel_pad,
        )
        panel.save(str(output_image))
        print(f'annotated_image_saved_to: {output_image}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        'PET single-image inference',
        parents=[get_args_parser()],
    )
    main(parser.parse_args())