nawta/PolarFree-PSD-4ch-weights

PolarFree PSD 4-channel Weights

Pre-trained weights for PolarFree model trained on PSD (Polarization Specular Dataset).

Important: 4-channel Polarization Format

PSD dataset originally has 12 polarization angles (30° intervals), but these weights were trained with 4-channel format (0°, 45°, 90°, 135°) to match standard polarization camera output.

Angle Conversion

PSD 12 angles → 4 angles:
- 0°   → idx-01 (direct)
- 45°  → interpolate(idx-02, idx-03)
- 90°  → idx-04 (direct)
- 135° → interpolate(idx-05, idx-06)

Model Structure

stage1/
├── net_g_100000.pth      # Generator (72MB)
└── net_le_100000.pth     # Latent Encoder (2.5MB)
stage2/
├── net_g_50000.pth       # Generator (72MB)
├── net_le_dm_50000.pth   # Latent Encoder DM (2.4MB)
└── net_d_50000.pth       # Denoising Network (10MB)

Performance

Test Set (50 samples)

Stage	Iteration	PSNR (dB)	SSIM
Stage 1	100,000	34.85	0.9918
Stage 2	50,000	34.45	0.9914

Inference Time (512x768)

Stage	GPU (RTX 6000)	CPU
Stage 1	325 ms	~31 sec
Stage 2	271 ms	~22 sec

Usage

from huggingface_hub import hf_hub_download
import torch

# Download Stage 1 weights
net_g_path = hf_hub_download(
    repo_id="nawta/PolarFree-PSD-4ch-weights",
    filename="stage1/net_g_100000.pth"
)

# Load
checkpoint = torch.load(net_g_path)
model.load_state_dict(checkpoint["params"])

License

Private repository - for research use only.