Initial commit.

[pysvrt.git] / test-svrt.py

diff --git a/test-svrt.py b/test-svrt.py
index 92fc554..9aa2d59 100755 (executable)
--- a/test-svrt.py
+++ b/test-svrt.py
@@ -24,16 +24,27 @@
 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 _ext import svrt
 
 from torchvision import datasets, transforms, utils
 
 from _ext import svrt
 

-train_set = svrt.generate_vignettes(12, 1234)
+labels = torch.LongTensor(12).zero_()
+labels.narrow(0, 0, labels.size(0)//2).fill_(1)
+
+train_set = svrt.generate_vignettes(4, labels)

 
 print(str(type(train_set)), train_set.size())
 
 print(str(type(train_set)), train_set.size())

+
+train_set = train_set.view(train_set.size(0), 1, train_set.size(1), train_set.size(2))
+
+train_set.div_(255)
+
+torchvision.utils.save_image(train_set, 'example.png')