import os
import pickle
import numpy as np
from sklearn.externals import joblib
from gala import imio, agglo, features, classify


fman = features.default.snemi3d()

def train(index):
    outfn = 'training-data-%i.h5' % index
    if os.path.exists(outfn):
        data, labels = classify.load_training_data_from_disk(outfn,
                                                             names=['data',
                                                                    'labels'])
    else:
        ws_tr = imio.read_image_stack('watershed-%i.lzf.h5' % index)
        pr_tr = imio.read_image_stack('probabilities-inv-norm-%i.lzf.h5' % index)
        gt_tr = imio.read_image_stack('ground-truth-%i.lzf.h5' % index)
        g = agglo.Rag(ws_tr, pr_tr,
                      feature_manager=fman)
        data, labels = g.learn_agglomerate(gt_tr, fman, min_num_epochs=4)[0][:2]
        classify.save_training_data_to_disk([data, labels],
                                            fn=outfn,
                                            names=['data', 'labels'])
    print('total training data:', data.shape)
    print('size in MB:', data.size * data.itemsize / 1e6)
    rf = classify.DefaultRandomForest(n_jobs=6)
    rf.fit(data, labels[:, 0])
    policy = agglo.classifier_probability(fman, rf)
    return policy


def test(index, policy):
    ws = imio.read_image_stack('watershed-%i.lzf.h5' % index)
    pr = imio.read_image_stack('probabilities-inv-norm-%i.lzf.h5' % index)
    g = agglo.Rag(ws, pr, merge_priority_function=policy,
                  feature_manager=fman)
    g.agglomerate(np.inf)
    return g.tree


def train_test_pair(training_index, testing_index):
    print('training %i' % training_index)
    policy = train(training_index)
    print('testing %i' % testing_index)
    tree = test(testing_index, policy)
    with open('results-%i-tr%i.pickle' % (testing_index, training_index),
              'wb') as fout:
        pickle.dump(tree, fout, protocol=-1)
    return tree


if __name__ == '__main__':
    index_pairs = [(tr, 3 - tr) for tr in range(4)]
    trees = joblib.Parallel(n_jobs=4)(joblib.delayed(train_test_pair)(*p)
                                      for p in index_pairs)
    trees = dict(zip(index_pairs, trees))
    with open('results.pickle', 'wb') as fout:
        pickle.dump(trees, fout, protocol=-1)

@@ -36,16 +36,28 @@
 y as np%0A
+import h5py%0A
 from skl
@@ -1997,16 +1997,304 @@
  tree%0A%0A%0A
+def write_saalfeld(fn, raw, labels, res=np.array(%5B12., 1, 1%5D)):%0A    imio.write_h5_stack(raw, fn, group='raw')%0A    imio.write_h5_stack(labels, fn, group='labels')%0A    f = h5py.File(fn, 'a')%0A    f%5B'/raw'%5D.attrs%5B'resolution'%5D = res%0A    f%5B'/labels'%5D.attrs%5B'resolution'%5D = res%0A    f.close()%0A%0A%0A
 if __nam
@@ -2336,26 +2336,26 @@
 = %5B(
-tr, 
 3 - t
-r
+s, ts
 ) for t
-r
+s
  in 
@@ -2613,28 +2613,713 @@
 p(trees, fout, protocol=-1)%0A
+    images = imio.read_image_stack('/groups/saalfeld/saalfeldlab/concha/sample_A/crop/raw/*.tiff')%0A    wss = %5Bimio.read_image_stack('watershed-%25i.lzf.h5' %25 i) for i in range(4)%5D%0A    maps = %5Bt.get_map(0.5) for t in trees%5D%0A    segs = %5Bm%5Bws%5D for m, ws in zip(maps, wss)%5D%0A%0A    seg = np.zeros(images.shape, dtype=np.uint64)%0A    seg%5B:, :625, :625%5D = segs%5B0%5D%0A    seg%5B:, :625, 625:%5D = segs%5B1%5D + np.max(segs%5B0%5D)%0A    seg%5B:, 625:, :625%5D = segs%5B2%5D + np.max(segs%5B0%5D) + np.max(segs%5B1%5D)%0A    seg%5B:, 625:, 625:%5D = segs%5B3%5D + np.max(segs%5B0%5D) + np.max(segs%5B1%5D) + np.max(segs%5B2%5D)%0A%0A    write_saalfeld('/groups/saalfeld/saalfeldlab/concha/sample_A/juan/corners-segments2.h5',%0A                   images, seg)%0A

Add capability to write to saalfeld h5 format