Lint

Author: AntoineSimoulin

test/test_cvcuda.py

@@ -1,15 +1,16 @@
-from torchvision import _is_cvcuda_available
-import torch
 import pytest
 import torch
+from torchvision import _is_cvcuda_available
 from torchvision.transforms.v2 import functional as F
+
 CVCUDA_AVAILABLE = _is_cvcuda_available()
 CUDA_AVAILABLE = torch.cuda.is_available()
 
 
 if CVCUDA_AVAILABLE:
     import nvcv
 
+
 @pytest.mark.skipif(CVCUDA_AVAILABLE is False, reason="test requires CVCUDA")
 @pytest.mark.skipif(CUDA_AVAILABLE is False, reason="test requires CUDA")
 class TestToNvcvTensor:
@@ -156,16 +157,24 @@ def transform_cls_to_functional(get_transform_cls):
     def wrapper(inpt):
         transform_cls = get_transform_cls()
         return transform_cls()(inpt)
+
     return wrapper
 
 
 @pytest.mark.skipif(CVCUDA_AVAILABLE is False, reason="test requires CVCUDA")
 @pytest.mark.skipif(CUDA_AVAILABLE is False, reason="test requires CUDA")
 class TestNVCVToTensor:
-    
     @pytest.mark.parametrize("num_channels", [1, 3, 4])
     @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32])
-    @pytest.mark.parametrize("fn", [F.nvcv_to_tensor, transform_cls_to_functional(lambda: __import__('torchvision.transforms.v2', fromlist=['NVCVToTensor']).NVCVToTensor)])
+    @pytest.mark.parametrize(
+        "fn",
+        [
+            F.nvcv_to_tensor,
+            transform_cls_to_functional(
+                lambda: __import__("torchvision.transforms.v2", fromlist=["NVCVToTensor"]).NVCVToTensor
+            ),
+        ],
+    )
     def test_functional_and_transform(self, num_channels, dtype, fn):
         input = make_nvcv_image(num_channels=num_channels, dtype=dtype)
         output = fn(input)

torchvision/transforms/v2/__init__.py

@@ -62,4 +62,4 @@
 from torchvision import _is_cvcuda_available
 
 if _is_cvcuda_available():
-    from ._cvcuda import ToNVCVTensor, NVCVToTensor
+    from ._cvcuda import NVCVToTensor, ToNVCVTensor

torchvision/transforms/v2/_cvcuda.py

@@ -1,5 +1,5 @@
-from torchvision.utils import _log_api_usage_once
 from torchvision.transforms.v2 import functional as F
+from torchvision.utils import _log_api_usage_once
 
 
 class ToNVCVTensor:

torchvision/transforms/v2/functional/_cvcuda.py

@@ -1,8 +1,9 @@
-import numpy as np
 import warnings
+
+import numpy as np
+import nvcv
 import torch
 from torchvision.utils import _log_api_usage_once
-import nvcv
 
 
 def _infer_nvcv_format(img_tensor: torch.Tensor):
@@ -100,8 +101,7 @@ def _validate_nvcv_format(format, num_channels: int) -> None:
     elif num_channels == 2:
         if format not in two_channel_formats:
             warnings.warn(
-                f"Format {format} may not be appropriate for 2-channel image. "
-                f"Common 2-channel formats: _2F32"
+                f"Format {format} may not be appropriate for 2-channel image. " f"Common 2-channel formats: _2F32"
             )
     elif num_channels == 3:
         if format not in three_channel_formats:
@@ -159,7 +159,7 @@ def to_nvcv_tensor(pic, format=None):
     # Ensure image has channel dimension for unbatched case
     if img_tensor.ndim == 2:
         img_tensor = img_tensor.unsqueeze(2)  # H W -> H W C
-    
+
     # Validate dimensions
     if img_tensor.ndim not in (3, 4):
         raise ValueError(f"pic should be 2/3/4 dimensional. Got {img_tensor.ndim} dimensions.")
@@ -201,6 +201,7 @@ def to_nvcv_tensor(pic, format=None):
 
     # Convert to NVCV tensor with NHWC layout
     import cvcuda
+
     return cvcuda.as_tensor(img_tensor.contiguous(), nvcv.TensorLayout.NHWC)
 
 
@@ -222,7 +223,7 @@ def nvcv_to_tensor(nvcv_img):
 
     # Convert NVCV Tensor to PyTorch tensor via CUDA array interface
     # NVCV tensors expose __cuda_array_interface__ which PyTorch can consume directly
-    cuda_tensor = torch.as_tensor(nvcv_img.cuda(), device='cuda')
+    cuda_tensor = torch.as_tensor(nvcv_img.cuda(), device="cuda")
 
     # Handle different dimensionalities
     if cuda_tensor.ndim == 4: