Skip to content

Commit 1b4455a

Browse files
ssnlfmassa
authored andcommitted
Add 0.5 in save_images so we round to nearest int (#754)
1 parent f75b814 commit 1b4455a

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

torchvision/utils.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ def norm_range(t, range):
7474
xmaps = min(nrow, nmaps)
7575
ymaps = int(math.ceil(float(nmaps) / xmaps))
7676
height, width = int(tensor.size(2) + padding), int(tensor.size(3) + padding)
77-
grid = tensor.new(3, height * ymaps + padding, width * xmaps + padding).fill_(pad_value)
77+
grid = tensor.new_full((3, height * ymaps + padding, width * xmaps + padding), pad_value)
7878
k = 0
7979
for y in irange(ymaps):
8080
for x in irange(xmaps):
@@ -99,6 +99,7 @@ def save_image(tensor, filename, nrow=8, padding=2,
9999
from PIL import Image
100100
grid = make_grid(tensor, nrow=nrow, padding=padding, pad_value=pad_value,
101101
normalize=normalize, range=range, scale_each=scale_each)
102-
ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
102+
# Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
103+
ndarr = grid.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
103104
im = Image.fromarray(ndarr)
104105
im.save(filename)

0 commit comments

Comments
 (0)