cattle-reid-hanwoo-vit-b16-phe

HanwooReID ViT-B/16 + PHE — Cattle Re-Identification

A supervised cattle re-identification model following the HanwooReID approach: a Vision Transformer (ViT-B/16) backbone with a Pose-Guided Heatmap Encoder (PHE), a BatchNorm embedding head and a linear classifier.

Hanwoo cattle have no distinctive coat markings, so appearance-only re-identification fails. The pose-guided encoder injects keypoint heatmaps into the ViT patch embeddings to make the embedding view-invariant.

Files

Training

• Architecture: ViT-B/16 (patch 16, image 256×256, timm vit_base_patch16_224) + PHE (16×16 Gaussian heatmaps, 12 keypoints)
  • BatchNorm1d embedding + linear classifier (514 identities)
• Dataset: 514 identities, 48,222 images (train/gallery/query splits)
• Hardware: RTX 5080 (16 GB), AMP fp16 + torch.compile
• Recipe (critical for convergence):
  • Backbone at 0.1× the head LR (--backbone-lr-factor 0.1, backbone 3e-5, head/PHE/BN 3e-4, AdamW, cosine schedule, 5-epoch warmup)
  • Classifier weighted low relative to the metric loss (--weight-ce 0.3, loss = 0.3·CE(focal γ=2, label smoothing 0.1)
    • 1.0·hard-triplet, margin 0.3)
  • Mixup disabled — mixing the image while feeding original labels to the triplet loss corrupts the metric supervision
• Without this recipe the eval collapses to chance (R1 19.5% at epoch 2, then ~1%): the CE head memorizes training identities and destroys the pretrained ViT features. With it, Rank-1 reaches 100% by epoch 42.

Results (query vs gallery, 3446 queries)

PHE adds no measurable gain at this scale (with/without PHE are equal); it is kept for view-invariance on CCTV footage.

Usage

The checkpoint is a raw torch.save dict:

import torch

ckpt = torch.load("hanwoo_reid_best.pth", map_location="cpu", weights_only=False)
state = ckpt["state"]                      # state dict (keys without _orig_mod.)
ncls  = ckpt["config"]["ncls"] if "ncls" in ckpt["config"] else None
# number of classes is inferred from the classifier head:
ncls  = state["classifier.weight"].shape[0]

import timm
class PHE(torch.nn.Module): ...            # see repository training/reid_common.py
model = ...                                # HanwooReID(ncls) — ViT + PHE + BN head
model.load_state_dict(state, strict=False)
model.eval()

Embedding recipe (see training/reid_common.py in the source repo):

1. Crop the cow, letterbox_crop to 256×256, normalize (x/255 - 0.5)/0.5
2. Build the 16×16 pose heatmap (12 keypoints → 16×16 Gaussian maps)
3. emb, logits = model(x, heatmap); use emb (cosine similarity)

The same checkpoint is used end-to-end in a cross-camera video pipeline (DeepSORT + global ID fusion across videos) — see the source repository training/video_reid_v6.py.

Contact

Model trained on a private cattle dataset (514 identities). For access to the training data or the full pipeline, open an issue on the source repo.