Mirrors/ECCV2022-RIFE
1
0
Fork 0
mirror of https://github.com/hzwer/ECCV2022-RIFE.git synced 2025-12-16 16:37:51 +01:00
Code Issues Packages Projects Releases Wiki Activity

Support cpu ssim

Browse Source
This commit is contained in:
hzwer
2021-01-18 17:41:27 +08:00
committed by GitHub
parent 1dc2dba7d1
commit e515b1d364
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
benchmark/pytorch_msssim/__init__.py
View File

@@ -3,6 +3,7 @@ import torch.nn.functional as F
from math import exp from math import exp
import numpy as np import numpy as np
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def gaussian(window_size, sigma): def gaussian(window_size, sigma):
gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)]) gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
@@ -11,7 +12,7 @@ def gaussian(window_size, sigma):
def create_window(window_size, channel=1): def create_window(window_size, channel=1):
_1D_window = gaussian(window_size, 1.5).unsqueeze(1) _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
_2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0).cuda() _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0).to(device)
window = _2D_window.expand(channel, 1, window_size, window_size).contiguous() window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
return window return window
@@ -19,7 +20,7 @@ def create_window_3d(window_size, channel=1):
_1D_window = gaussian(window_size, 1.5).unsqueeze(1) _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
_2D_window = _1D_window.mm(_1D_window.t()) _2D_window = _1D_window.mm(_1D_window.t())
_3D_window = _2D_window.unsqueeze(2) @ (_1D_window.t()) _3D_window = _2D_window.unsqueeze(2) @ (_1D_window.t())
window = _3D_window.expand(1, channel, window_size, window_size, window_size).contiguous().cuda() window = _3D_window.expand(1, channel, window_size, window_size, window_size).contiguous().to(device)
return window return window