Update.

[flatland.git] / test.py

diff --git a/test.py b/test.py
index c6b6c48..2853309 100755 (executable)
--- a/test.py
+++ b/test.py
@@ -8,7 +8,7 @@ from _ext import flatland
 
 ######################################################################
 
 
 ######################################################################
 

-def sequences_to_image(x):
+def sequences_to_image(x, gap=1, gap_color = (0, 128, 255)):

     from PIL import Image
 
     nb_sequences = x.size(0)
     from PIL import Image
 
     nb_sequences = x.size(0)


@@ -21,8 +21,6 @@ def sequences_to_image(x):
 
     height = x.size(3)
     width = x.size(4)
 
     height = x.size(3)
     width = x.size(4)

-    gap = 1
-    gap_color = (0, 128, 255)

 
     result = torch.ByteTensor(nb_channels,
                               gap + nb_sequences * (height + gap),
 
     result = torch.ByteTensor(nb_channels,
                               gap + nb_sequences * (height + gap),


@@ -34,7 +32,9 @@ def sequences_to_image(x):
 
     for s in range(0, nb_sequences):
         for i in range(0, nb_images_per_sequences):
 
     for s in range(0, nb_sequences):
         for i in range(0, nb_images_per_sequences):

-            result.narrow(1, gap + s * (height + gap), height).narrow(2, gap + i * (width + gap), width).copy_(x[s][i])
+            result.narrow(1, gap + s * (height + gap), height) \
+                  .narrow(2, gap + i * (width + gap), width) \
+                  .copy_(x[s][i])

 
     result_numpy = result.cpu().byte().transpose(0, 2).transpose(0, 1).numpy()
 
 
     result_numpy = result.cpu().byte().transpose(0, 2).transpose(0, 1).numpy()
 


@@ -42,6 +42,6 @@ def sequences_to_image(x):
 
 ######################################################################
 
 
 ######################################################################
 

-x = flatland.generate_sequence(5, 3, 128, 96)
+x = flatland.generate_sequence(1, 3, 80, 80, True, True)

 
 

-sequences_to_image(x).save('sequences.png')
+sequences_to_image(x, gap = 2, gap_color = (0, 0, 0)).save('sequences.png')