PST-PCQA / README.md

Update README.md

a9b4836 verified 27 days ago

7.85 kB

---
license: mit
library_name: pytorch
tags:
- point-cloud
- point-cloud-quality-assessment
- no-reference
- quality-assessment
- 3d
- pytorch-lightning
paper: https://ieeexplore.ieee.org/abstract/document/10948465
language:
- en
---
# PST-PCQA: Low-Complexity Patch-based No-Reference Point Cloud Quality Metric

PST-PCQA is a learning-based **no-reference point cloud quality assessment (NR-PCQA)** model. Given a distorted, coloured point cloud (no reference needed), it predicts a Mean Opinion Score (MOS) by splitting the cloud into local patches, extracting structure and texture features from each patch, predicting a quality score and an importance weight per patch, and combining them into a global score.

- **Paper:** M. Neri, F. Battisti, *"Low-Complexity Patch-Based No-Reference Point Cloud Quality Metric Exploiting Weighted Structure and Texture Features"*, IEEE Transactions on Broadcasting, vol. 71, no. 2, pp. 631–640, 2025. [DOI: 10.1109/TBC.2025.3553305](https://doi.org/10.1109/TBC.2025.3553305) · [arXiv:2503.15001](https://arxiv.org/abs/2503.15001)
- **Code:** https://github.com/michaelneri/PST-PCQA
- **Authors:** Michael Neri (Tampere University, Finland), Federica Battisti (University of Padova, Italy)
- **License:** MIT
- **Parameters:** ~1.8 M (paper; also matches the count of learnable weights in the released checkpoints: 1,797,697 for K=16)
- **Framework:** PyTorch + PyTorch Lightning (checkpoints saved with Lightning 2.3.0)

![PST-PCQA architecture](https://raw.githubusercontent.com/michaelneri/PST-PCQA/main/figures/Fig2.jpg)

## Released checkpoints

Both checkpoints were trained on the **WPC** dataset. No SJTU-PCQA or SIAT-PCQD checkpoints are provided in the repository. If needed, please contact Michael Neri at michael.neri@tuni.fi

| File | Patches K | Trained on | Saved at epoch |
|---|---|---|---|
| `PST_PCQAModule_K_16.ckpt` | 16 | WPC | 396 |
| `PST_PCQAModule_K_8.ckpt` | 8 | WPC | 233 |

The number of patches is baked into the network (two `BatchNorm1d(patches)` layers), so each checkpoint must be loaded with the matching `patches` value and fed exactly K patches.

## Model details

**Pre-processing (per point cloud)**
1. Point coordinates are normalised to the range [1, 2001].
2. K patch centres are chosen with Farthest Point Sampling, and each patch is formed with k-NN (14,900 points per patch).
3. From each patch, two inputs are built (each point = `x, y, z, r, g, b`):
   - **structure input**: 1,024 randomly sampled points → tensor `x_b` of shape `(B, K, 1024, 6)`
   - **texture input**: 8,192 points selected by k-NN → tensor `x_s` of shape `(B, K, 8192, 6)`

**Network** (`model.py`)
- Two feature extractors with the same design (`S_TFE`), one for structure and one for texture, each made of two sampling-grouping layers (512 then 256 points, 32 neighbours) with grouped 1×1 convolutions and ELU activations.
- The two feature maps are concatenated, fused with a grouped convolution, and pooled with **Global Variance Pooling (GVP)**.
- Two heads per patch: a quality score `ŷ_p` and a weight `w_p`. The global score is the mean over patches of `w_p · ŷ_p`.
- The forward pass returns `(mos, mos_per_patch)`.

**Training objective:** MSE on the global score + MSE between each patch score and the global MOS (weights 1 and 1).

## Training setup

From the training scripts (`training_val_test_{WPC,SJTU,SIAT}.py`) and the paper:

| Setting | Value |
|---|---|
| Optimizer | Adam, weight decay 1e-5 |
| Learning rate | 1e-3, cosine annealing (T_max = 400, minimum 1e-4) |
| Epochs | 400 |
| Batch size | 4 in the paper and SJTU/SIAT scripts; 2 in the WPC script |
| Dropout | 0 (WPC, SJTU), 0.2 (SIAT) |
| Checkpoint selection | top-k checkpoints by validation MSE, then the one with the best test PLCC is kept |
| Hardware | NVIDIA RTX 4070 |
### Datasets and protocols (from the paper)
| Dataset | References | Distorted | Protocol |
|---|---|---|---|
| WPC | 20 | 740 | train/test split as in prior work (ref. [9] of the paper) |
| SJTU-PCQA | 10 | 378 | leave-one-out, 10 folds |
| SIAT-PCQD | 20 | 340 | leave-one-out, 20 folds |
The datasets are not distributed with this model and must be obtained from their original sources.
## Evaluation results
**Within-dataset performance (K = 16)**
| Dataset | PLCC ↑ | SRCC ↑ | KRCC ↑ | RMSE ↓ |
|---|---|---|---|---|
| WPC | 0.8821 | 0.8624 | 0.6854 | 10.5769 |
| SJTU-PCQA | 0.9593 | 0.9514 | 0.8049 | 0.6630 |
| SIAT-PCQD | 0.8304 | 0.7931 | 0.5785 | 0.0183 |
**Effect of the number of patches K (WPC)**
| K | PLCC ↑ | SRCC ↑ | KRCC ↑ | RMSE ↓ |
|---|---|---|---|---|
| 2 | 0.8075 | 0.7906 | 0.6251 | 13.1335 |
| 4 | 0.8208 | 0.8087 | 0.6235 | 12.5506 |
| 8 | 0.8772 | 0.8779 | 0.6965 | 10.7112 |
| 16 | 0.8821 | 0.8624 | 0.6854 | 10.5769 |
| 32 | 0.8438 | 0.8266 | 0.6390 | 11.7893 |
**Cross-dataset generalisation**
| Train → Test | SRCC ↑ | PLCC ↑ |
|---|---|---|
| WPC → SJTU-PCQA | 0.7413 | 0.7522 |
| SJTU-PCQA → WPC | 0.2737 | 0.3797 |
**Speed:** On average 70ms, enabling real-time inference. This has been computed on a NVIDIA RTX 4070.
## How to use
Install the dependencies from `environment.yml` (Python 3.10, PyTorch 2.5.1, PyTorch Lightning 2.3.0, Open3D 0.18.0, SciPy, pandas), then:
```python
import numpy as np
import torch
from model import PST_PCQAModule          # from the GitHub repository
from data import patches_extraction      
K = 16
model = PST_PCQAModule.load_from_checkpoint(
    "PST_PCQAModule_K_16.ckpt",
    points_texture=8192, points_structure=1024,
    dropout=0.0, patches=K, lr=1e-3,
).eval()
x_b, x_s = patches_extraction(
    filename="example_pointcloud/house_gsigma_0_tsigma_16.ply",
    number_patches=K, points_per_patch=14900, small_points_per_patch=8192,
)
# structure input: random subset of 1,024 points per patch
np.random.seed(0)
idx = np.random.choice(x_b.shape[1], 1024, replace=False)
x_b = torch.tensor(np.array(x_b[:, idx]), device=model.device).float()
x_s = torch.tensor(np.array(x_s), device=model.device).float()
with torch.no_grad():
    mos, mos_per_patch = model(x_b.unsqueeze(0), x_s.unsqueeze(0))
print(mos, mos_per_patch)
```
This mirrors `example_inference.ipynb`, which also shows how to plot the per-patch scores on the point cloud.

![Example of patch-wise scores](https://raw.githubusercontent.com/michaelneri/PST-PCQA/main/figures/example_inference.png)

N.B.: The model outputs a normalised score: for the WPC checkpoints, multiply by 100 to get the MOS on the WPC scale (0–100)


Or download the checkpoint from the Hugging Face Hub:

```python
import torch
from huggingface_hub import hf_hub_download
from model import PST_PCQAModule   # from the cloned GitHub repository
K = 16  # use 8 with "PST_PCQAModule_K_8.ckpt"
ckpt_path = hf_hub_download(
    repo_id="michaelneri/PST-PCQA",          
    filename=f"PST_PCQAModule_K_{K}.ckpt",
)
model = PST_PCQAModule.load_from_checkpoint(
    ckpt_path,
    map_location="cuda" if torch.cuda.is_available() else "cpu",
    points_texture=8192, points_structure=1024,
    dropout=0.0, patches=K, lr=1e-3,
).eval()
# Shape check with dummy inputs: (batch, K, points, xyz+rgb)
x_b = torch.rand(1, K, 1024, 6, device=model.device)   # structure input
x_s = torch.rand(1, K, 8192, 6, device=model.device)   # texture input
with torch.no_grad():
    mos, mos_per_patch = model(x_b, x_s)
print(mos.shape, mos_per_patch.shape)  # torch.Size([1]) torch.Size([1, 16])
```


## Citation

```bibtex
@ARTICLE{Neri_TBC_2025,
   author={Neri, Michael and Battisti, Federica},
   journal={IEEE Transactions on Broadcasting},
   title={Low-Complexity Patch-Based No-Reference Point Cloud Quality Metric Exploiting Weighted Structure and Texture Features},
   year={2025},
   volume={71},
   number={2},
   pages={631-640},
   doi={10.1109/TBC.2025.3553305}
}
```

	---
	license: mit
	library_name: pytorch
	tags:
	- point-cloud
	- point-cloud-quality-assessment
	- no-reference
	- quality-assessment
	- 3d
	- pytorch-lightning
	paper: https://ieeexplore.ieee.org/abstract/document/10948465
	language:
	- en
	---

	# PST-PCQA: Low-Complexity Patch-based No-Reference Point Cloud Quality Metric

	PST-PCQA is a learning-based no-reference point cloud quality assessment (NR-PCQA) model. Given a distorted, coloured point cloud (no reference needed), it predicts a Mean Opinion Score (MOS) by splitting the cloud into local patches, extracting structure and texture features from each patch, predicting a quality score and an importance weight per patch, and combining them into a global score.

	- Paper: M. Neri, F. Battisti, "Low-Complexity Patch-Based No-Reference Point Cloud Quality Metric Exploiting Weighted Structure and Texture Features", IEEE Transactions on Broadcasting, vol. 71, no. 2, pp. 631–640, 2025. [DOI: 10.1109/TBC.2025.3553305](https://doi.org/10.1109/TBC.2025.3553305) · [arXiv:2503.15001](https://arxiv.org/abs/2503.15001)
	- Code: https://github.com/michaelneri/PST-PCQA
	- Authors: Michael Neri (Tampere University, Finland), Federica Battisti (University of Padova, Italy)
	- License: MIT
	- Parameters: ~1.8 M (paper; also matches the count of learnable weights in the released checkpoints: 1,797,697 for K=16)
	- Framework: PyTorch + PyTorch Lightning (checkpoints saved with Lightning 2.3.0)

	![PST-PCQA architecture](https://raw.githubusercontent.com/michaelneri/PST-PCQA/main/figures/Fig2.jpg)

	## Released checkpoints

	Both checkpoints were trained on the WPC dataset. No SJTU-PCQA or SIAT-PCQD checkpoints are provided in the repository. If needed, please contact Michael Neri at michael.neri@tuni.fi

	\| File \| Patches K \| Trained on \| Saved at epoch \|
	\|---\|---\|---\|---\|
	\| `PST_PCQAModule_K_16.ckpt` \| 16 \| WPC \| 396 \|
	\| `PST_PCQAModule_K_8.ckpt` \| 8 \| WPC \| 233 \|

	The number of patches is baked into the network (two `BatchNorm1d(patches)` layers), so each checkpoint must be loaded with the matching `patches` value and fed exactly K patches.

	## Model details

	Pre-processing (per point cloud)
	1. Point coordinates are normalised to the range [1, 2001].
	2. K patch centres are chosen with Farthest Point Sampling, and each patch is formed with k-NN (14,900 points per patch).
	3. From each patch, two inputs are built (each point = `x, y, z, r, g, b`):
	- structure input: 1,024 randomly sampled points → tensor `x_b` of shape `(B, K, 1024, 6)`
	- texture input: 8,192 points selected by k-NN → tensor `x_s` of shape `(B, K, 8192, 6)`

	Network (`model.py`)
	- Two feature extractors with the same design (`S_TFE`), one for structure and one for texture, each made of two sampling-grouping layers (512 then 256 points, 32 neighbours) with grouped 1×1 convolutions and ELU activations.
	- The two feature maps are concatenated, fused with a grouped convolution, and pooled with Global Variance Pooling (GVP).
	- Two heads per patch: a quality score `ŷ_p` and a weight `w_p`. The global score is the mean over patches of `w_p · ŷ_p`.
	- The forward pass returns `(mos, mos_per_patch)`.

	Training objective: MSE on the global score + MSE between each patch score and the global MOS (weights 1 and 1).

	## Training setup

	From the training scripts (`training_val_test_{WPC,SJTU,SIAT}.py`) and the paper:

	\| Setting \| Value \|
	\|---\|---\|
	\| Optimizer \| Adam, weight decay 1e-5 \|
	\| Learning rate \| 1e-3, cosine annealing (T_max = 400, minimum 1e-4) \|
	\| Epochs \| 400 \|
	\| Batch size \| 4 in the paper and SJTU/SIAT scripts; 2 in the WPC script \|
	\| Dropout \| 0 (WPC, SJTU), 0.2 (SIAT) \|
	\| Checkpoint selection \| top-k checkpoints by validation MSE, then the one with the best test PLCC is kept \|
	\| Hardware \| NVIDIA RTX 4070 \|


	### Datasets and protocols (from the paper)

	\| Dataset \| References \| Distorted \| Protocol \|
	\|---\|---\|---\|---\|
	\| WPC \| 20 \| 740 \| train/test split as in prior work (ref. [9] of the paper) \|
	\| SJTU-PCQA \| 10 \| 378 \| leave-one-out, 10 folds \|
	\| SIAT-PCQD \| 20 \| 340 \| leave-one-out, 20 folds \|

	The datasets are not distributed with this model and must be obtained from their original sources.

	## Evaluation results

	Within-dataset performance (K = 16)

	\| Dataset \| PLCC ↑ \| SRCC ↑ \| KRCC ↑ \| RMSE ↓ \|
	\|---\|---\|---\|---\|---\|
	\| WPC \| 0.8821 \| 0.8624 \| 0.6854 \| 10.5769 \|
	\| SJTU-PCQA \| 0.9593 \| 0.9514 \| 0.8049 \| 0.6630 \|
	\| SIAT-PCQD \| 0.8304 \| 0.7931 \| 0.5785 \| 0.0183 \|

	Effect of the number of patches K (WPC)

	\| K \| PLCC ↑ \| SRCC ↑ \| KRCC ↑ \| RMSE ↓ \|
	\|---\|---\|---\|---\|---\|
	\| 2 \| 0.8075 \| 0.7906 \| 0.6251 \| 13.1335 \|
	\| 4 \| 0.8208 \| 0.8087 \| 0.6235 \| 12.5506 \|
	\| 8 \| 0.8772 \| 0.8779 \| 0.6965 \| 10.7112 \|
	\| 16 \| 0.8821 \| 0.8624 \| 0.6854 \| 10.5769 \|
	\| 32 \| 0.8438 \| 0.8266 \| 0.6390 \| 11.7893 \|

	Cross-dataset generalisation

	\| Train → Test \| SRCC ↑ \| PLCC ↑ \|
	\|---\|---\|---\|
	\| WPC → SJTU-PCQA \| 0.7413 \| 0.7522 \|
	\| SJTU-PCQA → WPC \| 0.2737 \| 0.3797 \|

	Speed: On average 70ms, enabling real-time inference. This has been computed on a NVIDIA RTX 4070.

	## How to use

	Install the dependencies from `environment.yml` (Python 3.10, PyTorch 2.5.1, PyTorch Lightning 2.3.0, Open3D 0.18.0, SciPy, pandas), then:

	```python
	import numpy as np
	import torch
	from model import PST_PCQAModule # from the GitHub repository
	from data import patches_extraction

	K = 16
	model = PST_PCQAModule.load_from_checkpoint(
	"PST_PCQAModule_K_16.ckpt",
	points_texture=8192, points_structure=1024,
	dropout=0.0, patches=K, lr=1e-3,
	).eval()

	x_b, x_s = patches_extraction(
	filename="example_pointcloud/house_gsigma_0_tsigma_16.ply",
	number_patches=K, points_per_patch=14900, small_points_per_patch=8192,
	)

	# structure input: random subset of 1,024 points per patch
	np.random.seed(0)
	idx = np.random.choice(x_b.shape[1], 1024, replace=False)
	x_b = torch.tensor(np.array(x_b[:, idx]), device=model.device).float()
	x_s = torch.tensor(np.array(x_s), device=model.device).float()

	with torch.no_grad():
	mos, mos_per_patch = model(x_b.unsqueeze(0), x_s.unsqueeze(0))
	print(mos, mos_per_patch)
	```

	This mirrors `example_inference.ipynb`, which also shows how to plot the per-patch scores on the point cloud.

	![Example of patch-wise scores](https://raw.githubusercontent.com/michaelneri/PST-PCQA/main/figures/example_inference.png)

	N.B.: The model outputs a normalised score: for the WPC checkpoints, multiply by 100 to get the MOS on the WPC scale (0–100)


	Or download the checkpoint from the Hugging Face Hub:

	```python
	import torch
	from huggingface_hub import hf_hub_download
	from model import PST_PCQAModule # from the cloned GitHub repository

	K = 16 # use 8 with "PST_PCQAModule_K_8.ckpt"
	ckpt_path = hf_hub_download(
	repo_id="michaelneri/PST-PCQA",
	filename=f"PST_PCQAModule_K_{K}.ckpt",
	)

	model = PST_PCQAModule.load_from_checkpoint(
	ckpt_path,
	map_location="cuda" if torch.cuda.is_available() else "cpu",
	points_texture=8192, points_structure=1024,
	dropout=0.0, patches=K, lr=1e-3,
	).eval()

	# Shape check with dummy inputs: (batch, K, points, xyz+rgb)
	x_b = torch.rand(1, K, 1024, 6, device=model.device) # structure input
	x_s = torch.rand(1, K, 8192, 6, device=model.device) # texture input
	with torch.no_grad():
	mos, mos_per_patch = model(x_b, x_s)
	print(mos.shape, mos_per_patch.shape) # torch.Size([1]) torch.Size([1, 16])
	```


	## Citation

	```bibtex
	@ARTICLE{Neri_TBC_2025,
	author={Neri, Michael and Battisti, Federica},
	journal={IEEE Transactions on Broadcasting},
	title={Low-Complexity Patch-Based No-Reference Point Cloud Quality Metric Exploiting Weighted Structure and Texture Features},
	year={2025},
	volume={71},
	number={2},
	pages={631-640},
	doi={10.1109/TBC.2025.3553305}
	}
	```