Update.

diff --git a/build.py b/build.py
index cb77b38..9f03c37 100755 (executable)
--- a/build.py
+++ b/build.py
@@ -3,19 +3,16 @@
 import os
 from torch.utils.ffi import create_extension
 
-this_file = os.path.dirname(__file__)
-
-print('__file__', __file__)
+abs_path = os.path.dirname(os.path.abspath(__file__))
 
 ffi = create_extension(
     '_ext.flatland',
     headers = [ 'flatland.h' ],
     sources = [ 'flatland.c' ],
-    extra_objects = [ '/home/fleuret/sources/python/flatland/flatland_generator.so' ],
+    extra_objects = [ abs_path + '/flatland_generator.so' ],
     libraries = [ ],
     library_dirs = [ ],
     define_macros = [ ],
-    relative_to=this_file,
     with_cuda = False
 )
 

diff --git a/flatland.c b/flatland.c
index 13da6fd..71dee63 100644 (file)
--- a/flatland.c
+++ b/flatland.c
@@ -1,6 +1,3 @@
-
-#include <TH/TH.h>
-
 /*
 
   Example of FFI extension I started from:
@@ -17,38 +14,47 @@
 
  */
 
+#include <TH/TH.h>
+
 #include "sequence_generator.h"
 
-int generate_sequence(long nb_sequences, THByteTensor *output) {
-  long nb_images_per_sequence = 5;
-  long depth = 3;
-  long width = 64;
-  long height = 64;
-  long s;
+THByteTensor *generate_sequence(long nb_sequences, long nb_images_per_sequence, long image_width, long image_height) {
+  long nb_channels = 3;
   unsigned char *a, *b;
-  int c, k, i, j, st0, st1, st2, st3, st4;
+  long s, c, k, i, j, st0, st1, st2, st3, st4;
 
-  THByteTensor_resize5d(output, nb_sequences, nb_images_per_sequence, depth, height, width);
+  THLongStorage *size = THLongStorage_newWithSize(5);
+  size->data[0] = nb_sequences;
+  size->data[1] = nb_images_per_sequence;
+  size->data[2] = nb_channels;
+  size->data[3] = image_height;
+  size->data[4] = image_width;
 
-  st0 = THByteTensor_stride(output, 0);
-  st1 = THByteTensor_stride(output, 1);
-  st2 = THByteTensor_stride(output, 2);
-  st3 = THByteTensor_stride(output, 3);
-  st4 = THByteTensor_stride(output, 4);
+  THByteTensor *result = THByteTensor_newWithSize(size, NULL);
 
-  a =
-    THByteTensor_storage(output)->data + THByteTensor_storageOffset(output);
+  THLongStorage_free(size);
+
+  st0 = THByteTensor_stride(result, 0);
+  st1 = THByteTensor_stride(result, 1);
+  st2 = THByteTensor_stride(result, 2);
+  st3 = THByteTensor_stride(result, 3);
+  st4 = THByteTensor_stride(result, 4);
+
+  unsigned char tmp_buffer[nb_images_per_sequence * nb_channels * image_width * image_height];
 
   for(s = 0; s < nb_sequences; s++) {
-    unsigned char result[nb_images_per_sequence * depth * width * height];
-    unsigned char *r = result;
-    fl_generate_sequences(1, nb_images_per_sequence, width, height, result);
+    a =
+      THByteTensor_storage(result)->data + THByteTensor_storageOffset(result) + s * st0;
+
+    fl_generate_sequences(1, nb_images_per_sequence, image_width, image_height, tmp_buffer);
+
+    unsigned char *r = tmp_buffer;
     for(k = 0; k < nb_images_per_sequence; k++) {
-      for(c = 0; c < depth; c++) {
-        for(i = 0; i < height; i++) {
+      for(c = 0; c < nb_channels; c++) {
+        for(i = 0; i < image_height; i++) {
           b = a
-            + s * st0 + k * st1 + c * st2 + i * st3;
-          for(j = 0; j < width; j++) {
+            + k * st1 + c * st2 + i * st3;
+          for(j = 0; j < image_width; j++) {
             *b = (unsigned char) (*r);
             r++;
             b += st4;
@@ -58,5 +64,5 @@ int generate_sequence(long nb_sequences, THByteTensor *output) {
     }
   }
 
-  return 1;
+  return result;
 }

diff --git a/flatland.h b/flatland.h
index a0a255f..734add4 100644 (file)
--- a/flatland.h
+++ b/flatland.h
@@ -1,2 +1,4 @@
 
-int generate_sequence(long nb_sequences, THByteTensor *output);
+THByteTensor *generate_sequence(long nb_sequences,
+                                long nb_images_per_sequence,
+                                long image_width, long image_height);

diff --git a/test.py b/test.py
index 314e03d..c6b6c48 100755 (executable)
--- a/test.py
+++ b/test.py
@@ -4,6 +4,8 @@ import torch
 import torchvision
 from torchvision import datasets
 
+from _ext import flatland
+
 ######################################################################
 
 def sequences_to_image(x):
@@ -40,10 +42,6 @@ def sequences_to_image(x):
 
 ######################################################################
 
-from _ext import flatland
-
-x = torch.ByteTensor()
-
-flatland.generate_sequence(10, x)
+x = flatland.generate_sequence(5, 3, 128, 96)
 
 sequences_to_image(x).save('sequences.png')