Implementing Watershed Segmentation in Java
I'm trying to write my own implementation of Watershed Segmentation for a project. I have a version that returns something resembling the correct segmentation given really trivial pictures. Unfortunately, it's super-slow/inefficient and it may or may not terminate in all cases.
I've been working from the description in "Digital Image Processing," by Woods and Gonzales, and from the Watershed Wikipedia page. The general algorithm is coded and included below, but I have a feeling I'm looping over a lot of things I do not need to be. I appreciate any and all help here, thanks in advance.
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;
开发者_Go百科 } 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;
}
The algorithm as it seems is at most O(n^2). You have many many nested loops.. I was not able to find the Woods description. This code by Christopher Mei implements the algorithm: it is a really simple implementation.
WatershedPixel.java
/*
* Watershed algorithm
*
* Copyright (c) 2003 by Christopher Mei (christopher.mei@sophia.inria.fr)
*
* This plugin is free software; you can redistribute it and/or modify
* it under the terms of the GNU General Public License version 2
* as published by the Free Software Foundation.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this plugin; if not, write to the Free Software
* Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
*/
import java.lang.*;
import java.util.*;
import ij.*;
/**
* The aim of WatershedPixel is to enable
* sorting the pixels of an Image according
* to their grayscale value.
*
* This is the first step of the Vincent
* and Soille Watershed algorithm (1991)
*
**/
public class WatershedPixel implements Comparable {
/** Value used to initialise the image */
final static int INIT = -1;
/** Value used to indicate the new pixels that
* are going to be processed (intial value
* at each level)
**/
final static int MASK = -2;
/** Value indicating that the pixel belongs
* to a watershed.
**/
final static int WSHED = 0;
/** Fictitious pixel **/
final static int FICTITIOUS = -3;
/** x coordinate of the pixel **/
private int x;
/** y coordinate of the pixel **/
private int y;
/** grayscale value of the pixel **/
private byte height;
/** Label used in the Watershed immersion algorithm **/
private int label;
/** Distance used for working on pixels */
private int dist;
/** Neighbours **/
private Vector neighbours;
public WatershedPixel(int x, int y, byte height) {
this.x = x;
this.y = y;
this.height = height;
label = INIT;
dist = 0;
neighbours = new Vector(8);
}
public WatershedPixel() {
label = FICTITIOUS;
}
public void addNeighbour(WatershedPixel neighbour) {
/*IJ.write("In Pixel, adding :");
IJ.write(""+neighbour);
IJ.write("Add done");
*/
neighbours.add(neighbour);
}
public Vector getNeighbours() {
return neighbours;
}
public String toString() {
return new String("("+x+","+y+"), height : "+getIntHeight()+", label : "+label+", distance : "+dist);
}
public final byte getHeight() {
return height;
}
public final int getIntHeight() {
return (int) height&0xff;
}
public final int getX() {
return x;
}
public final int getY() {
return y;
}
/** Method to be able to use the Collections.sort static method. **/
public int compareTo(Object o) {
if(!(o instanceof WatershedPixel))
throw new ClassCastException();
WatershedPixel obj = (WatershedPixel) o;
if( obj.getIntHeight() < getIntHeight() )
return 1;
if( obj.getIntHeight() > getIntHeight() )
return -1;
return 0;
}
public void setLabel(int label) {
this.label = label;
}
public void setLabelToINIT() {
label = INIT;
}
public void setLabelToMASK() {
label = MASK;
}
public void setLabelToWSHED() {
label = WSHED;
}
public boolean isLabelINIT() {
return label == INIT;
}
public boolean isLabelMASK() {
return label == MASK;
}
public boolean isLabelWSHED() {
return label == WSHED;
}
public int getLabel() {
return label;
}
public void setDistance(int distance) {
dist = distance;
}
public int getDistance() {
return dist;
}
public boolean isFICTITIOUS() {
return label == FICTITIOUS;
}
public boolean allNeighboursAreWSHED() {
for(int i=0 ; i<neighbours.size() ; i++) {
WatershedPixel r = (WatershedPixel) neighbours.get(i);
if( !r.isLabelWSHED() )
return false;
}
return true;
}
}
WatershedFIFO.java
/*
* Watershed plugin
*
* Copyright (c) 2003 by Christopher Mei (christopher.mei@sophia.inria.fr)
*
* This plugin is free software; you can redistribute it and/or modify
* it under the terms of the GNU General Public License version 2
* as published by the Free Software Foundation.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this plugin; if not, write to the Free Software
* Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
*/
import java.util.*;
import ij.*;
/** This class implements a FIFO queue that
* uses the same formalism as the Vincent
* and Soille algorithm (1991)
**/
public class WatershedFIFO {
private LinkedList watershedFIFO;
public WatershedFIFO() {
watershedFIFO = new LinkedList();
}
public void fifo_add(WatershedPixel p) {
watershedFIFO.addFirst(p);
}
public WatershedPixel fifo_remove() {
return (WatershedPixel) watershedFIFO.removeLast();
}
public boolean fifo_empty() {
return watershedFIFO.isEmpty();
}
public void fifo_add_FICTITIOUS() {
watershedFIFO.addFirst(new WatershedPixel());
}
public String toString() {
StringBuffer ret = new StringBuffer();
for(int i=0; i<watershedFIFO.size(); i++) {
ret.append( ((WatershedPixel)watershedFIFO.get(i)).toString() );
ret.append( "\n" );
}
return ret.toString();
}
}
WatershedStructure.java
/*
* Watershed algorithm
*
* Copyright (c) 2003 by Christopher Mei (christopher.mei@sophia.inria.fr)
*
* This plugin is free software; you can redistribute it and/or modify
* it under the terms of the GNU General Public License version 2
* as published by the Free Software Foundation.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this plugin; if not, write to the Free Software
* Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
*/
import java.lang.*;
import java.util.*;
import ij.process.*;
import ij.*;
import java.awt.*;
/**
* WatershedStructure contains the pixels
* of the image ordered according to their
* grayscale value with a direct access to their
* neighbours.
*
**/
public class WatershedStructure {
private Vector watershedStructure;
public WatershedStructure(ImageProcessor ip) {
byte[] pixels = (byte[])ip.getPixels();
Rectangle r = ip.getRoi();
int width = ip.getWidth();
int offset, topOffset, bottomOffset, i;
watershedStructure = new Vector(r.width*r.height);
/** The structure is filled with the pixels of the image. **/
for (int y=r.y; y<(r.y+r.height); y++) {
offset = y*width;
IJ.showProgress(0.1+0.3*(y-r.y)/(r.height));
for (int x=r.x; x<(r.x+r.width); x++) {
i = offset + x;
int indiceY = y-r.y;
int indiceX = x-r.x;
watershedStructure.add(new WatershedPixel(indiceX, indiceY, pixels[i]));
}
}
/** The WatershedPixels are then filled with the reference to their neighbours. **/
for (int y=0; y<r.height; y++) {
offset = y*width;
topOffset = offset+width;
bottomOffset = offset-width;
IJ.showProgress(0.4+0.3*(y-r.y)/(r.height));
for (int x=0; x<r.width; x++) {
WatershedPixel currentPixel = (WatershedPixel)watershedStructure.get(x+offset);
if(x+1<r.width) {
currentPixel.addNeighbour((WatershedPixel)watershedStructure.get(x+1+offset));
if(y-1>=0)
currentPixel.addNeighbour((WatershedPixel)watershedStructure.get(x+1+bottomOffset));
if(y+1<r.height)
currentPixel.addNeighbour((WatershedPixel)watershedStructure.get(x+1+topOffset));
}
if(x-1>=0) {
currentPixel.addNeighbour((WatershedPixel)watershedStructure.get(x-1+offset));
if(y-1>=0)
currentPixel.addNeighbour((WatershedPixel)watershedStructure.get(x-1+bottomOffset));
if(y+1<r.height)
currentPixel.addNeighbour((WatershedPixel)watershedStructure.get(x-1+topOffset));
}
if(y-1>=0)
currentPixel.addNeighbour((WatershedPixel)watershedStructure.get(x+bottomOffset));
if(y+1<r.height)
currentPixel.addNeighbour((WatershedPixel)watershedStructure.get(x+topOffset));
}
}
Collections.sort(watershedStructure);
//IJ.showProgress(0.8);
}
public String toString() {
StringBuffer ret = new StringBuffer();
for(int i=0; i<watershedStructure.size() ; i++) {
ret.append( ((WatershedPixel) watershedStructure.get(i)).toString() );
ret.append( "\n" );
ret.append( "Neighbours :\n" );
Vector neighbours = ((WatershedPixel) watershedStructure.get(i)).getNeighbours();
for(int j=0 ; j<neighbours.size() ; j++) {
ret.append( ((WatershedPixel) neighbours.get(j)).toString() );
ret.append( "\n" );
}
ret.append( "\n" );
}
return ret.toString();
}
public int size() {
return watershedStructure.size();
}
public WatershedPixel get(int i) {
return (WatershedPixel) watershedStructure.get(i);
}
}
Watershed_Algorithm.java
/*
* Watershed algorithm
*
* Copyright (c) 2003 by Christopher Mei (christopher.mei@sophia.inria.fr)
*
* This plugin is free software; you can redistribute it and/or modify
* it under the terms of the GNU General Public License version 2
* as published by the Free Software Foundation.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this plugin; if not, write to the Free Software
* Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
*/
import ij.*;
import ij.plugin.filter.PlugInFilter;
import ij.process.*;
import ij.gui.*;
import ij.plugin.frame.PlugInFrame;
import java.awt.*;
import java.util.*;
/**
* This algorithm is an implementation of the watershed immersion algorithm
* written by Vincent and Soille (1991).
*
* @Article{Vincent/Soille:1991,
* author = "Lee Vincent and Pierre Soille",
* year = "1991",
* keywords = "IMAGE-PROC SKELETON SEGMENTATION GIS",
* institution = "Harvard/Paris+Louvain",
* title = "Watersheds in digital spaces: An efficient algorithm
* based on immersion simulations",
* journal = "IEEE PAMI, 1991",
* volume = "13",
* number = "6",
* pages = "583--598",
* annote = "Watershed lines (e.g. the continental divide) mark the
* boundaries of catchment regions in a topographical map.
* The height of a point on this map can have a direct
* correlation to its pixel intensity. WIth this analogy,
* the morphological operations of closing (or opening)
* can be understood as smoothing the ridges (or filling
* in the valleys). Develops a new algorithm for obtaining
* the watershed lines in a graph, and then uses this in
* developing a new segmentation approach based on the
* {"}depth of immersion{"}.",
* }
*
* A review of Watershed algorithms can be found at :
* http://www.cs.rug.nl/~roe/publications/parwshed.pdf
*
* @Article{RoeMei00,
* author = "Roerdink and Meijster",
* title = "The Watershed Transform: Definitions, Algorithms and
* Parallelization Strategies",
* journal = "FUNDINF: Fundamenta Informatica",
* volume = "41",
* publisher = "IOS Press",
* year = "2000",
* }
**/
public class Watershed_Algorithm implements PlugInFilter {
private int threshold;
final static int HMIN = 0;
final static int HMAX = 256;
public int setup(String arg, ImagePlus imp) {
if (arg.equals("about"))
{showAbout(); return DONE;}
return DOES_8G+DOES_STACKS+SUPPORTS_MASKING+NO_CHANGES;
}
public void run(ImageProcessor ip) {
boolean debug = false;
IJ.showStatus("Sorting pixels...");
IJ.showProgress(0.1);
/** First step : the pixels are sorted according to increasing grey values **/
WatershedStructure watershedStructure = new WatershedStructure(ip);
if(debug)
IJ.write(""+watershedStructure.toString());
IJ.showProgress(0.8);
IJ.showStatus("Start flooding...");
if(debug)
IJ.write("Starting algorithm...\n");
/** Start flooding **/
WatershedFIFO queue = new WatershedFIFO();
int curlab = 0;
int heightIndex1 = 0;
int heightIndex2 = 0;
for(int h=HMIN; h<HMAX; h++) /*Geodesic SKIZ of level h-1 inside level h */ {
for(int pixelIndex = heightIndex1 ; pixelIndex<watershedStructure.size() ; pixelIndex++) /*mask all pixels at level h*/ {
WatershedPixel p = watershedStructure.get(pixelIndex);
if(p.getIntHeight() != h) {
/** This pixel is at level h+1 **/
heightIndex1 = pixelIndex;
break;
}
p.setLabelToMASK();
Vector neighbours = p.getNeighbours();
for(int i=0 ; i<neighbours.size() ; i++) {
WatershedPixel q = (WatershedPixel) neighbours.get(i);
if(q.getLabel()>=0) {/*Initialise queue with neighbours at level h of current basins or watersheds*/
p.setDistance(1);
queue.fifo_add(p);
break;
} // end if
} // end for
} // end for
int curdist = 1;
queue.fifo_add_FICTITIOUS();
while(true) /** extend basins **/{
WatershedPixel p = queue.fifo_remove();
if(p.isFICTITIOUS())
if(queue.fifo_empty())
break;
else {
queue.fifo_add_FICTITIOUS();
curdist++;
p = queue.fifo_remove();
}
if(debug) {
IJ.write("\nWorking on :");
IJ.write(""+p);
}
Vector neighbours = p.getNeighbours();
for(int i=0 ; i<neighbours.size() ; i++) /* Labelling p by inspecting neighbours */{
WatershedPixel q = (WatershedPixel) neighbours.get(i);
if(debug)
IJ.write("Neighbour : "+q);
/* Original algorithm :
if( (q.getDistance() < curdist) &&
(q.getLabel()>0 || q.isLabelWSHED()) ) {*/
if( (q.getDistance() <= curdist) &&
(q.getLabel()>=0) ) {
/* q belongs to an existing basin or to a watershed */
if(q.getLabel() > 0) {
if( p.isLabelMASK() )
// Removed from original algorithm || p.isLabelWSHED() )
p.setLabel(q.getLabel());
else
if(p.getLabel() != q.getLabel())
p.setLabelToWSHED();
} // end if lab>0
else
if(p.isLabelMASK())
p.setLabelToWSHED();
}
else
if( q.isLabelMASK() &&
(q.getDistance() == 0) ) {
if(debug)
IJ.write("Adding value");
q.setDistance( curdist+1 );
queue.fifo_add( q );
}
} // end for, end processing neighbours
if(debug) {
IJ.write("End processing neighbours");
IJ.write("New val :\n"+p);
IJ.write("Queue :\n"+queue);
}
} // end while (loop)
/* Detect and process new minima at level h */
for(int pixelIndex = heightIndex2 ; pixelIndex<watershedStructure.size() ; pixelIndex++) {
WatershedPixel p = watershedStructure.get(pixelIndex);
if(p.getIntHeight() != h) {
/** This pixel is at level h+1 **/
heightIndex2 = pixelIndex;
break;
}
p.setDistance(0); /* Reset distance to zero */
if(p.isLabelMASK()) { /* the pixel is inside a new minimum */
curlab++;
p.setLabel(curlab);
queue.fifo_add(p);
while(!queue.fifo_empty()) {
WatershedPixel q = queue.fifo_remove();
Vector neighbours = q.getNeighbours();
for(int i=0 ; i<neighbours.size() ; i++) /* inspect neighbours of p2*/{
WatershedPixel r = (WatershedPixel) neighbours.get(i);
if( r.isLabelMASK() ) {
r.setLabel(curlab);
queue.fifo_add(r);
}
}
} // end while
} // end if
} // end for
} /** End of flooding **/
IJ.showProgress(0.9);
IJ.showStatus("Putting result in a new image...");
/** Put the result in a new image **/
int width = ip.getWidth();
ImageProcessor outputImage = new ByteProcessor(width, ip.getHeight());
byte[] newPixels = (byte[]) outputImage.getPixels();
for(int pixelIndex = 0 ; pixelIndex<watershedStructure.size() ; pixelIndex++) {
WatershedPixel p = watershedStructure.get(pixelIndex);
if(p.isLabelWSHED() && !p.allNeighboursAreWSHED())
newPixels[p.getX()+p.getY()*width] = (byte)255;
}
IJ.showProgress(1);
IJ.showStatus("Displaying result...");
new ImagePlus("Watershed", outputImage).show();
}
void showAbout() {
IJ.showMessage("About Watershed_Algorithm...",
"This plug-in filter calculates the watershed of a 8-bit images.\n" +
"It uses the immersion algorithm written by Vincent and Soille (1991)\n"
);
}
}
加载中,请稍侯......
精彩评论