Skip to content

Commit a0e5bbe

Browse files
authored
Remove torchvision.io from test_utils.py (#3092)
For now this simplifies a few things for fbcode integration. We can add this back in the future
1 parent ac288ea commit a0e5bbe

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

test/test_utils.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import numpy as np
12
import os
23
import sys
34
import tempfile
@@ -6,7 +7,6 @@
67
import unittest
78
from io import BytesIO
89
import torchvision.transforms.functional as F
9-
from torchvision.io.image import read_image, write_png
1010
from PIL import Image
1111

1212

@@ -90,9 +90,10 @@ def test_draw_boxes(self):
9090

9191
path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", "fakedata", "draw_boxes_util.png")
9292
if not os.path.exists(path):
93-
write_png(result, path)
93+
res = Image.fromarray(result.permute(1, 2, 0).contiguous().numpy())
94+
res.save(path)
9495

95-
expected = read_image(path)
96+
expected = torch.as_tensor(np.array(Image.open(path))).permute(2, 0, 1)
9697
self.assertTrue(torch.equal(result, expected))
9798

9899

0 commit comments

Comments
 (0)