<commit_before>import numpy as np
from gala import imio, agglo, features, classify


fman = features.default.snemi3d()

def train(index):
    ws_tr = imio.read_image_stack('watershed-%i.lzf.h5' % index)
    pr_tr = imio.read_image_stack('probabilities-%i.lzf.h5' % index) / 255
    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)
    rf = classify.DefaultRandomForest()
    rf.fit(data, labels[:, 0])
    policy = agglo.classifier_priority(fman, rf)
    return policy


def test(index, policy):
    ws = imio.read_image_stack('watershed-%i.lzf.h5' % index)
    pr = imio.read_image_stack('probabilities-%i.lzf.h5' % index) / 255
    g = agglo.Rag(ws, pr, merge_priority_function=policy,<commit_msg>Bug fix: correct indexing of learn_agglomerate output<commit_after>import numpy as np
from gala import imio, agglo, features, classify


fman = features.default.snemi3d()

def train(index):
    ws_tr = imio.read_image_stack('watershed-%i.lzf.h5' % index)
    pr_tr = imio.read_image_stack('probabilities-%i.lzf.h5' % index) / 255
    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]
    rf = classify.DefaultRandomForest()
    rf.fit(data, labels[:, 0])
    policy = agglo.classifier_priority(fman, rf)
    return policy


def test(index, policy):
    ws = imio.read_image_stack('watershed-%i.lzf.h5' % index)
    pr = imio.read_image_stack('probabilities-%i.lzf.h5' % index) / 255
    g = agglo.Rag(ws, pr, merge_priority_function=policy,
{'fraction_non_alphanumeric': 0.13416714042069358, 'fraction_numerical': 0.017055144968732235, 'mean_word_length': 3.8753462603878117, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3654977', 'n_tokens_mistral': 732, 'n_tokens_neox': 690, 'n_words': 121}