Skip to content

Commit e2bb6ba

Browse files
NicolasHugdatumbox
andauthored
[FBcode->GH] Fix accimage tests (#5545)
* Fix accimage tests * Adding workaround for accimage * Refactoring * restore channels Co-authored-by: Vasilis Vryniotis <[email protected]>
1 parent ea197e4 commit e2bb6ba

File tree

1 file changed

+8
-2
lines changed

1 file changed

+8
-2
lines changed

torchvision/transforms/functional_pil.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,10 @@ def _is_pil_image(img: Any) -> bool:
2323
@torch.jit.unused
2424
def get_dimensions(img: Any) -> List[int]:
2525
if _is_pil_image(img):
26-
channels = len(img.getbands())
26+
if hasattr(img, "getbands"):
27+
channels = len(img.getbands())
28+
else:
29+
channels = img.channels
2730
width, height = img.size
2831
return [channels, height, width]
2932
raise TypeError(f"Unexpected type {type(img)}")
@@ -39,7 +42,10 @@ def get_image_size(img: Any) -> List[int]:
3942
@torch.jit.unused
4043
def get_image_num_channels(img: Any) -> int:
4144
if _is_pil_image(img):
42-
return len(img.getbands())
45+
if hasattr(img, "getbands"):
46+
return len(img.getbands())
47+
else:
48+
return img.channels
4349
raise TypeError(f"Unexpected type {type(img)}")
4450

4551

0 commit comments

Comments
 (0)