Added storage compression / decompression functions to prepare for sets of 1M samples.

[pysvrt.git] / test-svrt.py

diff --git a/test-svrt.py b/test-svrt.py
index 92fc554..5f38fa9 100755 (executable)
--- a/test-svrt.py
+++ b/test-svrt.py
@@ -24,16 +24,29 @@
 import time
 
 import torch
 import time
 
 import torch

+import torchvision

 
 from torch import optim
 from torch import FloatTensor as Tensor
 from torch.autograd import Variable
 from torch import nn
 from torch.nn import functional as fn
 
 from torch import optim
 from torch import FloatTensor as Tensor
 from torch.autograd import Variable
 from torch import nn
 from torch.nn import functional as fn

+

 from torchvision import datasets, transforms, utils
 
 from torchvision import datasets, transforms, utils
 

-from _ext import svrt
+import svrt
+
+labels = torch.LongTensor(12).zero_()
+labels.narrow(0, 0, labels.size(0)//2).fill_(1)
+
+x = svrt.generate_vignettes(4, labels)
+
+print('compression factor {:f}'.format(x.storage().size() / svrt.compress(x.storage()).size()))
+
+x = x.view(x.size(0), 1, x.size(1), x.size(2))
+
+x.div_(255)

 
 

-train_set = svrt.generate_vignettes(12, 1234)
+torchvision.utils.save_image(x, 'example.png')

 
 

-print(str(type(train_set)), train_set.size())
+print('Wrote example.png')