WPMNet

Official implementation of the paper:

Weight-posture joint modeling for body weight estimation of unconstrained pigs
Peiguang Xin, Hao Wang, Juncheng Ma, Chao Liang, Haoyan Ma, Chaoyuan Wang, Jingguan Cao, Peipei Yang
Smart Agricultural Technology, Volume 15, 2026, Article 102679.

WPMNet is a weight-posture joint modeling framework for non-contact pig body weight estimation. It takes paired RGB-D images as input and combines intermediate feature-level fusion, hierarchical attention, and joint weight/posture learning to improve robustness to posture variation in unrestrained pigs.

Paper at a glance

Item	Value
Task	Pig body weight estimation under diverse postures
Input	Paired RGB and depth images
Model	Dual-stream RGB-D network with hierarchical attention fusion and joint weight/posture prediction
Reported result	R² = 0.759, MAPE = 7.122%

Repository Structure

.
├── configs/
│   ├── best.yaml             # Public WPMNet training config
│   └── test_config.yaml      # Test / evaluation config
├── data/
│   └── README.md             # Expected dataset and split-file format
├── scripts/
│   └── rewrite_split_paths.py # Rewrite absolute paths in split files
├── src/
│   ├── core/                 # Trainer, dataset, metrics, loss, model builder
│   ├── models/               # Backbones, fusion blocks, heads
│   └── utils/                # Label normalization and visualization helpers
├── train.py                  # Hydra training entry
├── test.py                   # Test / evaluation entry
├── run_experiments.py        # Internal sequential experiment runner
├── requirements.txt          # Python dependencies
├── CITATION.cff              # Machine-readable citation metadata
└── README.md

The released checkpoint is expected at:

weights/
└── wpmnet_best_model.pth

Installation

The code was developed with Python 3.11 and PyTorch 2.6.0.

conda create -n wpmnet python=3.11 -y
conda activate wpmnet

pip install -r requirements.txt

If you need to install PyTorch manually for a specific CUDA version, install PyTorch first from the official PyTorch instructions, then install the remaining packages:

pip install -r requirements.txt

Quick start

# 1) Create the environment
conda create -n wpmnet python=3.11 -y
conda activate wpmnet
pip install -r requirements.txt

# 2) Request access on Harvard Dataverse, then download:
#    Pigs_weight.7z.001 ... Pigs_weight.7z.010
#    train.txt, val.txt, test.txt
#    Extract the archives into data/700x700/.

# 3) Rewrite author-side absolute paths in the released split files
python scripts/rewrite_split_paths.py \
  --split-dir data/700x700 \
  --data-root data/700x700 \
  --inplace

# 4) Download the released checkpoint and place it at:
#    weights/wpmnet_best_model.pth

# 5) Evaluate the released checkpoint
python test.py

# 6) Train from scratch (optional; start MLflow first; see below)
torchrun --nproc_per_node=2 train.py --config-path=configs --config-name=best

test.py writes outputs/test_predictions.csv and outputs/test_metrics.json.

Dataset

The dataset is available on Harvard Dataverse:

Pigs_weight [Data set], Harvard Dataverse, 2024.
DOI: 10.7910/DVN/AZ5IBM

It contains RGB-D image pairs from 147 pigs covering standing, abdominal lying, and lateral lying postures. The dataset is released under CC0 1.0 and is not redistributed in this repository; please follow the access instructions on the Dataverse page.

The dataset archive contains 10 multi-volume parts (Pigs_weight.7z.001 ... Pigs_weight.7z.010) plus train.txt, val.txt, and test.txt. Download all parts, then extract from the directory containing them:

7z x Pigs_weight.7z.001 -odata/700x700

On Windows, select Pigs_weight.7z.001, choose 7-Zip → Extract, and use data/700x700 as the destination. 7-Zip will continue with the remaining volumes automatically. If the archive creates a top-level folder, move its contents so the layout below is satisfied.

The released split files may contain absolute paths from the authors' storage. Rewrite them to your local data root:

python scripts/rewrite_split_paths.py \
  --split-dir data/700x700 \
  --data-root data/700x700 \
  --inplace

This keeps .bak backups. Omit --inplace to write train.local.txt, val.local.txt, and test.local.txt instead; in that case, update the three data paths in configs/best.yaml accordingly.

After extraction and path rewriting, the layout should look like:

data/700x700/
├── train.txt
├── val.txt
├── test.txt
├── train/
│   ├── rgb/...
│   └── depth/...
├── val/
│   ├── rgb/...
│   └── depth/...
└── test/
    ├── rgb/...
    └── depth/...

weights/
└── wpmnet_best_model.pth

The split files use the following format:

rgb_path,depth_path,weight_kg,posture_name

where posture_name is one of standing, lyingonstomach, or lyingonside. See data/README.md for details. After preparing the data, update the paths in configs/best.yaml.

Pretrained Weights

Pretrained weights are available from Baidu Netdisk:

Link: https://pan.baidu.com/s/1PBvG9Uh4JY7ihBHErcUEbg?pwd=uq8a
code: uq8a

Place the released checkpoint here:

weights/
└── wpmnet_best_model.pth

The checkpoint must match the architecture in configs/best.yaml.

When evaluating a released checkpoint, set the checkpoint path in:

