Реализация сегментации водораздела на Java

Я пытаюсь написать свою собственную реализацию Watershed Segmentation для проекта. У меня есть версия, которая возвращает что-то похожее на правильную сегментацию при действительно тривиальных картинках. К сожалению, он очень медленный / неэффективный и может завершиться, а может и не завершиться во всех случаях.

Я работал с описанием в «Цифровой обработке изображений» Вудса и Гонсалеса и со страницы Википедии Watershed. Общий алгоритм закодирован и включен ниже, но я чувствую, что перебираю много вещей, в которых мне не нужно быть. Я ценю любую помощь здесь, заранее спасибо.

    public static Set<Point> watershed(double[][] im) {

    //Get the Gradient Magnitude to use as our "topographic map."

    double t1 = System.currentTimeMillis();
    double[][] edges = EdgeDetection.sobelMagnitude(im);


    //Only iterate over values in the gradient magnitude to speed up.

    double[] histogram = Image.getHistogram(edges);
    Image.drawHistogram(histogram, Color.red);
    int minPixelValue = 0;
    for (int i = 0; i < histogram.length; i++) {
        if (histogram[i] > 0) {
            minPixelValue = i;
            break;
        }
    }
    int h = im.length;
    int w = im[0].length;

    //SE is a 3x3 structuring element for morphological dilation.
    boolean[][] SE = {{true, true, true}, {true, true, true}, {true, true, true}};

    //Keeping track of last iteration's components to see if two flooded together.
    ArrayList<Set<Point>> lastComponents = connectedComponents(getSet(EdgeDetection.underThreshold(edges, minPixelValue + 1)));
    ArrayList<Set<Point>> components;
    Set<Point> boundary = new HashSet<Point>();

    for (int i = minPixelValue + 1; i < 256; i++) {
        if (histogram[i] != 0) {
            System.out.println("BEHHH " + i);
            t1 = System.currentTimeMillis();
            ArrayList<Integer> damLocations = new ArrayList<Integer>();
            HashMap<Integer, ArrayList<Integer>> correspondingSets = new HashMap<Integer, ArrayList<Integer>>();
            //Figure out which of the old sets the new sets incorporated.
            //Here is where we check if sets flooded into eachother.
            //System.out.println("Checking for flooding");
            components = connectedComponents(getSet(EdgeDetection.underThreshold(edges, i)));
            for (int nc = 0; nc < components.size(); nc++) {
                //System.out.println("Checking component " + nc);
                Set<Point> newComponent = components.get(nc);
                for (int oc = 0; oc < lastComponents.size(); oc++) {
                    //System.out.println("    Against component " + oc);
                    Set<Point> oldComponent = lastComponents.get(oc);
                    if (numberInCommon(newComponent, oldComponent) > 0) {
                        //System.out.println("     In there.");
                        ArrayList<Integer> oldSetsContained;
                        if (correspondingSets.containsKey(nc)) {
                            oldSetsContained = correspondingSets.get(nc);
                            damLocations.add(nc);
                        } else {
                            //System.out.println("     Nope.");
                            oldSetsContained = new ArrayList<Integer>();

                        }
                        oldSetsContained.add(oc);
                        correspondingSets.put(nc, oldSetsContained);
                    }
                }
            }
            System.out.println("Calculating overlapping sets: " + (System.currentTimeMillis() - t1));

            //System.out.println("Check done.");
            for (int key : correspondingSets.keySet()) {
                Integer[] cs = new Integer[correspondingSets.get(key).size()];
                correspondingSets.get(key).toArray(cs);
                if (cs.length == 1) {
                    //System.out.println("Set " + cs[0] + " has grown without flooding.");
                } else {
                    //System.out.println("The following sets have flooded together: " + Arrays.toString(cs));
                }
            }

            //Build Damns to prevent flooding

            for (int c : damLocations) {

                System.out.println("Building dam for component " + c);
                Set<Point> bigComponent = components.get(c);
                System.out.println("Total size: " + bigComponent.size());
                ArrayList<Set<Point>> littleComponent = new ArrayList<Set<Point>>();
                for (int lcindex : correspondingSets.get(c)) {
                    littleComponent.add(lastComponents.get(lcindex));
                }

                Set<Point> unionSet = new HashSet<Point>(boundary);
                for (Set<Point> lc : littleComponent) {
                    unionSet = union(unionSet, lc);

                }
                System.out.println("Building union sets: " + (System.currentTimeMillis() - t1));
                while (intersection(unionSet, bigComponent).size() < bigComponent.size()) {

                    for (int lIndex = 0; lIndex < littleComponent.size(); lIndex++) {
                        Set<Point> lc = littleComponent.get(lIndex);
                        Set<Point> lcBoundary = extractBoundaries(lc, SE, h, w);
                        Set<Point> toAdd = new HashSet<Point>();
                        Set<Point> otherComponents = new HashSet<Point>(unionSet);
                        otherComponents.removeAll(lc);
                        otherComponents.removeAll(boundary);
                        otherComponents = extractBoundaries(otherComponents, SE, h, w);
                        for (Point pt : lcBoundary) {
                            Set<Point> eightNbrs = get8Neighborhood(pt);
                            for (Point nbr : eightNbrs) {
                                if (bigComponent.contains(nbr) & !boundary.contains(nbr)) {
                                    Set<Point> neighborNbr = get8Neighborhood(nbr);
                                    if (intersection(neighborNbr, otherComponents).size() > 0) {
                                        boundary.add(nbr);
                                        edges[nbr.y][nbr.x] = 256;
                                        break;
                                    } else if (!lc.contains(nbr)) {
                                        toAdd.add(nbr);
                                        //if(i==65)System.out.println("Adding point " + nbr.y + " " + nbr.x);
                                    } else {
                                        //if(i==65)System.out.println("Already in here " + nbr.y + " " + nbr.x);
                                    }
                                }
                            }
                        }
                        t1 = System.currentTimeMillis();
                        lc.addAll(toAdd);
                        toAdd.removeAll(toAdd);

                        littleComponent.set(lIndex, lc);
                        unionSet = new HashSet<Point>(boundary);
                        for (Set<Point> ltc : littleComponent) {
                            unionSet = union(unionSet, ltc);
                        }
                        System.out.println("This is a donk " + intersection(unionSet, bigComponent).size());
                        otherComponents = new HashSet<Point>(unionSet);
                        otherComponents.removeAll(lc);
                        otherComponents.removeAll(boundary);
                    }
                }
            }
        }
    }
    boundary = close(boundary,h,w);
    Image.drawSet(boundary, h, w);
    return boundary;
}