depth first search

#!/usr/bin/python
"""
Author: Jeremy M. Stober
Program: RANSAC.PY
Date: Monday, September 29 2008
Description: Python implmentation of RANSAC.
"""
 
import os, sys, getopt, pdb, string
from numpy import *
import pylab
 
# First, lets generate some data based on an underlying model (y=x) with outliers.
x = random.uniform(low=0.0, high=10.0, size=20)
y = x + random.normal(scale=0.1,size=(20))
 
# Now lets generate some outliers uniformly.
ox = random.uniform(low=0.0, high=10.0, size=10)
oy = random.uniform(low=0.0, high=10.0, size=10)
 
# And combine the model with the outliers.
data = zeros((30,2))
 
data[0:20,0] = x
data[20:30,0] = ox
data[0:20,1] = y
data[20:30,1] = oy
random.shuffle(data) 
 
def choice(seq, n = 1):
    indx = arange(len(seq))
    random.shuffle(indx)
    return take(seq, indx[0:n],axis=0)
 
# Run RANSAC to see if we can recover the line (y=x) from the data.
def RANSAC(data,n,k,t,d):
    """
    Briefly:
    n -- the number of random points that are used to fit the initial model guess.
    k -- the number of iterations.
    t -- the threshold for deciding when a data point fits a model (i.e. is an inlier).
    d -- the number of inliers needed to justify the model.
    """
 
    model_params = array([0.0,0.0]) # slope and intercept
    model_error = 1e6
    model_inliers = []
 
    for i in range(k): # k iterations
        inliers = choice(data,n)
 
        # Fit a line (or generally any model) using lstsq.
        a = transpose(array([inliers[:,0], ones(n)]))
        b = inliers[:,1]
        (p, residuals, rank, s) = linalg.lstsq(a,b)
 
        compatible = []
        for pt in data:
            if abs(pt[1] - dot(p, [pt[0], 1.0])) < t:
                compatible.append(pt)
 
        if len(compatible) > d:
            # The current model is good enough so we should recompute it using all compatible points.
            compatible = array(compatible)
            a = transpose(array([compatible[:,0], ones(len(compatible))]))
            b = compatible[:,1]
            (p, residuals, rank, s) = linalg.lstsq(a,b)
 
            if residuals < model_error:
                model_params = p
                model_error = residuals
                model_inliers = compatible
 
    return (model_params, model_error, model_inliers)
 
(params, error, inliers) =  RANSAC(data,2,100,0.5,20)
 
# Diagnostic: divide the points into inliers and outliers given the model.
outliers = []
for pt in data:
    if not pt in inliers:
        outliers.append(pt)
 
# Diagnostic: plot inliers, outliers, and the model.
pylab.figure(0)
pylab.plot(inliers[:,0],inliers[:,1], '+', color='blue')
outliers = array(outliers)
pylab.plot(outliers[:,0],outliers[:,1],'+', color='red')
 
x = [0.0, 10.0]
y = [params[1], 10.0 * params[0] + params[1]]
pylab.plot(x,y,ls='-',color='green')
 
# Diagnostic: plot original data.
pylab.figure(1)
pylab.plot(data[:,0],data[:,1],'+')
 
pylab.show()