Compute the ssim. Args: arr: (*b_shape, *d_shape) ref: (*b_shape, *d_shape) ndim_b: number of dimension of b_shape. If None, = 0. # valid_mask: (*b_shape, *d_shape) Returns: ssim_scores: (*b_shape,) Note: This function is
(
arr: torch.Tensor,
ref: torch.Tensor,
ndim_b: int = None,
# valid_mask: torch.Tensor = None,
)
| 216 | |
| 217 | |
| 218 | def compute_ssim( |
| 219 | arr: torch.Tensor, |
| 220 | ref: torch.Tensor, |
| 221 | ndim_b: int = None, |
| 222 | # valid_mask: torch.Tensor = None, |
| 223 | ): |
| 224 | """ |
| 225 | Compute the ssim. |
| 226 | Args: |
| 227 | arr: (*b_shape, *d_shape) |
| 228 | ref: (*b_shape, *d_shape) |
| 229 | ndim_b: |
| 230 | number of dimension of b_shape. If None, = 0. |
| 231 | # valid_mask: (*b_shape, *d_shape) |
| 232 | |
| 233 | Returns: |
| 234 | ssim_scores: (*b_shape,) |
| 235 | |
| 236 | Note: |
| 237 | This function is NOT differentiable |
| 238 | """ |
| 239 | if ndim_b is None: |
| 240 | ndim_b = 0 |
| 241 | |
| 242 | ori_shape = arr.shape |
| 243 | b_shape = ori_shape[:ndim_b] |
| 244 | d_shape = ori_shape[ndim_b:] |
| 245 | arr = arr.reshape(-1, *d_shape) # (b, *d) |
| 246 | ref = ref.reshape(-1, *d_shape) # (b, *d) |
| 247 | b = arr.size(0) |
| 248 | |
| 249 | assert len(d_shape) == 3 |
| 250 | assert d_shape[-1] >= 3 |
| 251 | |
| 252 | ssim_scores = [] |
| 253 | for ib in range(b): |
| 254 | # if valid_mask is None: |
| 255 | ssim_score = metrics.ssim(rgb=arr[ib], gts=ref[ib]) # float |
| 256 | ssim_scores.append(ssim_score) |
| 257 | ssim_scores = torch.tensor(ssim_scores, dtype=torch.float, device=arr.device) |
| 258 | ssim_scores = ssim_scores.reshape(*b_shape) |
| 259 | return ssim_scores |
| 260 | |
| 261 | |
| 262 | def compute_lpips( |