Skip to content

feat: Device option for image loading when the image is a Pytorch Tensor #17

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
patsch007 merged 8 commits into from
Feb 18, 2026
Merged

feat: Device option for image loading when the image is a Pytorch Tensor #17

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
036ebe5
fix: add check for image dimensions during loading
Apr 9, 2025
c573452
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 9, 2025
c6766d2
refactor: refactor to comply with ruff
lukasbehammer Apr 9, 2025
44f7d2e
docs(load_data): add Exception to function docstring
lukasbehammer Apr 9, 2025
c12d261
Device can be set for torch tensor images
Apr 22, 2025
b9959d7
Merge remote-tracking branch 'refs/remotes/origin/dev' into dev
Apr 22, 2025
5025415
refactor: ensure ruff rules
Apr 23, 2025
2592d37
Merge branch 'dev' into dev
lukasbehammer Apr 23, 2025
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
2 changes: 1 addition & 1 deletion src/viqa/_metrics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ def __init__(self, data_range, normalize, **kwargs):
"normalize": normalize,
"chromatic": False,
"roi": None,
**kwargs,
"device": "cpu" ** kwargs,
}
self.score_val = None
self._name = None
Expand Down
8 changes: 8 additions & 0 deletions src/viqa/fr_metrics/fsim.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,9 @@ class FSIM(FullReferenceMetricsInterface):
If True, the input images are expected to be RGB images and FSIMc is
calculated. See [1]_. Passed to
:py:func:`piq.fsim`. See the documentation under [2]_.
device : Union[str, torch.device], default 'cpu'
Determines the device if the image is a PyTorch tensor,
e.g. "cuda", "cpu", "cuda:0", ...

