From 24368498f3065e8a4be34c5e8e2b68f9d1220f7d Mon Sep 17 00:00:00 2001
From: Francois Fleuret <francois@fleuret.org>
Date: Wed, 14 Jun 2017 18:27:51 +0200
Subject: [PATCH] svrt.generate_vignettes now takes a 1d label tensor as
 arguments.

---
 svrt.c            | 19 ++++++++++++++++---
 svrt.h            |  2 +-
 svrt_generator.cc |  5 +++--
 svrt_generator.h  |  3 ++-
 test-svrt.py      |  9 +++++++--
 5 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/svrt.c b/svrt.c
index fdee66f..1a2449b 100644
--- a/svrt.c
+++ b/svrt.c
@@ -27,14 +27,25 @@
 
 #include "svrt_generator.h"
 
-THByteTensor *generate_vignettes(long n_problem, long nb_vignettes) {
+THByteTensor *generate_vignettes(long n_problem, THLongTensor *labels) {
   struct VignetteSet vs;
+  long nb_vignettes;
   long st0, st1, st2;
   long v, i, j;
+  long *m, *l;
   unsigned char *a, *b;
 
-  svrt_generate_vignettes(n_problem, nb_vignettes, &vs);
-  printf("SANITY %d %d %d\n", vs.nb_vignettes, vs.width, vs.height);
+  nb_vignettes = THLongTensor_size(labels, 0);
+  m = THLongTensor_storage(labels)->data + THLongTensor_storageOffset(labels);
+  st0 = THLongTensor_stride(labels, 0);
+  l = (long *) malloc(sizeof(long) * nb_vignettes);
+  for(v = 0; v < nb_vignettes; v++) {
+    l[v] = *m;
+    m += st0;
+  }
+
+  svrt_generate_vignettes(n_problem, nb_vignettes, l, &vs);
+  free(l);
 
   THLongStorage *size = THLongStorage_newWithSize(3);
   size->data[0] = vs.nb_vignettes;
@@ -61,5 +72,7 @@ THByteTensor *generate_vignettes(long n_problem, long nb_vignettes) {
     }
   }
 
+  free(vs.data);
+
   return result;
 }
diff --git a/svrt.h b/svrt.h
index 4335df3..77b8b46 100644
--- a/svrt.h
+++ b/svrt.h
@@ -23,4 +23,4 @@
  *
  */
 
-THByteTensor *generate_vignettes(long n_problem, long nb_images);
+THByteTensor *generate_vignettes(long n_problem, THLongTensor *labels);
diff --git a/svrt_generator.cc b/svrt_generator.cc
index 80cfd12..90f781d 100644
--- a/svrt_generator.cc
+++ b/svrt_generator.cc
@@ -153,7 +153,8 @@ struct VignetteSet {
   unsigned char *data;
 };
 
-void svrt_generate_vignettes(int n_problem, int nb_vignettes, VignetteSet *result) {
+void svrt_generate_vignettes(int n_problem, int nb_vignettes, long *labels,
+                             VignetteSet *result) {
   Vignette tmp;
 
   VignetteGenerator *vg = new_generator(n_problem);
@@ -165,7 +166,7 @@ void svrt_generate_vignettes(int n_problem, int nb_vignettes, VignetteSet *resul
 
   unsigned char *s = result->data;
   for(int i = 0; i < nb_vignettes; i++) {
-    vg->generate(drand48() < 0.5 ? 1 : 0, &tmp);
+    vg->generate(labels[i], &tmp);
     int *r = tmp.content;
     for(int k = 0; k < Vignette::width * Vignette::height; k++) {
       *s++ = *r++;
diff --git a/svrt_generator.h b/svrt_generator.h
index bdfe5c1..7f6a3ad 100644
--- a/svrt_generator.h
+++ b/svrt_generator.h
@@ -35,7 +35,8 @@ struct VignetteSet {
   unsigned char *data;
 };
 
-void svrt_generate_vignettes(int n_problem, int nb_vignettes, struct VignetteSet *result);
+  void svrt_generate_vignettes(int n_problem, int nb_vignettes, long *labels,
+                               struct VignetteSet *result);
 
 #ifdef __cplusplus
 }
diff --git a/test-svrt.py b/test-svrt.py
index 6b5f826..9aa2d59 100755
--- a/test-svrt.py
+++ b/test-svrt.py
@@ -36,10 +36,15 @@ from torchvision import datasets, transforms, utils
 
 from _ext import svrt
 
-train_set = svrt.generate_vignettes(12, 64)
+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())
 
+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.view(train_set.size(0), 1, train_set.size(1), train_set.size(2)), 'example.png')
+torchvision.utils.save_image(train_set, 'example.png')
-- 
2.39.5