<commit_before>#!/usr/bin/env python

import sys, os, argparse
from numpy import   shape, reshape, \
                    array, zeros, zeros_like, ones, ones_like, arange, \
                    double, \
                    int8, int16, int32, int64, uint8, uint16, uint32, uint64, \
                    iinfo, isscalar, \
                    unique, \
                    where, unravel_index, newaxis, \
                    ceil, floor, prod, cumprod, \
                    concatenate
import itertools
from collections import deque as queue
from scipy.ndimage import filters
from scipy.ndimage.measurements import label
#from scipy.spatial.distance import cityblock as manhattan_distance
import iterprogress as ip
import imio

zero3d = array([0,0,0])

arguments = argparse.ArgumentParser(add_help=False)
arggroup = arguments.add_argument_group('Morphological operations options')
arggroup.add_argument('-S', '--save-watershed', metavar='FILE',
    help='Write the watershed result to FILE (overwrites).'
)
arggroup.add_argument('-w', '--watershed', metavar='WS_FN',
    type=imio.read_image_stack_single_arg,
    help='Use a precomputed watershed volume from file.'
)
arggroup.add_argument('--seed', metavar='FN', 
    type=imio.read_image_stack_single_arg,
    help='''use the volume in FN to seed the watershed. By default, connected
        components of 0-valued pixels will be used as the seeds.'''
)

def manhattan_distance(a, b):
    return sum(abs(a-b))

def diamondse(sz, dimension):
    d = floor(sz/2)+1 # ceil has problems if sz.dtype is int
    sz = sz * ones(dimension)
    se = zeros(sz)
    ctr = floor(sz/2)
    for i in [array(unravel_index(j,sz)) for j in range(se.size)]:
        if manhattan_distance(i, ctr) < d:
            se[tuple(i)] = 1
    return se


def watershed(a, seeds=None, show_progress=False):
    seeded = seeds is not None
    if not seeded:
        ws = zeros(shape(a), uint32)
    else:
        ws = seeds
    a = pad(a, a.max()+1)
    ws = pad(ws, 0)
    maxlabel = iinfo(ws.dtype).max
    sel = diamondse(3, a.ndim)
    current_label = 0
    neighbors = build_neighbors_array(a)
    level_pixels = build_levels_dict(a)
    if show_progress:
        def with_progress(it): 
            return ip.with_progress(
                it, pbar=ip.StandardProgressBar('Watershed... ')
            )
    else:
        def with_progress(it):
            return ip.with_progress(it, pbar=ip.NoProgressBar())
    for level in with_progress(sorted(level_pixels.keys())[:-1]):
        idxs_adjacent_to_labels = queue([idx for idx in level_pixels[level] if
                                            any(ws.ravel()[neighbors[idx]])])
        while len(idxs_adjacent_to_labels) > 0:
            idx = idxs_adjacent_to_labels.popleft()
            adj_labels = ws.ravel()[neighbors[idx]]
            adj_labels = unique(adj_labels[
                ((adj_labels != 0) * (adj_labels != maxlabel)).astype(bool)
            ])
            if len(adj_labels) > 1:
                ws.ravel()[idx] = maxlabel # build a dam
            elif len(adj_labels) == 1 and ws.ravel()[idx] == 0:
                ws.ravel()[idx] = adj_labels[0]
                ns = neighbors[idx]
                idxs_adjacent_to_labels.extend(ns[((ws.ravel()[ns] == 0) * 
                                    (a.ravel()[ns] == level)).astype(bool) ])
        if seeded:
            not_adj = where((ws == 0) * (a == level))[0]
            try:
                level_pixels[level+1].extend(not_adj)
            except KeyError:
                level_pixels[level+1] = list(not_adj)
            a.ravel()[not_adj] = level+1
        else:
            new_labels, num_new = label((ws == 0) * (a == level), sel)
            new_labels = (current_label + new_labels) * (new_labels != 0)
            current_label += num_new
            ws += new_labels
    ws[ws==maxlabel] = 0
    return juicy_center(ws)

def smallest_int_dtype(number, signed=False, mindtype=None):
    if number < 0: signed = True
    if not signed:
        if number <= iinfo(uint8).max:
            return uint8<commit_msg>hotfix: bug in seeded watershed.<commit_after>#!/usr/bin/env python