configs/test_config.yaml

testing:
  checkpoint: weights/wpmnet_best_model.pth

MLflow Logging

Training logs metrics and checkpoints with MLflow. The default config expects an MLflow server at:

http://localhost:2000

Start a local MLflow server before training:

mlflow server \
  --host 127.0.0.1 \
  --port 2000 \
  --backend-store-uri /path/to/mlflow/backend \
  --default-artifact-root /path/to/mlflow/artifacts

Then open:

http://127.0.0.1:2000

If you use a different MLflow URI, update:

configs/best.yaml

Training

Training uses Hydra configs. The public release exposes the best model config as configs/best.yaml, and train.py uses it by default.

For one GPU:

torchrun --nproc_per_node=1 train.py --config-path=configs --config-name=best

For two GPUs:

torchrun --nproc_per_node=2 train.py --config-path=configs --config-name=best

You can override training parameters from the command line:

torchrun --nproc_per_node=2 train.py \
  --config-path=configs \
  --config-name=best \
  training.parameter.epochs=100 \
  training.parameter.batch_size=64 \
  optimizer.lr=0.00005

Useful visualization/logging overrides for faster training:

torchrun --nproc_per_node=2 train.py \
  --config-path=configs \
  --config-name=best \
  visualization.log_images=false \
  visualization.log_model_graph=false \
  visualization.visualization_interval=0 \
  visualization.log_interval=0

Checkpoints are saved by the trainer as MLflow artifacts, depending on the active run configuration.

configs/best.yaml enables DDP and uses seed: 20. Use torchrun, which sets LOCAL_RANK. For a single process without DDP, override training.parameter.ddp=false:

python train.py \
  --config-path=configs \
  --config-name=best \
  training.parameter.ddp=false

Training requires a running MLflow server because mlflow.tracking_uri defaults to http://localhost:2000.

Testing

Set the test config and checkpoint in:

configs/test_config.yaml

Example:

defaults:
  - best
  - _self_

testing:
  checkpoint: weights/wpmnet_best_model.pth
  output_csv: outputs/test_predictions.csv
  metrics:
    regression:
      - mae
      - rmse
      - r2
      - mape
    classification:
      - accuracy
      - f1

Run:

python test.py

or explicitly:

python test.py --config-path=configs --config-name=test_config

The script writes:

• outputs/test_predictions.csv: per-image true/predicted weight, error, and posture labels
• outputs/test_metrics.json: MAE, RMSE, R², MAPE, posture accuracy, and macro-F1

Important notes:

• The experiment config in configs/test_config.yaml must match the checkpoint architecture.
• The test script uses the training split weights to initialize label normalization, matching the training procedure.
• Reported regression metrics include MAE, RMSE, R2, and MAPE.
• Reported posture metrics include accuracy and macro-F1.
• The reference test path uses single-view inference; no test-time augmentation is applied.
• By default, outputs are written under outputs/, which is ignored by git.

Reference results

The paper reports R² = 0.759 and MAPE = 7.122% for WPMNet on the posture-diverse test set. Exact values may vary slightly with the released checkpoint and software/hardware environment.

Reproducibility notes

• Obtain the dataset as described in the paper and download all 10 archive parts plus train.txt, val.txt, and test.txt.
• Run scripts/rewrite_split_paths.py if the released split files contain absolute paths.
• The public config uses seed: 20, 100 epochs, batch size 64 per process, and DDP. Bitwise reproduction is not guaranteed because of GPU nondeterminism and random augmentation.
• Training requires a running MLflow server; testing does not.
• test.py performs single-view inference and writes per-image predictions and metrics to outputs/.
• Reported values can vary slightly with checkpoint version, CUDA/cuDNN, and hardware.

Troubleshooting

Symptom	Fix
FileNotFoundError: weights/wpmnet_best_model.pth	Download the released checkpoint and place it at that path.
FileNotFoundError: data/700x700/train.txt	Download/extract the dataset and run the split-path rewrite script.
LOCAL_RANK / torch.distributed error	Use torchrun, or set training.parameter.ddp=false.
MLflow connection error	Start the MLflow server or change mlflow.tracking_uri in configs/best.yaml.
Address already in use on port 2000	Stop the old MLflow process or use another port and update mlflow.tracking_uri.
CUDA out of memory	Reduce training.parameter.batch_size or use more GPUs.

Citation

If you use this code, the released checkpoint, or the dataset, please cite the paper:

@article{xin2026weight,
  title   = {Weight-posture joint modeling for body weight estimation of unconstrained pigs},
  author  = {Xin, Peiguang and Wang, Hao and Ma, Juncheng and Liang, Chao and Ma, Haoyan and Wang, Chaoyuan and Cao, Jingguan and Yang, Peipei},
  journal = {Smart Agricultural Technology},
  volume  = {15},
  pages   = {102679},
  year    = {2026},
  doi     = {10.1016/j.atech.2026.102679}
}

A machine-readable citation is available in CITATION.cff.

License

The published article is available under the CC BY 4.0 license. If you use this code, the released model, or the dataset, please cite the paper above.

Contact

• Corresponding author: Chao Liang, China Agricultural University, liangchao@cau.edu.cn
• Code questions: open an issue or contact B20243090880@cau.edu.cn.