Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -7,3 +7,12 @@ tmp.sh
log*
datadir*
debug*

# Package build artifacts
dist/
build/
*.egg-info/
*.egg

# Plan file
plan.md
66 changes: 62 additions & 4 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,71 @@ University of Oxford.


## Table of Content
- [Environment](#Environment)
- [Installation](#installation)
- [Quick Start](#quick-start)
- [Environment (conda, legacy)](#environment-conda-legacy)
- [Data](#Data)
- [Training](#Training)
- [Inferencing](#Inferencing)

## Environment
We provide a `environment.yaml` file to set up a `conda` environment:
## Installation

### GPU (with CUDA) — one line
```bash
git clone https://github.com/ActiveVisionLab/CrossScore.git
cd CrossScore
conda env create -f environment_gpu.yaml && conda activate CrossScore
```
This installs Python, PyTorch with CUDA 12.1, and all CrossScore dependencies in a single command.

### CPU only — one line
```bash
git clone https://github.com/ActiveVisionLab/CrossScore.git
cd CrossScore
conda env create -f environment_cpu.yaml && conda activate CrossScore
```

> **Note:** If you use the CPU install, CrossScore will print a reminder at runtime on how to switch to the GPU version for faster inference.

## Quick Start

### Python API
```python
from crossscore.api import score

# Score query images against reference images
# Model checkpoint is auto-downloaded on first use (~129MB)
results = score(
query_dir="path/to/query/images",
reference_dir="path/to/reference/images",
)

# Per-image mean scores
print(results["scores"]) # [0.82, 0.91, 0.76, ...]

# Score map tensors (pixel-level quality maps)
for score_map in results["score_maps"]:
print(score_map.shape) # (batch_size, H, W)

# Colorized score map PNGs are written to results["out_dir"]
```

### Command Line
```bash
python -m crossscore.cli --query-dir path/to/queries --reference-dir path/to/references

# With options
python -m crossscore.cli --query-dir renders/ --reference-dir gt/ --metric-type mae --batch-size 4

# Force CPU mode
python -m crossscore.cli --query-dir renders/ --reference-dir gt/ --cpu
```

### Environment Variables
- `CROSSSCORE_CKPT_PATH`: Use a specific local checkpoint instead of auto-downloading

## Environment (conda, legacy)
For training and development, we provide the full pinned `environment.yaml`:
```bash
git clone https://github.com/ActiveVisionLab/CrossScore.git
cd CrossScore
Expand Down Expand Up @@ -86,7 +144,7 @@ on our project page.
- [ ] Create a HuggingFace demo page.
- [ ] Release ECCV quantitative results related scripts.
- [x] Release [data processing scripts](https://github.com/ziruiw-dev/CrossScore-3DGS-Preprocessing)
- [ ] Release PyPI and Conda package.
- [x] Release Conda package.

## Acknowledgement
This research is supported by an
Expand Down
38 changes: 38 additions & 0 deletions crossscore/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
"""CrossScore: Towards Multi-View Image Evaluation and Scoring.

A pip-installable package for neural image quality assessment using
cross-reference scoring with DINOv2 backbone.

Example:
>>> import crossscore
>>> results = crossscore.score(
... query_dir="path/to/query/images",
... reference_dir="path/to/reference/images",
... )
>>> print(results["scores"]) # per-image mean scores
"""

__version__ = "1.0.0"


def score(*args, **kwargs):
"""Score query images against reference images using CrossScore.

See crossscore.api.score for full documentation.
"""
from crossscore.api import score as _score

return _score(*args, **kwargs)


def get_checkpoint_path():
"""Get path to the CrossScore checkpoint, downloading if necessary.

See crossscore._download.get_checkpoint_path for full documentation.
"""
from crossscore._download import get_checkpoint_path as _get

return _get()


__all__ = ["score", "get_checkpoint_path"]
33 changes: 33 additions & 0 deletions crossscore/_download.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
"""Utilities for downloading CrossScore model checkpoints."""

import os
from pathlib import Path

HF_REPO_ID = "ActiveVisionLab/CrossScore"
CHECKPOINT_FILENAME = "CrossScore-v1.0.0.ckpt"


def get_checkpoint_path() -> str:
"""Get path to the CrossScore checkpoint, downloading it if necessary.

Downloads from HuggingFace Hub on first use and caches locally.
Set environment variables to customize:
CROSSSCORE_CKPT_PATH - use a specific local checkpoint file

Returns:
Path to the checkpoint file.
"""
# Allow user to override with a custom path
custom_path = os.environ.get("CROSSSCORE_CKPT_PATH")
if custom_path:
if not Path(custom_path).exists():
raise FileNotFoundError(f"Checkpoint not found at CROSSSCORE_CKPT_PATH={custom_path}")
return custom_path

from huggingface_hub import hf_hub_download

path = hf_hub_download(
repo_id=HF_REPO_ID,
filename=CHECKPOINT_FILENAME,
)
return path
179 changes: 179 additions & 0 deletions crossscore/api.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,179 @@
"""High-level API for CrossScore image quality assessment."""

from pathlib import Path
from typing import Optional, Union, List

import numpy as np
import torch
from torch.utils.data import DataLoader
from torchvision.transforms import v2 as T
from omegaconf import OmegaConf
from tqdm import tqdm

from crossscore._download import get_checkpoint_path
from crossscore.utils.io.images import ImageNetMeanStd
from crossscore.dataloading.dataset.simple_reference import SimpleReference


def _write_score_maps(score_maps, query_paths, out_dir, metric_type, metric_min, metric_max):
"""Write score maps to disk as colorized PNGs."""
from PIL import Image
from crossscore.utils.misc.image import gray2rgb

vrange_vis = [metric_min, metric_max]
out_dir = Path(out_dir) / "score_maps"
out_dir.mkdir(parents=True, exist_ok=True)

idx = 0
for batch_maps, batch_paths in zip(score_maps, query_paths):
for score_map, qpath in zip(batch_maps, batch_paths):
fname = Path(qpath).stem + ".png"
rgb = gray2rgb(score_map.cpu().numpy(), vrange_vis)
Image.fromarray(rgb).save(out_dir / fname)
idx += 1
return str(out_dir)


def score(
query_dir: str,
reference_dir: str,
ckpt_path: Optional[str] = None,
metric_type: str = "ssim",
batch_size: int = 8,
num_workers: int = 4,
resize_short_side: int = 518,
device: Optional[str] = None,
out_dir: Optional[str] = None,
write_score_maps: bool = True,
) -> dict:
"""Score query images against reference images using CrossScore.

Args:
query_dir: Directory containing query images (e.g., NVS rendered images).
reference_dir: Directory containing reference images (e.g., real captured images).
ckpt_path: Path to model checkpoint. Auto-downloads if not provided.
metric_type: Metric type to predict. One of "ssim", "mae", "mse".
batch_size: Batch size for inference.
num_workers: Number of data loading workers.
resize_short_side: Resize images so short side equals this value. -1 to disable.
device: Device string ("cuda", "cuda:0", "cpu"). Auto-detected if None.
out_dir: Output directory for score maps. Defaults to "./crossscore_output".
write_score_maps: Whether to write colorized score map PNGs to disk.

Returns:
Dictionary with:
- "score_maps": List of score map tensors, each (B, H, W)
- "scores": List of per-image mean scores (float)
- "out_dir": Output directory path (if write_score_maps=True)

Example:
>>> import crossscore
>>> results = crossscore.score(
... query_dir="path/to/query/images",
... reference_dir="path/to/reference/images",
... )
>>> print(results["scores"]) # per-image mean scores
"""
from crossscore.task.core import load_model

# Determine device
if device is None:
if torch.cuda.is_available():
device = "cuda"
else:
device = "cpu"
print(
"Note: CUDA not available, running on CPU. "
"For GPU acceleration, install with:\n"
" conda env create -f environment_gpu.yaml"
)

# Get checkpoint
if ckpt_path is None:
ckpt_path = get_checkpoint_path()

# Load model
model = load_model(ckpt_path, device=device)

# Set up data transforms
img_norm_stat = ImageNetMeanStd()
transforms = {
"img": T.Normalize(mean=img_norm_stat.mean, std=img_norm_stat.std),
}
if resize_short_side > 0:
transforms["resize"] = T.Resize(
resize_short_side,
interpolation=T.InterpolationMode.BILINEAR,
antialias=True,
)

# Build dataset and dataloader
neighbour_config = {"strategy": "random", "cross": 5, "deterministic": False}
dataset = SimpleReference(
query_dir=query_dir,
reference_dir=reference_dir,
transforms=transforms,
neighbour_config=neighbour_config,
return_item_paths=True,
zero_reference=False,
)

dataloader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=(device != "cpu"),
persistent_workers=False,
)

# Run inference
all_score_maps = []
all_scores = []
all_query_paths = []

with torch.no_grad():
for batch in tqdm(dataloader, desc="CrossScore"):
query_img = batch["query/img"].to(device)
ref_imgs = batch.get("reference/cross/imgs")
if ref_imgs is not None:
ref_imgs = ref_imgs.to(device)

outputs = model(
query_img=query_img,
ref_cross_imgs=ref_imgs,
norm_img=False,
)

score_map = outputs["score_map_ref_cross"] # (B, H, W)
all_score_maps.append(score_map.cpu())

# Per-image mean score
for i in range(score_map.shape[0]):
all_scores.append(score_map[i].mean().item())

# Track query paths for output naming
if "item_paths" in batch and "query/img" in batch["item_paths"]:
all_query_paths.append(batch["item_paths"]["query/img"])

# Build results
metric_min = -1 if metric_type == "ssim" else 0
if metric_type == "ssim":
metric_min = 0 # CrossScore predicts SSIM in [0, 1] by default

results = {
"score_maps": all_score_maps,
"scores": all_scores,
}

# Write outputs
if write_score_maps and all_score_maps:
if out_dir is None:
out_dir = "./crossscore_output"
written_dir = _write_score_maps(
all_score_maps, all_query_paths, out_dir,
metric_type, metric_min, metric_max=1,
)
results["out_dir"] = written_dir

return results
Loading