Compute the LPIPS metric. Args: arr: (*b_shape, *d_shape) Assumes the RGB image is in [0,1] ref: (*b_shape, *d_shape) Assumes the RGB image is in [0,1] ndim_b: number of dimension of b_shape. If None, = 0. Returns: lpips_score: (*b_shape,)
(
arr: torch.Tensor,
ref: torch.Tensor,
ndim_b: int = None,
lpips_model: torch.nn.Module = None,
device: torch.device = torch.device('cuda'),
)
| 260 | |
| 261 | |
| 262 | def compute_lpips( |
| 263 | arr: torch.Tensor, |
| 264 | ref: torch.Tensor, |
| 265 | ndim_b: int = None, |
| 266 | lpips_model: torch.nn.Module = None, |
| 267 | device: torch.device = torch.device('cuda'), |
| 268 | ): |
| 269 | """ |
| 270 | Compute the LPIPS metric. |
| 271 | Args: |
| 272 | arr: (*b_shape, *d_shape) Assumes the RGB image is in [0,1] |
| 273 | ref: (*b_shape, *d_shape) Assumes the RGB image is in [0,1] |
| 274 | ndim_b: |
| 275 | number of dimension of b_shape. If None, = 0. |
| 276 | |
| 277 | Returns: |
| 278 | lpips_score: (*b_shape,) |
| 279 | |
| 280 | Note: |
| 281 | This function is NOT differentiable |
| 282 | """ |
| 283 | if ndim_b is None: |
| 284 | ndim_b = 0 |
| 285 | |
| 286 | ori_shape = arr.shape |
| 287 | b_shape = ori_shape[:ndim_b] |
| 288 | d_shape = ori_shape[ndim_b:] |
| 289 | arr = arr.reshape(-1, *d_shape) # (b, *d) |
| 290 | ref = ref.reshape(-1, *d_shape) # (b, *d) |
| 291 | b = arr.size(0) |
| 292 | |
| 293 | arr_device = arr.device |
| 294 | |
| 295 | arr = arr.to(device=device) |
| 296 | ref = ref.to(dtype=arr.dtype, device=device) |
| 297 | |
| 298 | assert len(d_shape) == 3 |
| 299 | assert d_shape[-1] >= 3 |
| 300 | |
| 301 | scores = [] |
| 302 | for ib in range(b): |
| 303 | score = metrics.lpips(rgb=arr[ib], gts=ref[ib], lpips_model=lpips_model) # float |
| 304 | scores.append(score) |
| 305 | scores = torch.tensor(scores, dtype=torch.float, device=arr_device) |
| 306 | scores = scores.reshape(*b_shape) |
| 307 | return scores.to(device=arr_device) |
| 308 | |
| 309 | |
| 310 | def compute_l1( |