Raises
------
Expand Down Expand Up @@ -197,6 +200,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r[im_slice, :, :],
img_m[im_slice, :, :],
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = fsim(
img_r_tensor,
Expand All @@ -210,6 +214,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r[:, im_slice, :],
img_m[:, im_slice, :],
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = fsim(
img_r_tensor,
Expand All @@ -223,6 +228,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r[:, :, im_slice],
img_m[:, :, im_slice],
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = fsim(
img_r_tensor,
Expand All @@ -246,6 +252,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r,
img_m,
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = fsim(
img_r_tensor,
Expand All @@ -268,6 +275,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r,
img_m,
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = fsim(
img_r_tensor,
Expand Down
8 changes: 8 additions & 0 deletions src/viqa/fr_metrics/msssim.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,9 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
Algorithm parameter, K2 (small constant, see [3]_).
Try a larger K2 constant (e.g. 0.4) if you get a negative or NaN
results.
device : Union[str, torch.device], default 'cpu'
Determines the device if the image is a PyTorch tensor,
e.g. "cuda", "cpu", "cuda:0", ...

.. seealso::
See :py:func:`.viqa.fr_metrics.ssim.structural_similarity` for more
Expand Down Expand Up @@ -202,6 +205,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r[im_slice, :, :],
img_m[im_slice, :, :],
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = multi_scale_ssim(
img_r_tensor,
Expand All @@ -214,6 +218,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r[:, im_slice, :],
img_m[:, im_slice, :],
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = multi_scale_ssim(
img_r_tensor,
Expand All @@ -226,6 +231,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r[:, :, im_slice],
img_m[:, :, im_slice],
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = multi_scale_ssim(
img_r_tensor,
Expand All @@ -249,6 +255,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r,
img_m,
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = multi_scale_ssim(
img_r_tensor,
Expand All @@ -270,6 +277,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r,
img_m,
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = multi_scale_ssim(
img_r_tensor,
Expand Down
8 changes: 8 additions & 0 deletions src/viqa/fr_metrics/vif.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,9 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
HVS model parameter (variance of the visual noise). See [3]_.
reduction : str, default='mean'
Specifies the reduction type: 'none', 'mean' or 'sum'.
device : Union[str, torch.device], default 'cpu'
Determines the device if the image is a PyTorch tensor,
e.g. "cuda", "cpu", "cuda:0", ...

Returns
-------
Expand Down Expand Up @@ -176,6 +179,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r[im_slice, :, :],
img_m[im_slice, :, :],
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = vif_p(
img_r_tensor,
Expand All @@ -188,6 +192,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r[:, im_slice, :],
img_m[:, im_slice, :],
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = vif_p(
img_r_tensor,
Expand All @@ -200,6 +205,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r[:, :, im_slice],
img_m[:, :, im_slice],
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = vif_p(
img_r_tensor,
Expand All @@ -222,6 +228,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r,
img_m,
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = vif_p(
img_r_tensor,
Expand All @@ -243,6 +250,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r,
img_m,
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = vif_p(
img_r_tensor,
Expand Down
8 changes: 8 additions & 0 deletions src/viqa/fr_metrics/vsi.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,9 @@ class VSI(FullReferenceMetricsInterface):
----------------
chromatic : bool, default False
If True, the input images are expected to be RGB images.
device : Union[str, torch.device], default 'cpu'
Determines the device if the image is a PyTorch tensor,
e.g. "cuda", "cpu", "cuda:0", ...

Raises
------
Expand Down Expand Up @@ -197,6 +200,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r[im_slice, :, :],
img_m[im_slice, :, :],
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = vsi(
img_r_tensor,
Expand All @@ -209,6 +213,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r[:, im_slice, :],
img_m[:, im_slice, :],
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = vsi(
img_r_tensor,
Expand All @@ -221,6 +226,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r[:, :, im_slice],
img_m[:, :, im_slice],
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = vsi(
img_r_tensor,
Expand All @@ -243,6 +249,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r,
img_m,
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = vsi(
img_r_tensor,
Expand All @@ -264,6 +271,7 @@ def score(self, img_r, img_m, dim=None, im_slice=None, **kwargs):
img_r,
img_m,
self.parameters["chromatic"],
self.parameters["device"],
)
score_val = vsi(
img_r_tensor,
Expand Down
18 changes: 11 additions & 7 deletions src/viqa/utils/misc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -417,25 +417,29 @@ def _get_range(cols_rows):
return res


def _check_chromatic(img_r, img_m, chromatic):
def _check_chromatic(img_r, img_m, chromatic, device):
"""Permute image based on dimensions and chromaticity."""
img_r = _to_float(img_r, np.float32)
img_m = _to_float(img_m, np.float32)
# check if chromatic
if chromatic is False:
if img_r.ndim == 3:
# 3D images
img_r_tensor = torch.tensor(img_r).unsqueeze(0).permute(3, 0, 1, 2)
img_m_tensor = torch.tensor(img_m).unsqueeze(0).permute(3, 0, 1, 2)
img_r_tensor = (
torch.tensor(img_r).unsqueeze(0).permute(3, 0, 1, 2).to(device)
)
img_m_tensor = (
torch.tensor(img_m).unsqueeze(0).permute(3, 0, 1, 2).to(device)
)
elif img_r.ndim == 2:
# 2D images
img_r_tensor = torch.tensor(img_r).unsqueeze(0).unsqueeze(0)
img_m_tensor = torch.tensor(img_m).unsqueeze(0).unsqueeze(0)
img_r_tensor = torch.tensor(img_r).unsqueeze(0).unsqueeze(0).to(device)
img_m_tensor = torch.tensor(img_m).unsqueeze(0).unsqueeze(0).to(device)
else:
raise ValueError("Image format not supported.")
else:
img_r_tensor = torch.tensor(img_r).permute(2, 0, 1).unsqueeze(0)
img_m_tensor = torch.tensor(img_m).permute(2, 0, 1).unsqueeze(0)
img_r_tensor = torch.tensor(img_r).permute(2, 0, 1).unsqueeze(0).to(device)
img_m_tensor = torch.tensor(img_m).permute(2, 0, 1).unsqueeze(0).to(device)
return img_r_tensor, img_m_tensor


Expand Down
Loading