import sys, os, argparse
from numpy import   shape, reshape, \
                    array, zeros, zeros_like, ones, ones_like, arange, \
                    double, \
                    int8, int16, int32, int64, uint8, uint16, uint32, uint64, \
                    iinfo, isscalar, \
                    unique, \
                    where, unravel_index, newaxis, \
                    ceil, floor, prod, cumprod, \
                    concatenate
import itertools
from collections import deque as queue
from scipy.ndimage import filters
from scipy.ndimage.measurements import label
#from scipy.spatial.distance import cityblock as manhattan_distance
import iterprogress as ip
import imio

zero3d = array([0,0,0])

arguments = argparse.ArgumentParser(add_help=False)
arggroup = arguments.add_argument_group('Morphological operations options')
arggroup.add_argument('-S', '--save-watershed', metavar='FILE',
    help='Write the watershed result to FILE (overwrites).'
)
arggroup.add_argument('-w', '--watershed', metavar='WS_FN',
    type=imio.read_image_stack_single_arg,
    help='Use a precomputed watershed volume from file.'
)
arggroup.add_argument('--seed', metavar='FN', 
    type=imio.read_image_stack_single_arg,
    help='''use the volume in FN to seed the watershed. By default, connected
        components of 0-valued pixels will be used as the seeds.'''
)

def manhattan_distance(a, b):
    return sum(abs(a-b))

def diamondse(sz, dimension):
    d = floor(sz/2)+1 # ceil has problems if sz.dtype is int
    sz = sz * ones(dimension)
    se = zeros(sz)
    ctr = floor(sz/2)
    for i in [array(unravel_index(j,sz)) for j in range(se.size)]:
        if manhattan_distance(i, ctr) < d:
            se[tuple(i)] = 1
    return se


def watershed(a, seeds=None, show_progress=False):
    seeded = seeds is not None
    if not seeded:
        ws = zeros(shape(a), uint32)
    else:
        ws = seeds
    levels = unique(a)
    a = pad(a, a.max()+1)
    ws = pad(ws, 0)
    maxlabel = iinfo(ws.dtype).max
    sel = diamondse(3, a.ndim)
    current_label = 0
    neighbors = build_neighbors_array(a)
    level_pixels = build_levels_dict(a)
    if show_progress:
        def with_progress(it, *args, **kwargs): 
            return ip.with_progress(
                it, *args, pbar=ip.StandardProgressBar('Watershed...'), **kwargs
            )
    else:
        def with_progress(it, *args, **kwargs):
            return ip.with_progress(it, *args,pbar=ip.NoProgressBar(),**kwargs)
    for i, level in with_progress(enumerate(levels), length=len(levels)):
        idxs_adjacent_to_labels = queue([idx for idx in level_pixels[level] if
                                            any(ws.ravel()[neighbors[idx]])])
        while len(idxs_adjacent_to_labels) > 0:
            idx = idxs_adjacent_to_labels.popleft()
            adj_labels = ws.ravel()[neighbors[idx]]
            adj_labels = unique(adj_labels[
                ((adj_labels != 0) * (adj_labels != maxlabel)).astype(bool)
            ])
            if len(adj_labels) > 1:
                ws.ravel()[idx] = maxlabel # build a dam
            elif len(adj_labels) == 1 and ws.ravel()[idx] == 0:
                ws.ravel()[idx] = adj_labels[0]
                ns = neighbors[idx]
                idxs_adjacent_to_labels.extend(ns[((ws.ravel()[ns] == 0) * 
                                    (a.ravel()[ns] == level)).astype(bool) ])
        if seeded:
            if i+1 < len(levels):
                not_adj = where((ws.ravel() == 0) * (a.ravel() == level))[0]
                level_pixels[levels[i+1]].extend(not_adj)
                a.ravel()[not_adj] = levels[i+1]
        else:
            new_labels, num_new = label((ws == 0) * (a == level), sel)
            new_labels = (current_label + new_labels) * (new_labels != 0)
            current_label += num_new
            ws += new_labels
    ws[ws==maxlabel] = 0
    return juicy_center(ws)

def smallest_int_dtype(number, signed=False, mindtype=None):
    if number < 0: signed = True
    if not signed:
        if number <= iinfo(uint8).max:
            return uint8
{'fraction_non_alphanumeric': 0.12049945447933083, 'fraction_numerical': 0.011516547460298218, 'mean_word_length': 2.268621236133122, 'pattern_counts': {'":': 0, '<': 10, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5024529', 'n_tokens_mistral': 2759, 'n_tokens_neox': 2589, 'n_words': 627}