!7

lib 1400166 // jython 2.7

import org.python.core.*;
import org.python.modules.cPickle;

static LS colors = splitAtComma("red, blue, green, orange, yellow, gray");

p-exp {
  File f = userDir("dev/Relational_Network_Tensorflow-master/data/sort-of-clevr.100-10.pickle");

  PyFile pyFile = new(bufferedInputStream(f));
  O o = cPickle.load(pyFile);
  print("Loaded: " + className(o));
  
  L trainAndTest = cast o;
  for (L x : (L<L>) trainAndTest) {
    print(l(x) + " examples");
    L entry = cast first(x);
    print(l(entry));
    L<LL<Double>> img = (L) entry.get(0);
    L<L> relations = cast entry.get(1);
    L<L> norelations = cast entry.get(2);
    //print("img: " + sfu(img)); // img is a LLL<Double>
    showImage(rgbImageFromLLLDouble255((L) img));
    LL<Double> relQuestions = (L) first(relations);
    L<Double> relAnswers = (L) second(relations);
    LL<Double> norelQuestions = (L) first(norelations);
    L<Int> norelAnswers = (L) second(norelations);
    print("relations:" + sfu(relations));
    print("norelations:" + sfu(norelations));
    for i over norelQuestions: {
      L<Double> q = norelQuestions.get(i);
      Int a = norelAnswers.get(i);
      S color = get(colors, indexOf(q, 1.0));
      S subtype = get(
        ll("query shape->rectangle/circle",
           "query horizontal position->yes/no",
           "query vertical position->yes/no"),
        indexOf(q, 8, 1.0)-8);
      print("norel: " + q + " => " + a + ". color: " + color + ", subtype: " + subtype);
    }
  }
}